1use crate::{Backend, Module, VecZnx, znx_base::ZnxLayout};
  2use rand_distr::{Distribution, Normal};
  3use sampling::source::Source;
  4
  5pub trait Sampling {
  6    /// Fills the first `size` size with uniform values in \[-2^{log_base2k-1}, 2^{log_base2k-1}\]
  7    fn fill_uniform(&self, log_base2k: usize, a: &mut VecZnx, col_i: usize, size: usize, source: &mut Source);
  8
  9    /// Adds vector sampled according to the provided distribution, scaled by 2^{-log_k} and bounded to \[-bound, bound\].
 10    fn add_dist_f64<D: Distribution<f64>>(
 11        &self,
 12        log_base2k: usize,
 13        a: &mut VecZnx,
 14        col_i: usize,
 15        log_k: usize,
 16        source: &mut Source,
 17        dist: D,
 18        bound: f64,
 19    );
 20
 21    /// Adds a discrete normal vector scaled by 2^{-log_k} with the provided standard deviation and bounded to \[-bound, bound\].
 22    fn add_normal(
 23        &self,
 24        log_base2k: usize,
 25        a: &mut VecZnx,
 26        col_i: usize,
 27        log_k: usize,
 28        source: &mut Source,
 29        sigma: f64,
 30        bound: f64,
 31    );
 32}
 33
 34impl<B: Backend> Sampling for Module<B> {
 35    fn fill_uniform(&self, log_base2k: usize, a: &mut VecZnx, col_a: usize, size: usize, source: &mut Source) {
 36        let base2k: u64 = 1 << log_base2k;
 37        let mask: u64 = base2k - 1;
 38        let base2k_half: i64 = (base2k >> 1) as i64;
 39        (0..size).for_each(|j| {
 40            a.at_mut(col_a, j)
 41                .iter_mut()
 42                .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
 43        })
 44    }
 45
 46    fn add_dist_f64<D: Distribution<f64>>(
 47        &self,
 48        log_base2k: usize,
 49        a: &mut VecZnx,
 50        col_a: usize,
 51        log_k: usize,
 52        source: &mut Source,
 53        dist: D,
 54        bound: f64,
 55    ) {
 56        assert!(
 57            (bound.log2().ceil() as i64) < 64,
 58            "invalid bound: ceil(log2(bound))={} > 63",
 59            (bound.log2().ceil() as i64)
 60        );
 61
 62        let limb: usize = (log_k + log_base2k - 1) / log_base2k - 1;
 63        let log_base2k_rem: usize = log_k % log_base2k;
 64
 65        if log_base2k_rem != 0 {
 66            a.at_mut(col_a, limb).iter_mut().for_each(|a| {
 67                let mut dist_f64: f64 = dist.sample(source);
 68                while dist_f64.abs() > bound {
 69                    dist_f64 = dist.sample(source)
 70                }
 71                *a += (dist_f64.round() as i64) << log_base2k_rem;
 72            });
 73        } else {
 74            a.at_mut(col_a, limb).iter_mut().for_each(|a| {
 75                let mut dist_f64: f64 = dist.sample(source);
 76                while dist_f64.abs() > bound {
 77                    dist_f64 = dist.sample(source)
 78                }
 79                *a += dist_f64.round() as i64
 80            });
 81        }
 82    }
 83
 84    fn add_normal(
 85        &self,
 86        log_base2k: usize,
 87        a: &mut VecZnx,
 88        col_a: usize,
 89        log_k: usize,
 90        source: &mut Source,
 91        sigma: f64,
 92        bound: f64,
 93    ) {
 94        self.add_dist_f64(
 95            log_base2k,
 96            a,
 97            col_a,
 98            log_k,
 99            source,
100            Normal::new(0.0, sigma).unwrap(),
101            bound,
102        );
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::Sampling;
109    use crate::{FFT64, Module, Stats, VecZnx, VecZnxOps, znx_base::ZnxLayout};
110    use sampling::source::Source;
111
112    #[test]
113    fn fill_uniform() {
114        let n: usize = 4096;
115        let module: Module<FFT64> = Module::<FFT64>::new(n);
116        let log_base2k: usize = 17;
117        let size: usize = 5;
118        let mut source: Source = Source::new([0u8; 32]);
119        let cols: usize = 2;
120        let zero: Vec<i64> = vec![0; n];
121        let one_12_sqrt: f64 = 0.28867513459481287;
122        (0..cols).for_each(|col_i| {
123            let mut a: VecZnx = module.new_vec_znx(cols, size);
124            module.fill_uniform(log_base2k, &mut a, col_i, size, &mut source);
125            (0..cols).for_each(|col_j| {
126                if col_j != col_i {
127                    (0..size).for_each(|limb_i| {
128                        assert_eq!(a.at(col_j, limb_i), zero);
129                    })
130                } else {
131                    let std: f64 = a.std(col_i, log_base2k);
132                    assert!(
133                        (std - one_12_sqrt).abs() < 0.01,
134                        "std={} ~!= {}",
135                        std,
136                        one_12_sqrt
137                    );
138                }
139            })
140        });
141    }
142
143    #[test]
144    fn add_normal() {
145        let n: usize = 4096;
146        let module: Module<FFT64> = Module::<FFT64>::new(n);
147        let log_base2k: usize = 17;
148        let log_k: usize = 2 * 17;
149        let size: usize = 5;
150        let sigma: f64 = 3.2;
151        let bound: f64 = 6.0 * sigma;
152        let mut source: Source = Source::new([0u8; 32]);
153        let cols: usize = 2;
154        let zero: Vec<i64> = vec![0; n];
155        let k_f64: f64 = (1u64 << log_k as u64) as f64;
156        (0..cols).for_each(|col_i| {
157            let mut a: VecZnx = module.new_vec_znx(cols, size);
158            module.add_normal(log_base2k, &mut a, col_i, log_k, &mut source, sigma, bound);
159            (0..cols).for_each(|col_j| {
160                if col_j != col_i {
161                    (0..size).for_each(|limb_i| {
162                        assert_eq!(a.at(col_j, limb_i), zero);
163                    })
164                } else {
165                    let std: f64 = a.std(col_i, log_base2k) * k_f64;
166                    assert!((std - sigma).abs() < 0.1, "std={} ~!= {}", std, sigma);
167                }
168            })
169        });
170    }
171}