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