diff --git a/benches/sample_z.rs b/benches/sample_z.rs index 6ecccd9c..6921171c 100644 --- a/benches/sample_z.rs +++ b/benches/sample_z.rs @@ -12,6 +12,7 @@ use criterion::*; use qfall_math::{ integer::{MatZ, Z}, rational::Q, + utils::sample::discrete_gauss::{DiscreteGaussianIntegerSampler, LookupTableSetting}, }; /// benchmark creating a matrix of size 100x100 sampled by a comparatively wide discrete Gaussian distribution. @@ -54,10 +55,29 @@ pub fn bench_sample_z_narrow_single(c: &mut Criterion) { }); } +/// benchmark discrete Gaussian sampling using [`DiscreteGaussianIntegerSampler::sample_z`] for a variety of widths. +pub fn bench_sample_z(c: &mut Criterion) { + let center = 0; + let gaussian_widths = [ + 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, + ]; + + for s in gaussian_widths { + let mut dgis = + DiscreteGaussianIntegerSampler::init(center, s, 6.0, LookupTableSetting::Precompute) + .unwrap(); + + c.bench_function("DiscreteGauss RejectionSampling", |b| { + b.iter(|| dgis.sample_z()) + }); + } +} + criterion_group!( benches, bench_sample_z_wide, bench_sample_z_narrow, bench_sample_z_wide_single, - bench_sample_z_narrow_single + bench_sample_z_narrow_single, + bench_sample_z, ); diff --git a/src/integer/mat_poly_over_z/sample/binomial.rs b/src/integer/mat_poly_over_z/sample/binomial.rs index d4d39618..acc6c5df 100644 --- a/src/integer/mat_poly_over_z/sample/binomial.rs +++ b/src/integer/mat_poly_over_z/sample/binomial.rs @@ -13,8 +13,8 @@ use crate::{ error::MathError, integer::{MatPolyOverZ, PolyOverZ, Z}, rational::Q, - traits::{MatrixDimensions, MatrixSetEntry}, - utils::index::evaluate_index, + traits::{MatrixDimensions, MatrixSetEntry, SetCoefficient}, + utils::{index::evaluate_index, sample::binomial::BinomialSampler}, }; use std::fmt::Display; @@ -107,14 +107,20 @@ impl MatPolyOverZ { ) -> Result { let max_degree = evaluate_index(max_degree)?; let offset: Z = offset.into(); - let n: Z = n.into(); - let p: Q = p.into(); + let mut bin_sampler = BinomialSampler::init(n, p)?; let mut matrix = MatPolyOverZ::new(num_rows, num_cols); for row in 0..matrix.get_num_rows() { for col in 0..matrix.get_num_columns() { - let sample = PolyOverZ::sample_binomial_with_offset(max_degree, &offset, &n, &p)?; - unsafe { matrix.set_entry_unchecked(row, col, sample) }; + let mut poly_z = PolyOverZ::default(); + + for index in 0..=max_degree { + let mut sample = bin_sampler.sample(); + sample += &offset; + unsafe { poly_z.set_coeff_unchecked(index, sample) }; + } + + unsafe { matrix.set_entry_unchecked(row, col, poly_z) }; } } diff --git a/src/integer/mat_z/sample/binomial.rs b/src/integer/mat_z/sample/binomial.rs index 0ea49833..faccf1ad 100644 --- a/src/integer/mat_z/sample/binomial.rs +++ b/src/integer/mat_z/sample/binomial.rs @@ -14,7 +14,7 @@ use crate::{ integer::{MatZ, Z}, rational::Q, traits::{MatrixDimensions, MatrixSetEntry}, - utils::sample::binomial::sample_binomial, + utils::sample::binomial::BinomialSampler, }; use std::fmt::Display; @@ -102,14 +102,14 @@ impl MatZ { p: impl Into, ) -> Result { let offset: Z = offset.into(); - let n: Z = n.into(); - let p: Q = p.into(); + let mut bin_sampler = BinomialSampler::init(n, p)?; let mut matrix = MatZ::new(num_rows, num_cols); for row in 0..matrix.get_num_rows() { for col in 0..matrix.get_num_columns() { - let sample = sample_binomial(&n, &p)?; - unsafe { matrix.set_entry_unchecked(row, col, &offset + sample) }; + let mut sample = bin_sampler.sample(); + sample += &offset; + unsafe { matrix.set_entry_unchecked(row, col, sample) }; } } diff --git a/src/integer/poly_over_z/sample/binomial.rs b/src/integer/poly_over_z/sample/binomial.rs index e04cd2f8..b53ccb8d 100644 --- a/src/integer/poly_over_z/sample/binomial.rs +++ b/src/integer/poly_over_z/sample/binomial.rs @@ -6,15 +6,14 @@ // the terms of the Mozilla Public License Version 2.0 as published by the // Mozilla Foundation. See . -//! This module contains algorithms for sampling -//! according to the binomial distribution. +//! This module contains algorithms for sampling according to the binomial distribution. use crate::{ error::MathError, integer::{PolyOverZ, Z}, rational::Q, traits::SetCoefficient, - utils::{index::evaluate_index, sample::binomial::sample_binomial}, + utils::{index::evaluate_index, sample::binomial::BinomialSampler}, }; use std::fmt::Display; @@ -99,14 +98,15 @@ impl PolyOverZ { ) -> Result { let max_degree = evaluate_index(max_degree)?; let offset: Z = offset.into(); - let n: Z = n.into(); - let p: Q = p.into(); let mut poly_z = PolyOverZ::default(); + let mut bin_sampler = BinomialSampler::init(n, p)?; + for index in 0..=max_degree { - let sample = sample_binomial(&n, &p)?; - unsafe { poly_z.set_coeff_unchecked(index, &offset + sample) }; + let mut sample = bin_sampler.sample(); + sample += &offset; + unsafe { poly_z.set_coeff_unchecked(index, sample) }; } Ok(poly_z) diff --git a/src/integer/z/sample/binomial.rs b/src/integer/z/sample/binomial.rs index 29dc0058..4dca6008 100644 --- a/src/integer/z/sample/binomial.rs +++ b/src/integer/z/sample/binomial.rs @@ -9,7 +9,7 @@ //! This module contains algorithms for sampling //! according to the binomial distribution. -use crate::{error::MathError, integer::Z, rational::Q, utils::sample::binomial::sample_binomial}; +use crate::{error::MathError, integer::Z, rational::Q, utils::sample::binomial::BinomialSampler}; impl Z { /// Chooses a [`Z`] instance according to the binomial distribution @@ -38,11 +38,9 @@ impl Z { /// - Returns a [`MathError`] of type [`ConversionError`](MathError::ConversionError) /// if `n` does not fit into an [`i64`]. pub fn sample_binomial(n: impl Into, p: impl Into) -> Result { - let n: Z = n.into(); - let p: Q = p.into(); + let mut bin_sampler = BinomialSampler::init(n, p)?; - let sample = sample_binomial(&n, &p)?; - Ok(Z::from(sample)) + Ok(bin_sampler.sample()) } } diff --git a/src/integer_mod_q/mat_zq/sample/binomial.rs b/src/integer_mod_q/mat_zq/sample/binomial.rs index 65913bb2..1bb73966 100644 --- a/src/integer_mod_q/mat_zq/sample/binomial.rs +++ b/src/integer_mod_q/mat_zq/sample/binomial.rs @@ -15,7 +15,7 @@ use crate::{ integer_mod_q::{MatZq, Modulus}, rational::Q, traits::{MatrixDimensions, MatrixSetEntry}, - utils::sample::binomial::sample_binomial, + utils::sample::binomial::BinomialSampler, }; use std::fmt::Display; @@ -109,14 +109,14 @@ impl MatZq { p: impl Into, ) -> Result { let offset: Z = offset.into(); - let n: Z = n.into(); - let p: Q = p.into(); + let mut bin_sampler = BinomialSampler::init(n, p)?; let mut matrix = MatZq::new(num_rows, num_cols, modulus); for row in 0..matrix.get_num_rows() { for col in 0..matrix.get_num_columns() { - let sample = sample_binomial(&n, &p)?; - unsafe { matrix.set_entry_unchecked(row, col, &offset + sample) }; + let mut sample = bin_sampler.sample(); + sample += &offset; + unsafe { matrix.set_entry_unchecked(row, col, sample) }; } } diff --git a/src/integer_mod_q/poly_over_zq/sample/binomial.rs b/src/integer_mod_q/poly_over_zq/sample/binomial.rs index 4f9b8b4a..fafd3c7d 100644 --- a/src/integer_mod_q/poly_over_zq/sample/binomial.rs +++ b/src/integer_mod_q/poly_over_zq/sample/binomial.rs @@ -15,7 +15,7 @@ use crate::{ integer_mod_q::{Modulus, PolyOverZq}, rational::Q, traits::SetCoefficient, - utils::{index::evaluate_index, sample::binomial::sample_binomial}, + utils::{index::evaluate_index, sample::binomial::BinomialSampler}, }; use std::fmt::Display; @@ -108,14 +108,14 @@ impl PolyOverZq { let max_degree = evaluate_index(max_degree)?; let offset: Z = offset.into(); let modulus: Modulus = modulus.into(); - let n: Z = n.into(); - let p: Q = p.into(); + let mut bin_sampler = BinomialSampler::init(n, p)?; let mut poly_z = PolyOverZq::from(&modulus); for index in 0..=max_degree { - let sample = sample_binomial(&n, &p)?; - unsafe { poly_z.set_coeff_unchecked(index, &offset + sample) }; + let mut sample = bin_sampler.sample(); + sample += &offset; + unsafe { poly_z.set_coeff_unchecked(index, sample) }; } Ok(poly_z) diff --git a/src/integer_mod_q/z_q/sample/binomial.rs b/src/integer_mod_q/z_q/sample/binomial.rs index 4243f419..50f0dc18 100644 --- a/src/integer_mod_q/z_q/sample/binomial.rs +++ b/src/integer_mod_q/z_q/sample/binomial.rs @@ -14,7 +14,7 @@ use crate::{ integer::Z, integer_mod_q::{Modulus, Zq}, rational::Q, - utils::sample::binomial::sample_binomial, + utils::sample::binomial::BinomialSampler, }; impl Zq { @@ -53,10 +53,9 @@ impl Zq { p: impl Into, ) -> Result { let modulus: Modulus = modulus.into(); - let n: Z = n.into(); - let p: Q = p.into(); + let mut bin_sampler = BinomialSampler::init(n, p)?; - let sample = sample_binomial(&n, &p)?; + let sample = bin_sampler.sample(); Ok(Zq::from((sample, modulus))) } } diff --git a/src/utils/sample.rs b/src/utils/sample.rs index bebf5f46..ab3e752e 100644 --- a/src/utils/sample.rs +++ b/src/utils/sample.rs @@ -8,6 +8,6 @@ //! This module includes core functionality to sample according to random distributions. -pub(crate) mod binomial; +pub mod binomial; pub mod discrete_gauss; pub mod uniform; diff --git a/src/utils/sample/binomial.rs b/src/utils/sample/binomial.rs index f4f5b9c0..93efe8f8 100644 --- a/src/utils/sample/binomial.rs +++ b/src/utils/sample/binomial.rs @@ -1,4 +1,4 @@ -// Copyright © 2023 Niklas Siemer +// Copyright © 2025 Niklas Siemer // // This file is part of qFALL-math. // @@ -10,97 +10,138 @@ //! binomial distribution. use crate::{error::MathError, integer::Z, rational::Q}; +use rand::rngs::ThreadRng; use rand_distr::{Binomial, Distribution}; -/// Chooses a sample according to the binomial distribution parameterized by `n` and `p`. +/// Enables sampling a [`Z`] according to the binomial distribution `Bin(n, p)`. /// -/// Parameters: -/// - `n`: specifies the number of trials -/// - `p`: specifies the probability of success -/// -/// Returns a sample as a [`u64`] chosen from the specified binomial distribution -/// or a [`MathError`] if `n < 0`, `p ∉ (0,1)`, or `n` does not fit into an [`i64`]. +/// Attributes: +/// - `distr`: defines the binomial distribution with parameters `n` and `p` to sample from +/// - `rng`: defines the [`ThreadRng`] that's used to sample from `distr` /// /// # Examples -/// ```compile_fail -/// use qfall_math::{integer::Z, rational::Q}; -/// use qfall_math::utils::sample::binomial::sample_binomial; -/// let n = Z::from(2); -/// let p = Q::from((1, 4)); -/// -/// let sample = sample_binomial(&n, &p).unwrap(); /// ``` +/// use qfall_math::utils::sample::binomial::BinomialSampler; +/// let n = 2; +/// let p = 0.5; /// -/// # Errors and Failures -/// - Returns a [`MathError`] of type [`InvalidIntegerInput`](MathError::InvalidIntegerInput) -/// if `n < 0`. -/// - Returns a [`MathError`] of type [`InvalidInterval`](MathError::InvalidInterval) -/// if `p ∉ (0,1)`. -/// - Returns a [`MathError`] of type [`ConversionError`](MathError::ConversionError) -/// if `n` does not fit into an [`i64`]. -pub(crate) fn sample_binomial(n: &Z, p: &Q) -> Result { - if p <= &Q::ZERO || p >= &Q::ONE { - return Err(MathError::InvalidInterval(format!( - "p (the probability of success for binomial sampling) must be chosen between 0 and 1. \ - Currently it is {p}. \ - Hence, the interval to sample from is invalid and contains only exactly one number." - ))); - } - if n < &Z::ZERO { - return Err(MathError::InvalidIntegerInput(format!( - "n (the number of trials for binomial sampling) must be no smaller than 0. Currently it is {n}." - ))); - } +/// let mut bin_sampler = BinomialSampler::init(n, p).unwrap(); +/// +/// let sample = bin_sampler.sample(); +/// +/// assert!(0 <= sample); +/// assert!(sample < n); +/// ``` +pub struct BinomialSampler { + distr: Binomial, + rng: ThreadRng, +} - let n = i64::try_from(n)? as u64; - let p = f64::from(p); +impl BinomialSampler { + /// Initializes the [`BinomialSampler`] with + /// - `distr` as the binomial distribution with `n` tries and success probability `p` for each try, and + /// - `rng` as a fresh [`ThreadRng`]. + /// + /// Parameters: + /// - `n`: specifies the number of tries + /// - `p`: specifies the success probability + /// + /// Returns a [`BinomialSampler`] or a [`MathError`] if `n < 0`, + /// `p ∉ (0,1)`, or `n` does not fit into an [`i64`]. + /// + /// # Examples + /// ``` + /// use qfall_math::utils::sample::binomial::BinomialSampler; + /// let n = 2; + /// let p = 0.5; + /// + /// let mut bin_sampler = BinomialSampler::init(n, p).unwrap(); + /// ``` + /// + /// # Errors and Failures + /// - Returns a [`MathError`] of type [`InvalidIntegerInput`](MathError::InvalidIntegerInput) + /// if `n < 0`. + /// - Returns a [`MathError`] of type [`InvalidInterval`](MathError::InvalidInterval) + /// if `p ∉ (0,1)`. + /// - Returns a [`MathError`] of type [`ConversionError`](MathError::ConversionError) + /// if `n` does not fit into an [`i64`]. + pub fn init(n: impl Into, p: impl Into) -> Result { + let n = n.into(); + let p = p.into(); + + if p <= Q::ZERO || p >= Q::ONE { + return Err(MathError::InvalidInterval(format!( + "p (the probability of success for binomial sampling) must be chosen between 0 and 1. \ + Currently it is {p}. \ + Hence, the interval to sample from is invalid and contains only exactly one number." + ))); + } + if n < Z::ZERO { + return Err(MathError::InvalidIntegerInput(format!( + "n (the number of trials for binomial sampling) must be no smaller than 0. Currently it is {n}." + ))); + } + + let n = i64::try_from(n)? as u64; + let p = f64::from(&p); - let distr = Binomial::new(n, p).unwrap(); - let mut rng = rand::rng(); + let distr = Binomial::new(n, p).unwrap(); + let rng = rand::rng(); - let sample = distr.sample(&mut rng); + Ok(Self { distr, rng }) + } - Ok(sample) + /// Computes a uniformly chosen [`Z`] sample in `[0, interval_size)` + /// using rejection sampling that accepts samples with probability at least 1/2. + /// + /// # Examples + /// ``` + /// use qfall_math::utils::sample::binomial::BinomialSampler; + /// let n = 2; + /// let p = 0.5; + /// + /// let mut bin_sampler = BinomialSampler::init(n, p).unwrap(); + /// + /// let sample = bin_sampler.sample(); + /// + /// assert!(0 <= sample); + /// assert!(sample < n); + /// ``` + pub fn sample(&mut self) -> Z { + Z::from(self.distr.sample(&mut self.rng)) + } } #[cfg(test)] -mod test_sample_binomial { - use super::{Q, Z, sample_binomial}; - - /// Ensures that the doc tests works correctly. - #[test] - fn doc_test() { - let n = Z::from(2); - let p = Q::from((1, 4)); - - let _ = sample_binomial(&n, &p).unwrap(); - } +mod test_binomial_sampler { + use super::{BinomialSampler, Q, Z}; /// Checks whether the range is kept, /// i.e. if any result is at least 0 and smaller than or equal to `n`. #[test] fn keeps_range() { - let n_u64 = 16; - let n = Z::from(n_u64); - let p = Q::from((1, 2)); + let n = 16; + let p = 0.5; + let mut bin_sampler = BinomialSampler::init(n, p).unwrap(); for _ in 0..16 { - let sample = sample_binomial(&n, &p).unwrap(); + let sample = bin_sampler.sample(); // sample >= 0 check is not required as sample is a u64 - assert!(sample <= n_u64); + assert!(sample <= n); } } /// Roughly checks that the collected samples are distributed according to the binomial distribution. #[test] fn distribution() { - let n = Z::from(2); - let p = Q::from((1, 2)); + let n = 2; + let p = 0.5; + let mut bin_sampler = BinomialSampler::init(n, p).unwrap(); let mut counts = [0; 3]; // count sampled instances for _ in 0..1000 { - let sample = sample_binomial(&n, &p).unwrap() as usize; + let sample = u64::try_from(bin_sampler.sample()).unwrap() as usize; counts[sample] += 1; } @@ -121,20 +162,20 @@ mod test_sample_binomial { /// Checks whether invalid choices for n result in an error. #[test] fn invalid_n() { - let p = Q::from((1, 2)); + let p = 0.5; - assert!(sample_binomial(&Z::MINUS_ONE, &p).is_err()); - assert!(sample_binomial(&Z::from(i64::MIN), &p).is_err()); + assert!(BinomialSampler::init(&Z::MINUS_ONE, p).is_err()); + assert!(BinomialSampler::init(&Z::from(i64::MIN), p).is_err()); } /// Checks whether invalid choices for p result in an error. #[test] fn invalid_p() { - let n = Z::from(2); + let n = 2; - assert!(sample_binomial(&n, &Q::MINUS_ONE).is_err()); - assert!(sample_binomial(&n, &Q::ZERO).is_err()); - assert!(sample_binomial(&n, &Q::ONE).is_err()); - assert!(sample_binomial(&n, &Q::from(5)).is_err()); + assert!(BinomialSampler::init(n, &Q::MINUS_ONE).is_err()); + assert!(BinomialSampler::init(n, &Q::ZERO).is_err()); + assert!(BinomialSampler::init(n, &Q::ONE).is_err()); + assert!(BinomialSampler::init(n, &Q::from(5)).is_err()); } } diff --git a/src/utils/sample/discrete_gauss.rs b/src/utils/sample/discrete_gauss.rs index 3f968758..6a58b80a 100644 --- a/src/utils/sample/discrete_gauss.rs +++ b/src/utils/sample/discrete_gauss.rs @@ -156,18 +156,13 @@ impl DiscreteGaussianIntegerSampler { let mut table = HashMap::new(); - if lookup_table_setting == LookupTableSetting::FillOnTheFly && interval_size > u16::MAX { + if lookup_table_setting != LookupTableSetting::NoLookup && interval_size > u16::MAX { println!( "WARNING: A completely filled lookup table will exceed 2^16 entries. You should reconsider your sampling method for discrete Gaussians." ) } if lookup_table_setting == LookupTableSetting::Precompute { - assert!( - interval_size <= u16::MAX, - "The interval size {interval_size} for discrete Gaussian sampling exceeds 2^16 entries. You should reconsider your sampling method." - ); - let mut i = lower_bound.clone(); while i <= upper_bound { let evaluated_gauss_function = gaussian_function(&i, ¢er, &s); diff --git a/src/utils/sample/uniform.rs b/src/utils/sample/uniform.rs index 6a134b3f..4a8fc646 100644 --- a/src/utils/sample/uniform.rs +++ b/src/utils/sample/uniform.rs @@ -35,10 +35,6 @@ use rand::{RngCore, rngs::ThreadRng}; /// assert!(Z::ZERO <= sample); /// assert!(sample < interval_size); /// ``` -/// -/// # Errors and Failures -/// - Returns a [`MathError`] of type [`InvalidInterval`](MathError::InvalidInterval) -/// if the interval is chosen smaller than or equal to `1`. pub struct UniformIntegerSampler { interval_size: Z, two_pow_32: u64,