poulpy/backend/src/sampling.rs
backend/src/sampling.rsRaw
1use crate::znx_base::ZnxViewMut;
2use crate::{FFT64, VecZnx, VecZnxBig, VecZnxBigToMut, VecZnxToMut};
3use rand_distr::{Distribution, Normal};
4use sampling::source::Source;
5
6pub trait FillUniform {
7 /// Fills the first `size` size with uniform values in \[-2^{basek-1}, 2^{basek-1}\]
8 fn fill_uniform(&mut self, basek: usize, col_i: usize, size: usize, source: &mut Source);
9}
10
11pub trait FillDistF64 {
12 fn fill_dist_f64<D: Distribution<f64>>(
13 &mut self,
14 basek: usize,
15 col_i: usize,
16 k: usize,
17 source: &mut Source,
18 dist: D,
19 bound: f64,
20 );
21}
22
23pub trait AddDistF64 {
24 /// Adds vector sampled according to the provided distribution, scaled by 2^{-k} and bounded to \[-bound, bound\].
25 fn add_dist_f64<D: Distribution<f64>>(
26 &mut self,
27 basek: usize,
28 col_i: usize,
29 k: usize,
30 source: &mut Source,
31 dist: D,
32 bound: f64,
33 );
34}
35
36pub trait FillNormal {
37 fn fill_normal(&mut self, basek: usize, col_i: usize, k: usize, source: &mut Source, sigma: f64, bound: f64);
38}
39
40pub trait AddNormal {
41 /// Adds a discrete normal vector scaled by 2^{-k} with the provided standard deviation and bounded to \[-bound, bound\].
42 fn add_normal(&mut self, basek: usize, col_i: usize, k: usize, source: &mut Source, sigma: f64, bound: f64);
43}
44
45impl<T: AsMut<[u8]> + AsRef<[u8]>> FillUniform for VecZnx<T>
46where
47 VecZnx<T>: VecZnxToMut,
48{
49 fn fill_uniform(&mut self, basek: usize, col_i: usize, size: usize, source: &mut Source) {
50 let mut a: VecZnx<&mut [u8]> = self.to_mut();
51 let base2k: u64 = 1 << basek;
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
62impl<T: AsMut<[u8]> + AsRef<[u8]>> FillDistF64 for VecZnx<T>
63where
64 VecZnx<T>: VecZnxToMut,
65{
66 fn fill_dist_f64<D: Distribution<f64>>(
67 &mut self,
68 basek: usize,
69 col_i: usize,
70 k: usize,
71 source: &mut Source,
72 dist: D,
73 bound: f64,
74 ) {
75 let mut a: VecZnx<&mut [u8]> = self.to_mut();
76 assert!(
77 (bound.log2().ceil() as i64) < 64,
78 "invalid bound: ceil(log2(bound))={} > 63",
79 (bound.log2().ceil() as i64)
80 );
81
82 let limb: usize = (k + basek - 1) / basek - 1;
83 let basek_rem: usize = (limb + 1) * basek - k;
84
85 if basek_rem != 0 {
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) << basek_rem;
92 });
93 } else {
94 a.at_mut(col_i, limb).iter_mut().for_each(|a| {
95 let mut dist_f64: f64 = dist.sample(source);
96 while dist_f64.abs() > bound {
97 dist_f64 = dist.sample(source)
98 }
99 *a = dist_f64.round() as i64
100 });
101 }
102 }
103}
104
105impl<T: AsMut<[u8]> + AsRef<[u8]>> AddDistF64 for VecZnx<T>
106where
107 VecZnx<T>: VecZnxToMut,
108{
109 fn add_dist_f64<D: Distribution<f64>>(
110 &mut self,
111 basek: usize,
112 col_i: usize,
113 k: usize,
114 source: &mut Source,
115 dist: D,
116 bound: f64,
117 ) {
118 let mut a: VecZnx<&mut [u8]> = self.to_mut();
119 assert!(
120 (bound.log2().ceil() as i64) < 64,
121 "invalid bound: ceil(log2(bound))={} > 63",
122 (bound.log2().ceil() as i64)
123 );
124
125 let limb: usize = (k + basek - 1) / basek - 1;
126 let basek_rem: usize = (limb + 1) * basek - k;
127
128 if basek_rem != 0 {
129 a.at_mut(col_i, limb).iter_mut().for_each(|a| {
130 let mut dist_f64: f64 = dist.sample(source);
131 while dist_f64.abs() > bound {
132 dist_f64 = dist.sample(source)
133 }
134 *a += (dist_f64.round() as i64) << basek_rem;
135 });
136 } else {
137 a.at_mut(col_i, limb).iter_mut().for_each(|a| {
138 let mut dist_f64: f64 = dist.sample(source);
139 while dist_f64.abs() > bound {
140 dist_f64 = dist.sample(source)
141 }
142 *a += dist_f64.round() as i64
143 });
144 }
145 }
146}
147
148impl<T: AsMut<[u8]> + AsRef<[u8]>> FillNormal for VecZnx<T>
149where
150 VecZnx<T>: VecZnxToMut,
151{
152 fn fill_normal(&mut self, basek: usize, col_i: usize, k: usize, source: &mut Source, sigma: f64, bound: f64) {
153 self.fill_dist_f64(
154 basek,
155 col_i,
156 k,
157 source,
158 Normal::new(0.0, sigma).unwrap(),
159 bound,
160 );
161 }
162}
163
164impl<T: AsMut<[u8]> + AsRef<[u8]>> AddNormal for VecZnx<T>
165where
166 VecZnx<T>: VecZnxToMut,
167{
168 fn add_normal(&mut self, basek: usize, col_i: usize, k: usize, source: &mut Source, sigma: f64, bound: f64) {
169 self.add_dist_f64(
170 basek,
171 col_i,
172 k,
173 source,
174 Normal::new(0.0, sigma).unwrap(),
175 bound,
176 );
177 }
178}
179
180impl<T: AsMut<[u8]> + AsRef<[u8]>> FillDistF64 for VecZnxBig<T, FFT64>
181where
182 VecZnxBig<T, FFT64>: VecZnxBigToMut<FFT64>,
183{
184 fn fill_dist_f64<D: Distribution<f64>>(
185 &mut self,
186 basek: usize,
187 col_i: usize,
188 k: usize,
189 source: &mut Source,
190 dist: D,
191 bound: f64,
192 ) {
193 let mut a: VecZnxBig<&mut [u8], FFT64> = self.to_mut();
194 assert!(
195 (bound.log2().ceil() as i64) < 64,
196 "invalid bound: ceil(log2(bound))={} > 63",
197 (bound.log2().ceil() as i64)
198 );
199
200 let limb: usize = (k + basek - 1) / basek - 1;
201 let basek_rem: usize = (limb + 1) * basek - k;
202
203 if basek_rem != 0 {
204 a.at_mut(col_i, limb).iter_mut().for_each(|a| {
205 let mut dist_f64: f64 = dist.sample(source);
206 while dist_f64.abs() > bound {
207 dist_f64 = dist.sample(source)
208 }
209 *a = (dist_f64.round() as i64) << basek_rem;
210 });
211 } else {
212 a.at_mut(col_i, limb).iter_mut().for_each(|a| {
213 let mut dist_f64: f64 = dist.sample(source);
214 while dist_f64.abs() > bound {
215 dist_f64 = dist.sample(source)
216 }
217 *a = dist_f64.round() as i64
218 });
219 }
220 }
221}
222
223impl<T: AsMut<[u8]> + AsRef<[u8]>> AddDistF64 for VecZnxBig<T, FFT64>
224where
225 VecZnxBig<T, FFT64>: VecZnxBigToMut<FFT64>,
226{
227 fn add_dist_f64<D: Distribution<f64>>(
228 &mut self,
229 basek: usize,
230 col_i: usize,
231 k: usize,
232 source: &mut Source,
233 dist: D,
234 bound: f64,
235 ) {
236 let mut a: VecZnxBig<&mut [u8], FFT64> = self.to_mut();
237 assert!(
238 (bound.log2().ceil() as i64) < 64,
239 "invalid bound: ceil(log2(bound))={} > 63",
240 (bound.log2().ceil() as i64)
241 );
242
243 let limb: usize = (k + basek - 1) / basek - 1;
244 let basek_rem: usize = (limb + 1) * basek - k;
245
246 if basek_rem != 0 {
247 a.at_mut(col_i, limb).iter_mut().for_each(|a| {
248 let mut dist_f64: f64 = dist.sample(source);
249 while dist_f64.abs() > bound {
250 dist_f64 = dist.sample(source)
251 }
252 *a += (dist_f64.round() as i64) << basek_rem;
253 });
254 } else {
255 a.at_mut(col_i, limb).iter_mut().for_each(|a| {
256 let mut dist_f64: f64 = dist.sample(source);
257 while dist_f64.abs() > bound {
258 dist_f64 = dist.sample(source)
259 }
260 *a += dist_f64.round() as i64
261 });
262 }
263 }
264}
265
266impl<T: AsMut<[u8]> + AsRef<[u8]>> FillNormal for VecZnxBig<T, FFT64>
267where
268 VecZnxBig<T, FFT64>: VecZnxBigToMut<FFT64>,
269{
270 fn fill_normal(&mut self, basek: usize, col_i: usize, k: usize, source: &mut Source, sigma: f64, bound: f64) {
271 self.fill_dist_f64(
272 basek,
273 col_i,
274 k,
275 source,
276 Normal::new(0.0, sigma).unwrap(),
277 bound,
278 );
279 }
280}
281
282impl<T: AsMut<[u8]> + AsRef<[u8]>> AddNormal for VecZnxBig<T, FFT64>
283where
284 VecZnxBig<T, FFT64>: VecZnxBigToMut<FFT64>,
285{
286 fn add_normal(&mut self, basek: usize, col_i: usize, k: usize, source: &mut Source, sigma: f64, bound: f64) {
287 self.add_dist_f64(
288 basek,
289 col_i,
290 k,
291 source,
292 Normal::new(0.0, sigma).unwrap(),
293 bound,
294 );
295 }
296}
297
298#[cfg(test)]
299mod tests {
300 use super::{AddNormal, FillUniform};
301 use crate::vec_znx_ops::*;
302 use crate::znx_base::*;
303 use crate::{FFT64, Module, Stats, VecZnx};
304 use sampling::source::Source;
305
306 #[test]
307 fn vec_znx_fill_uniform() {
308 let n: usize = 4096;
309 let module: Module<FFT64> = Module::<FFT64>::new(n);
310 let basek: usize = 17;
311 let size: usize = 5;
312 let mut source: Source = Source::new([0u8; 32]);
313 let cols: usize = 2;
314 let zero: Vec<i64> = vec![0; n];
315 let one_12_sqrt: f64 = 0.28867513459481287;
316 (0..cols).for_each(|col_i| {
317 let mut a: VecZnx<_> = module.new_vec_znx(cols, size);
318 a.fill_uniform(basek, col_i, size, &mut source);
319 (0..cols).for_each(|col_j| {
320 if col_j != col_i {
321 (0..size).for_each(|limb_i| {
322 assert_eq!(a.at(col_j, limb_i), zero);
323 })
324 } else {
325 let std: f64 = a.std(col_i, basek);
326 assert!(
327 (std - one_12_sqrt).abs() < 0.01,
328 "std={} ~!= {}",
329 std,
330 one_12_sqrt
331 );
332 }
333 })
334 });
335 }
336
337 #[test]
338 fn vec_znx_add_normal() {
339 let n: usize = 4096;
340 let module: Module<FFT64> = Module::<FFT64>::new(n);
341 let basek: usize = 17;
342 let k: usize = 2 * 17;
343 let size: usize = 5;
344 let sigma: f64 = 3.2;
345 let bound: f64 = 6.0 * sigma;
346 let mut source: Source = Source::new([0u8; 32]);
347 let cols: usize = 2;
348 let zero: Vec<i64> = vec![0; n];
349 let k_f64: f64 = (1u64 << k as u64) as f64;
350 (0..cols).for_each(|col_i| {
351 let mut a: VecZnx<_> = module.new_vec_znx(cols, size);
352 a.add_normal(basek, col_i, k, &mut source, sigma, bound);
353 (0..cols).for_each(|col_j| {
354 if col_j != col_i {
355 (0..size).for_each(|limb_i| {
356 assert_eq!(a.at(col_j, limb_i), zero);
357 })
358 } else {
359 let std: f64 = a.std(col_i, basek) * k_f64;
360 assert!((std - sigma).abs() < 0.1, "std={} ~!= {}", std, sigma);
361 }
362 })
363 });
364 }
365}