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