ixa/random/
context_ext.rs

1use std::any::TypeId;
2use std::cell::RefMut;
3use std::hash::Hasher;
4
5use log::trace;
6
7use crate::hashing::{hash_str, DeterministicHasher};
8use crate::rand::distr::uniform::{SampleRange, SampleUniform};
9use crate::rand::distr::weighted::{Weight, WeightedIndex};
10use crate::rand::distr::Distribution;
11use crate::rand::{Rng, RngExt, SeedableRng};
12use crate::random::{RngHolder, RngPlugin};
13use crate::{Context, ContextBase, RngId};
14
15/// Gets a mutable reference to the random number generator associated with the given
16/// [`RngId`]. If the Rng has not been used before, one will be created with the base seed
17/// you defined in `init`. Note that this will panic if `init` was not called yet.
18fn get_rng<R: RngId + 'static>(context: &impl ContextBase) -> RefMut<R::RngType> {
19    let data_container = context.get_data(RngPlugin);
20
21    let rng_holders = data_container.rng_holders.try_borrow_mut().unwrap();
22    RefMut::map(rng_holders, |holders| {
23        holders
24            .entry(TypeId::of::<R>())
25            // Create a new rng holder if it doesn't exist yet
26            .or_insert_with(|| {
27                trace!(
28                    "creating new RNG (seed={}) for type id {:?}",
29                    data_container.base_seed,
30                    TypeId::of::<R>()
31                );
32                let base_seed = data_container.base_seed;
33                let seed_offset = hash_str(R::get_name());
34                RngHolder {
35                    rng: Box::new(R::RngType::seed_from_u64(
36                        base_seed.wrapping_add(seed_offset),
37                    )),
38                }
39            })
40            .rng
41            .downcast_mut::<R::RngType>()
42            .unwrap()
43    })
44}
45
46// This is a trait extension on Context for
47// random number generation functionality.
48pub trait ContextRandomExt: ContextBase {
49    /// Initializes the `RngPlugin` data container to store rngs as well as a base
50    /// seed. Note that rngs are created lazily when `get_rng` is called.
51    fn init_random(&mut self, base_seed: u64) {
52        trace!("initializing random module");
53        let data_container = self.get_data_mut(RngPlugin);
54        data_container.base_seed = base_seed;
55
56        // Clear any existing Rngs to ensure they get re-seeded when `get_rng` is called
57        let mut rng_map = data_container.rng_holders.try_borrow_mut().unwrap();
58        rng_map.clear();
59    }
60
61    /// Gets a random sample from the random number generator associated with the given
62    /// [`RngId`] by applying the specified sampler function. If the Rng has not been used
63    /// before, one will be created with the base seed you defined in `set_base_random_seed`.
64    /// Note that this will panic if `set_base_random_seed` was not called yet.
65    #[must_use]
66    fn sample<R: RngId + 'static, T>(
67        &self,
68        _rng_type: R,
69        sampler: impl FnOnce(&mut R::RngType) -> T,
70    ) -> T {
71        let mut rng = get_rng::<R>(self);
72        sampler(&mut rng)
73    }
74
75    /// Returns a deterministic debug fingerprint for the current position of the random number
76    /// generator associated with the given [`RngId`].
77    ///
78    /// This clones the current generator and hashes a small number of draws from the clone, so it
79    /// does not advance the generator used by sampling. The returned value is intended only for
80    /// comparing checkpoints across runs of the same code and dependency versions. It is not a
81    /// serialization format, and it does not guarantee reproducibility across RNG algorithm or
82    /// version changes.
83    #[must_use]
84    fn debug_rng_state<R: RngId + 'static>(&self, _rng_id: R) -> u64
85    where
86        R::RngType: Clone + Rng,
87    {
88        let rng = get_rng::<R>(self);
89        let mut rng_snapshot = (*rng).clone();
90        let mut hasher = DeterministicHasher::default();
91
92        for _ in 0..4 {
93            hasher.write(&rng_snapshot.next_u64().to_le_bytes());
94        }
95
96        hasher.finish()
97    }
98
99    /// Gets a random sample from the specified distribution using a random number generator
100    /// associated with the given [`RngId`]. If the Rng has not been used before, one will be
101    /// created with the base seed you defined in `set_base_random_seed`.
102    /// Note that this will panic if `set_base_random_seed` was not called yet.
103    #[must_use]
104    fn sample_distr<R: RngId + 'static, T>(
105        &self,
106        _rng_type: R,
107        distribution: impl Distribution<T>,
108    ) -> T
109    where
110        R::RngType: Rng,
111    {
112        let mut rng = get_rng::<R>(self);
113        distribution.sample::<R::RngType>(&mut rng)
114    }
115
116    /// Gets a random sample within the range provided by `range`
117    /// using the generator associated with the given [`RngId`].
118    /// Note that this will panic if `set_base_random_seed` was not called yet.
119    #[must_use]
120    fn sample_range<R: RngId + 'static, S, T>(&self, rng_id: R, range: S) -> T
121    where
122        R::RngType: Rng,
123        S: SampleRange<T>,
124        T: SampleUniform,
125    {
126        self.sample(rng_id, |rng| rng.random_range(range))
127    }
128
129    /// Gets a random boolean value which is true with probability `p` using the
130    /// generator associated with the given [`RngId`]. The supplied probability
131    /// is converted to `f64` before sampling.
132    ///
133    /// ```
134    /// use ixa::{define_rng, Context, ContextRandomExt};
135    ///
136    /// define_rng!(ExampleRng);
137    ///
138    /// struct Probability(f64);
139    ///
140    /// impl From<Probability> for f64 {
141    ///     fn from(probability: Probability) -> Self {
142    ///         probability.0
143    ///     }
144    /// }
145    ///
146    /// let mut context = Context::new();
147    /// context.init_random(42);
148    /// assert!(context.sample_bool(ExampleRng, Probability(1.0)));
149    /// ```
150    ///
151    /// # Panics
152    ///
153    /// Panics if the converted probability is outside `[0.0, 1.0]` or is NaN.
154    #[must_use]
155    fn sample_bool<R: RngId + 'static>(&self, rng_id: R, p: impl Into<f64>) -> bool
156    where
157        R::RngType: Rng,
158    {
159        let p = p.into();
160        self.sample(rng_id, |rng| rng.random_bool(p))
161    }
162
163    /// Draws a random entry out of the list provided in `weights`
164    /// with the given weights using the generator associated with the
165    /// given [`RngId`].  Note that this will panic if
166    /// `set_base_random_seed` was not called yet.
167    #[must_use]
168    fn sample_weighted<R: RngId + 'static, T>(&self, _rng_id: R, weights: &[T]) -> usize
169    where
170        R::RngType: Rng,
171        T: Clone
172            + Default
173            + SampleUniform
174            + for<'a> std::ops::AddAssign<&'a T>
175            + PartialOrd
176            + Weight,
177    {
178        let index = WeightedIndex::new(weights).unwrap();
179        let mut rng = get_rng::<R>(self);
180        index.sample(&mut *rng)
181    }
182}
183
184impl ContextRandomExt for Context {}
185
186#[cfg(test)]
187mod test {
188    use crate::context::Context;
189    use crate::rand::distr::weighted::WeightedIndex;
190    use crate::rand::distr::Distribution;
191    use crate::rand::Rng;
192    use crate::random::context_ext::ContextRandomExt;
193    use crate::{define_data_plugin, define_rng};
194
195    define_rng!(FooRng);
196    define_rng!(BarRng);
197
198    struct Probability(f64);
199
200    impl From<Probability> for f64 {
201        fn from(probability: Probability) -> Self {
202            probability.0
203        }
204    }
205
206    #[test]
207    fn get_rng_basic() {
208        let mut context = Context::new();
209        context.init_random(42);
210
211        assert_ne!(
212            context.sample(FooRng, Rng::next_u64),
213            context.sample(FooRng, Rng::next_u64)
214        );
215    }
216
217    #[test]
218    fn multiple_rng_types() {
219        let mut context = Context::new();
220        context.init_random(42);
221
222        assert_ne!(
223            context.sample(FooRng, Rng::next_u64),
224            context.sample(BarRng, Rng::next_u64)
225        );
226    }
227
228    #[test]
229    fn reset_seed() {
230        let mut context = Context::new();
231        context.init_random(42);
232
233        let run_0 = context.sample(FooRng, Rng::next_u64);
234        let run_1 = context.sample(FooRng, Rng::next_u64);
235
236        // Reset with same seed, ensure we get the same values
237        context.init_random(42);
238        assert_eq!(run_0, context.sample(FooRng, Rng::next_u64));
239        assert_eq!(run_1, context.sample(FooRng, Rng::next_u64));
240
241        // Reset with different seed, ensure we get different values
242        context.init_random(88);
243        assert_ne!(run_0, context.sample(FooRng, Rng::next_u64));
244        assert_ne!(run_1, context.sample(FooRng, Rng::next_u64));
245    }
246
247    #[test]
248    fn debug_rng_state_matches_for_same_seed_and_progress() {
249        let mut context_0 = Context::new();
250        context_0.init_random(42);
251        let mut context_1 = Context::new();
252        context_1.init_random(42);
253
254        for _ in 0..3 {
255            let _ = context_0.sample(FooRng, Rng::next_u64);
256            let _ = context_1.sample(FooRng, Rng::next_u64);
257        }
258
259        assert_eq!(
260            context_0.debug_rng_state(FooRng),
261            context_1.debug_rng_state(FooRng)
262        );
263    }
264
265    #[test]
266    fn debug_rng_state_changes_with_rng_progress() {
267        let mut context = Context::new();
268        context.init_random(42);
269
270        let initial = context.debug_rng_state(FooRng);
271        let _ = context.sample(FooRng, Rng::next_u64);
272
273        assert_ne!(initial, context.debug_rng_state(FooRng));
274    }
275
276    #[test]
277    fn debug_rng_state_is_stable_without_sampling() {
278        let mut context = Context::new();
279        context.init_random(42);
280
281        assert_eq!(
282            context.debug_rng_state(FooRng),
283            context.debug_rng_state(FooRng)
284        );
285    }
286
287    #[test]
288    fn debug_rng_state_does_not_affect_next_sample() {
289        let mut with_debug = Context::new();
290        with_debug.init_random(42);
291        let mut without_debug = Context::new();
292        without_debug.init_random(42);
293
294        let _ = with_debug.debug_rng_state(FooRng);
295
296        assert_eq!(
297            with_debug.sample(FooRng, Rng::next_u64),
298            without_debug.sample(FooRng, Rng::next_u64)
299        );
300    }
301
302    #[test]
303    fn debug_rng_state_resets_with_seed() {
304        let mut context = Context::new();
305        context.init_random(42);
306
307        let initial = context.debug_rng_state(FooRng);
308        let _ = context.sample(FooRng, Rng::next_u64);
309        assert_ne!(initial, context.debug_rng_state(FooRng));
310
311        context.init_random(42);
312        assert_eq!(initial, context.debug_rng_state(FooRng));
313    }
314
315    #[test]
316    fn debug_rng_state_is_independent_by_rng_id() {
317        let mut context_0 = Context::new();
318        context_0.init_random(42);
319        let mut context_1 = Context::new();
320        context_1.init_random(42);
321
322        let foo_initial = context_0.debug_rng_state(FooRng);
323        let bar_initial = context_0.debug_rng_state(BarRng);
324        assert_ne!(foo_initial, bar_initial);
325
326        let _ = context_0.sample(FooRng, Rng::next_u64);
327        let _ = context_1.sample(BarRng, Rng::next_u64);
328
329        assert_ne!(context_0.debug_rng_state(FooRng), foo_initial);
330        assert_eq!(context_0.debug_rng_state(BarRng), bar_initial);
331        assert_eq!(context_1.debug_rng_state(FooRng), foo_initial);
332        assert_ne!(context_1.debug_rng_state(BarRng), bar_initial);
333    }
334
335    define_data_plugin!(
336        SamplerData,
337        WeightedIndex<f64>,
338        WeightedIndex::new(vec![1.0]).unwrap()
339    );
340
341    #[test]
342    fn sampler_function_closure_capture() {
343        let mut context = Context::new();
344        context.init_random(42);
345
346        // Initialize weighted sampler. Zero is selected with probability 1/3, one with a
347        // probability of 2/3.
348        *context.get_data_mut(SamplerData) = WeightedIndex::new(vec![1.0, 2.0]).unwrap();
349
350        let parameters = context.get_data(SamplerData);
351        let n_samples = 3000;
352        let mut zero_counter = 0;
353        for _ in 0..n_samples {
354            let sample = context.sample(FooRng, |rng| parameters.sample(rng));
355            if sample == 0 {
356                zero_counter += 1;
357            }
358        }
359        // The expected value of `zero_counter` is 1000.
360        assert!((zero_counter - 1000_i32).abs() < 100);
361    }
362
363    #[test]
364    fn sample_distribution() {
365        let mut context = Context::new();
366        context.init_random(42);
367
368        // Initialize weighted sampler. Zero is selected with probability 1/3, one with a
369        // probability of 2/3.
370        *context.get_data_mut(SamplerData) = WeightedIndex::new(vec![1.0, 2.0]).unwrap();
371
372        let parameters = context.get_data(SamplerData);
373        let n_samples = 3000;
374        let mut zero_counter = 0;
375        for _ in 0..n_samples {
376            let sample = context.sample_distr(FooRng, parameters);
377            if sample == 0 {
378                zero_counter += 1;
379            }
380        }
381        // The expected value of `zero_counter` is 1000.
382        assert!((zero_counter - 1000_i32).abs() < 100);
383    }
384
385    #[test]
386    fn sample_range() {
387        let mut context = Context::new();
388        context.init_random(42);
389        let result = context.sample_range(FooRng, 0..10);
390        assert!((0..10).contains(&result));
391    }
392
393    #[test]
394    fn sample_bool() {
395        let mut context = Context::new();
396        context.init_random(42);
397        let _r: bool = context.sample_bool(FooRng, 0.5);
398    }
399
400    #[test]
401    fn sample_bool_accepts_into_f64_without_extra_turbofish() {
402        let mut context = Context::new();
403        context.init_random(42);
404
405        assert!(!context.sample_bool::<FooRng>(FooRng, Probability(0.0)));
406        assert!(context.sample_bool::<FooRng>(FooRng, Probability(1.0)));
407    }
408
409    #[test]
410    #[should_panic(expected = "outside range [0.0, 1.0]")]
411    fn sample_bool_rejects_wrapped_invalid_probability() {
412        let mut context = Context::new();
413        context.init_random(42);
414
415        let _ = context.sample_bool(FooRng, Probability(1.1));
416    }
417
418    #[test]
419    fn sample_weighted() {
420        let mut context = Context::new();
421        context.init_random(42);
422        let r: usize = context.sample_weighted(FooRng, &[0.1, 0.3, 0.4]);
423        assert!(r < 3);
424    }
425}
ixa/random/context_ext.rs

ixa/random/
context_ext.rs
