use crate::modular::mod_reduce;
use crate::params::Params;
use crate::poly::Poly;
use rand::Rng;
use rand_distr::{Distribution, Normal};
pub fn sample_gaussian<R: Rng>(params: &Params, rng: &mut R) -> Poly {
let sigma = params.sigma_x1000 as f64 / 1000.0;
let normal = Normal::new(0.0, sigma).unwrap();
let coeffs = (0..params.n)
.map(|_| {
let sample = normal.sample(rng).round() as i64;
mod_reduce(sample, params.q)
})
.collect();
Poly { coeffs }
}
pub fn sample_uniform<R: Rng>(params: &Params, rng: &mut R) -> Poly {
let coeffs = (0..params.n).map(|_| rng.gen_range(0..params.q)).collect();
Poly { coeffs }
}
pub fn sample_ternary<R: Rng>(params: &Params, rng: &mut R) -> Poly {
let coeffs = (0..params.n)
.map(|_| {
let v: i64 = rng.gen_range(-1..=1);
mod_reduce(v, params.q)
})
.collect();
Poly { coeffs }
}
pub fn sample_binary<R: Rng>(params: &Params, rng: &mut R) -> Poly {
let coeffs = (0..params.n)
.map(|_| rng.gen_range(0..=1u64))
.collect();
Poly { coeffs }
}
#[cfg(test)]
mod tests {
use super::*;
use rand::SeedableRng;
use rand::rngs::StdRng;
#[test]
fn test_gaussian_in_range() {
let p = Params::test_tiny();
let mut rng = StdRng::seed_from_u64(42);
let poly = sample_gaussian(&p, &mut rng);
for &c in &poly.coeffs {
assert!(c < p.q);
}
assert_eq!(poly.coeffs.len(), p.n);
}
#[test]
fn test_uniform_in_range() {
let p = Params::test_tiny();
let mut rng = StdRng::seed_from_u64(42);
let poly = sample_uniform(&p, &mut rng);
for &c in &poly.coeffs {
assert!(c < p.q);
}
}
#[test]
fn test_ternary_values() {
let p = Params::test_tiny();
let mut rng = StdRng::seed_from_u64(42);
let poly = sample_ternary(&p, &mut rng);
for &c in &poly.coeffs {
assert!(c == 0 || c == 1 || c == p.q - 1, "got {c}");
}
}
#[test]
fn test_gaussian_centered() {
let p = Params::test_small();
let mut rng = StdRng::seed_from_u64(123);
let poly = sample_gaussian(&p, &mut rng);
let centered: Vec<i64> = poly
.coeffs
.iter()
.map(|&c| crate::modular::center(c, p.q))
.collect();
let mean: f64 = centered.iter().map(|&x| x as f64).sum::<f64>() / p.n as f64;
assert!(
mean.abs() < 2.0,
"mean {mean} too far from 0 for Gaussian"
);
}
}