qfall_math/utils/sample/binomial.rs
1// Copyright 2025 Jan Niklas Siemer
2//
3// This file is part of qFALL-math.
4//
5// qFALL-math is free software: you can redistribute it and/or modify it under
6// the terms of the Mozilla Public License Version 2.0 as published by the
7// Mozilla Foundation. See <https://mozilla.org/en-US/MPL/2.0/>.
8
9//! This module includes core functionality to sample according to the
10//! binomial distribution.
11
12use super::uniform::SamplerRng;
13use crate::{error::MathError, integer::Z, rational::Q};
14use rand_distr::{Binomial, Distribution};
15
16/// Enables sampling a [`Z`] according to the binomial distribution `Bin(n, p)`.
17///
18/// Attributes:
19/// - `distr`: defines the binomial distribution with parameters `n` and `p` to sample from
20/// - `rng`: defines the [`ThreadRng`](rand::rngs::ThreadRng) or seeded
21/// [`StdRng`](rand::rngs::StdRng) that's used to sample from `distr`
22///
23/// # Examples
24/// ```
25/// use qfall_math::utils::sample::binomial::BinomialSampler;
26/// let n = 2;
27/// let p = 0.5;
28///
29/// let mut bin_sampler = BinomialSampler::init(n, p, None).unwrap();
30///
31/// let sample = bin_sampler.sample();
32///
33/// assert!(0 <= sample);
34/// assert!(sample <= n);
35/// ```
36pub struct BinomialSampler {
37 pub distr: Binomial,
38 pub rng: SamplerRng,
39}
40
41impl BinomialSampler {
42 /// Initializes the [`BinomialSampler`] with
43 /// - `distr` as the binomial distribution with `n` tries and success probability `p` for each try, and
44 /// - `rng` as a [`StdRng`](rand::rngs::StdRng) seeded with `seed` if `seed` is provided,
45 /// and as a fresh [`ThreadRng`](rand::rngs::ThreadRng) otherwise.
46 ///
47 /// Parameters:
48 /// - `n`: specifies the number of tries
49 /// - `p`: specifies the success probability
50 /// - `seed`: specifies an optional 256-bit seed for the internal [`StdRng`](rand::rngs::StdRng).
51 /// If `None` is provided, a fresh [`ThreadRng`](rand::rngs::ThreadRng) is used instead.
52 ///
53 /// Returns a [`BinomialSampler`] or a [`MathError`] if `n < 0`,
54 /// `p ∉ (0,1)`, or `n` does not fit into an [`i64`].
55 ///
56 /// # Examples
57 /// ```
58 /// use qfall_math::utils::sample::binomial::BinomialSampler;
59 /// let n = 2;
60 /// let p = 0.5;
61 ///
62 /// let bin_sampler = BinomialSampler::init(n, p, None).unwrap();
63 ///
64 /// let mut bin_sampler_seeded_0 = BinomialSampler::init(n, p, Some([42; 32])).unwrap();
65 /// let mut bin_sampler_seeded_1 = BinomialSampler::init(n, p, Some([42; 32])).unwrap();
66 /// assert_eq!(bin_sampler_seeded_0.sample(), bin_sampler_seeded_1.sample());
67 /// ```
68 ///
69 /// # Errors and Failures
70 /// - Returns a [`MathError`] of type [`InvalidIntegerInput`](MathError::InvalidIntegerInput)
71 /// if `n < 0`.
72 /// - Returns a [`MathError`] of type [`InvalidInterval`](MathError::InvalidInterval)
73 /// if `p ∉ (0,1)`.
74 /// - Returns a [`MathError`] of type [`ConversionError`](MathError::ConversionError)
75 /// if `n` does not fit into an [`i64`].
76 pub fn init(
77 n: impl Into<Z>,
78 p: impl Into<Q>,
79 seed: Option<[u8; 32]>,
80 ) -> Result<Self, MathError> {
81 let n = n.into();
82 let p = p.into();
83
84 if p <= Q::ZERO || p >= Q::ONE {
85 return Err(MathError::InvalidInterval(format!(
86 "p (the probability of success for binomial sampling) must be chosen between 0 and 1. \
87 Currently it is {p}. \
88 Hence, the interval to sample from is invalid and contains only exactly one number."
89 )));
90 }
91 if n < Z::ZERO {
92 return Err(MathError::InvalidIntegerInput(format!(
93 "n (the number of trials for binomial sampling) must be no smaller than 0. Currently it is {n}."
94 )));
95 }
96
97 let n = i64::try_from(n)? as u64;
98 let p = f64::from(&p);
99
100 let distr = Binomial::new(n, p).unwrap();
101 let rng = SamplerRng::new(seed);
102
103 Ok(Self { distr, rng })
104 }
105
106 /// Samples a [`Z`] according to the binomial distribution `Bin(n, p)`.
107 ///
108 /// # Examples
109 /// ```
110 /// use qfall_math::utils::sample::binomial::BinomialSampler;
111 /// let n = 2;
112 /// let p = 0.5;
113 ///
114 /// let mut bin_sampler = BinomialSampler::init(n, p, None).unwrap();
115 ///
116 /// let sample = bin_sampler.sample();
117 ///
118 /// assert!(0 <= sample);
119 /// assert!(sample <= n);
120 /// ```
121 pub fn sample(&mut self) -> Z {
122 Z::from(self.distr.sample(&mut self.rng))
123 }
124}
125
126#[cfg(test)]
127mod test_binomial_sampler {
128 use super::{BinomialSampler, Q, Z};
129 use crate::utils::sample::test_seed;
130
131 /// Checks whether the range is kept,
132 /// i.e. if any result is at least 0 and smaller than or equal to `n`.
133 #[test]
134 fn keeps_range() {
135 for seeded in [false, true] {
136 let n = 16;
137 let p = 0.5;
138 let mut bin_sampler = BinomialSampler::init(n, p, seeded.then(test_seed)).unwrap();
139
140 for _ in 0..16 {
141 let sample = bin_sampler.sample();
142 // sample >= 0 check is not required as sample is a u64
143 assert!(sample <= n);
144 }
145 }
146 }
147
148 /// Roughly checks that the collected samples are distributed according to the binomial distribution.
149 #[test]
150 fn distribution() {
151 for seeded in [false, true] {
152 let n = 2;
153 let p = 0.5;
154 let mut bin_sampler = BinomialSampler::init(n, p, seeded.then(test_seed)).unwrap();
155
156 let mut counts = [0; 3];
157 // count sampled instances
158 for _ in 0..1000 {
159 let sample = u64::try_from(bin_sampler.sample()).unwrap() as usize;
160 counts[sample] += 1;
161 }
162
163 let expl_text = String::from("This test can fail with probability close to 0.
164 It fails if the sampled occurrences do not look like a typical binomial random distribution.
165 If this happens, rerun the tests several times and check whether this issue comes up again.");
166
167 // Check that the sampled occurrences roughly look
168 // like a binomial distribution
169 assert!(counts[0] > 100, "{expl_text}");
170 assert!(counts[0] < 400, "{expl_text}");
171 assert!(counts[1] > 300, "{expl_text}");
172 assert!(counts[1] < 700, "{expl_text}");
173 assert!(counts[2] > 100, "{expl_text}");
174 assert!(counts[2] < 400, "{expl_text}");
175 }
176 }
177
178 /// Checks whether invalid choices for n result in an error.
179 #[test]
180 fn invalid_n() {
181 for seeded in [false, true] {
182 let p = 0.5;
183
184 assert!(BinomialSampler::init(&Z::MINUS_ONE, p, seeded.then(test_seed)).is_err());
185 assert!(BinomialSampler::init(Z::from(i64::MIN), p, seeded.then(test_seed)).is_err());
186 }
187 }
188
189 /// Checks whether invalid choices for p result in an error.
190 #[test]
191 fn invalid_p() {
192 for seeded in [false, true] {
193 let n = 2;
194
195 assert!(BinomialSampler::init(n, &Q::MINUS_ONE, seeded.then(test_seed)).is_err());
196 assert!(BinomialSampler::init(n, &Q::ZERO, seeded.then(test_seed)).is_err());
197 assert!(BinomialSampler::init(n, &Q::ONE, seeded.then(test_seed)).is_err());
198 assert!(BinomialSampler::init(n, Q::from(5), seeded.then(test_seed)).is_err());
199 }
200 }
201
202 /// Checks whether two samplers with the same seed output the same samples.
203 #[test]
204 fn same_seed_same_samples() {
205 let n = 1024;
206 let p = 0.5;
207 let mut bin_sampler_0 = BinomialSampler::init(n, p, Some([42; 32])).unwrap();
208 let mut bin_sampler_1 = BinomialSampler::init(n, p, Some([42; 32])).unwrap();
209
210 for _ in 0..u8::MAX {
211 assert_eq!(bin_sampler_0.sample(), bin_sampler_1.sample());
212 }
213 }
214
215 /// Checks whether the seeded sampler outputs the known samples for a fixed seed.
216 #[test]
217 fn known_answer() {
218 let mut bin_sampler = BinomialSampler::init(1024, 0.5, Some([42; 32])).unwrap();
219
220 let samples: Vec<String> = (0..8).map(|_| bin_sampler.sample().to_string()).collect();
221
222 assert_eq!("495, 492, 500, 489, 491, 510, 489, 555", samples.join(", "));
223 }
224}