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

ixa/random/
context_ext.rs
