poulpy/base2k/src/sampling.rs
base2k/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^{log_base2k-1}, 2^{log_base2k-1}\]
8 fn fill_uniform(&mut self, log_base2k: 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 log_base2k: usize,
15 col_i: usize,
16 log_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^{-log_k} and bounded to \[-bound, bound\].
25 fn add_dist_f64<D: Distribution<f64>>(
26 &mut self,
27 log_base2k: usize,
28 col_i: usize,
29 log_k: usize,
30 source: &mut Source,
31 dist: D,
32 bound: f64,
33 );
34}
35
36pub trait FillNormal {
37 fn fill_normal(&mut self, log_base2k: usize, col_i: usize, log_k: usize, source: &mut Source, sigma: f64, bound: f64);
38}
39
40pub trait AddNormal {
41 /// Adds a discrete normal vector scaled by 2^{-log_k} with the provided standard deviation and bounded to \[-bound, bound\].
42 fn add_normal(&mut self, log_base2k: usize, col_i: usize, log_k: usize, source: &mut Source, sigma: f64, bound: f64);
43}
44
45impl<T> FillUniform for VecZnx<T>
46where
47 VecZnx<T>: VecZnxToMut,
48{
49 fn fill_uniform(&mut self, log_base2k: usize, col_i: usize, size: usize, source: &mut Source) {
50 let mut a: VecZnx<&mut [u8]> = self.to_mut();
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
62impl<T> FillDistF64 for VecZnx<T>
63where
64 VecZnx<T>: VecZnxToMut,
65{
66 fn fill_dist_f64<D: Distribution<f64>>(
67 &mut self,
68 log_base2k: usize,
69 col_i: usize,
70 log_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 = (log_k + log_base2k - 1) / log_base2k - 1;
83 let log_base2k_rem: usize = log_k % log_base2k;
84
85 if log_base2k_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) << log_base2k_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> AddDistF64 for VecZnx<T>
106where
107 VecZnx<T>: VecZnxToMut,
108{
109 fn add_dist_f64<D: Distribution<f64>>(
110 &mut self,
111 log_base2k: usize,
112 col_i: usize,
113 log_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 = (log_k + log_base2k - 1) / log_base2k - 1;
126 let log_base2k_rem: usize = log_k % log_base2k;
127
128 if log_base2k_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) << log_base2k_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> FillNormal for VecZnx<T>
149where
150 VecZnx<T>: VecZnxToMut,
151{
152 fn fill_normal(&mut self, log_base2k: usize, col_i: usize, log_k: usize, source: &mut Source, sigma: f64, bound: f64) {
153 self.fill_dist_f64(
154 log_base2k,
155 col_i,
156 log_k,
157 source,
158 Normal::new(0.0, sigma).unwrap(),
159 bound,
160 );
161 }
162}
163
164impl<T> AddNormal for VecZnx<T>
165where
166 VecZnx<T>: VecZnxToMut,
167{
168 fn add_normal(&mut self, log_base2k: usize, col_i: usize, log_k: usize, source: &mut Source, sigma: f64, bound: f64) {
169 self.add_dist_f64(
170 log_base2k,
171 col_i,
172 log_k,
173 source,
174 Normal::new(0.0, sigma).unwrap(),
175 bound,
176 );
177 }
178}
179
180impl<T> 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 log_base2k: usize,
187 col_i: usize,
188 log_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 = (log_k + log_base2k - 1) / log_base2k - 1;
201 let log_base2k_rem: usize = log_k % log_base2k;
202
203 if log_base2k_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) << log_base2k_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> 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 log_base2k: usize,
230 col_i: usize,
231 log_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 = (log_k + log_base2k - 1) / log_base2k - 1;
244 let log_base2k_rem: usize = log_k % log_base2k;
245
246 if log_base2k_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) << log_base2k_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> FillNormal for VecZnxBig<T, FFT64>
267where
268 VecZnxBig<T, FFT64>: VecZnxBigToMut<FFT64>,
269{
270 fn fill_normal(&mut self, log_base2k: usize, col_i: usize, log_k: usize, source: &mut Source, sigma: f64, bound: f64) {
271 self.fill_dist_f64(
272 log_base2k,
273 col_i,
274 log_k,
275 source,
276 Normal::new(0.0, sigma).unwrap(),
277 bound,
278 );
279 }
280}
281
282impl<T> AddNormal for VecZnxBig<T, FFT64>
283where
284 VecZnxBig<T, FFT64>: VecZnxBigToMut<FFT64>,
285{
286 fn add_normal(&mut self, log_base2k: usize, col_i: usize, log_k: usize, source: &mut Source, sigma: f64, bound: f64) {
287 self.add_dist_f64(
288 log_base2k,
289 col_i,
290 log_k,
291 source,
292 Normal::new(0.0, sigma).unwrap(),
293 bound,
294 );
295 }
296}
297
298#[cfg(test)]
299mod tests {
300 use std::fmt::Display;
301
302 use super::{AddNormal, FillUniform};
303 use crate::vec_znx_ops::*;
304 use crate::znx_base::*;
305 use crate::{FFT64, Module, Stats, VecZnx};
306 use sampling::source::Source;
307
308 #[test]
309 fn vec_znx_fill_uniform() {
310 let n: usize = 4096;
311 let module: Module<FFT64> = Module::<FFT64>::new(n);
312 let log_base2k: usize = 17;
313 let size: usize = 5;
314 let mut source: Source = Source::new([0u8; 32]);
315 let cols: usize = 2;
316 let zero: Vec<i64> = vec![0; n];
317 let one_12_sqrt: f64 = 0.28867513459481287;
318 (0..cols).for_each(|col_i| {
319 let mut a: VecZnx<_> = module.new_vec_znx(cols, size);
320 a.fill_uniform(log_base2k, col_i, size, &mut source);
321 (0..cols).for_each(|col_j| {
322 if col_j != col_i {
323 (0..size).for_each(|limb_i| {
324 assert_eq!(a.at(col_j, limb_i), zero);
325 })
326 } else {
327 let std: f64 = a.std(col_i, log_base2k);
328 assert!(
329 (std - one_12_sqrt).abs() < 0.01,
330 "std={} ~!= {}",
331 std,
332 one_12_sqrt
333 );
334 }
335 })
336 });
337 }
338
339 #[test]
340 fn vec_znx_add_normal() {
341 let n: usize = 4096;
342 let module: Module<FFT64> = Module::<FFT64>::new(n);
343 let log_base2k: usize = 17;
344 let log_k: usize = 2 * 17;
345 let size: usize = 5;
346 let sigma: f64 = 3.2;
347 let bound: f64 = 6.0 * sigma;
348 let mut source: Source = Source::new([0u8; 32]);
349 let cols: usize = 2;
350 let zero: Vec<i64> = vec![0; n];
351 let k_f64: f64 = (1u64 << log_k as u64) as f64;
352 (0..cols).for_each(|col_i| {
353 let mut a: VecZnx<_> = module.new_vec_znx(cols, size);
354 a.add_normal(log_base2k, col_i, log_k, &mut source, sigma, bound);
355 (0..cols).for_each(|col_j| {
356 if col_j != col_i {
357 (0..size).for_each(|limb_i| {
358 assert_eq!(a.at(col_j, limb_i), zero);
359 })
360 } else {
361 let std: f64 = a.std(col_i, log_base2k) * k_f64;
362 assert!((std - sigma).abs() < 0.1, "std={} ~!= {}", std, sigma);
363 }
364 })
365 });
366 }
367}