ixa/random/
context_ext.rs

1use std::any::TypeId;
2use std::cell::RefMut;
3
4use log::trace;
5
6use crate::hashing::hash_str;
7use crate::rand::distr::uniform::{SampleRange, SampleUniform};
8use crate::rand::distr::weighted::{Weight, WeightedIndex};
9use crate::rand::distr::Distribution;
10use crate::rand::{Rng, SeedableRng};
11use crate::random::{RngHolder, RngPlugin};
12use crate::{Context, ContextBase, RngId};
13
14/// Gets a mutable reference to the random number generator associated with the given
15/// [`RngId`]. If the Rng has not been used before, one will be created with the base seed
16/// you defined in `init`. Note that this will panic if `init` was not called yet.
17fn get_rng<R: RngId + 'static>(context: &impl ContextBase) -> RefMut<R::RngType> {
18    let data_container = context.get_data(RngPlugin);
19
20    let rng_holders = data_container.rng_holders.try_borrow_mut().unwrap();
21    RefMut::map(rng_holders, |holders| {
22        holders
23            .entry(TypeId::of::<R>())
24            // Create a new rng holder if it doesn't exist yet
25            .or_insert_with(|| {
26                trace!(
27                    "creating new RNG (seed={}) for type id {:?}",
28                    data_container.base_seed,
29                    TypeId::of::<R>()
30                );
31                let base_seed = data_container.base_seed;
32                let seed_offset = hash_str(R::get_name());
33                RngHolder {
34                    rng: Box::new(R::RngType::seed_from_u64(
35                        base_seed.wrapping_add(seed_offset),
36                    )),
37                }
38            })
39            .rng
40            .downcast_mut::<R::RngType>()
41            .unwrap()
42    })
43}
44
45// This is a trait extension on Context for
46// random number generation functionality.
47pub trait ContextRandomExt: ContextBase {
48    /// Initializes the `RngPlugin` data container to store rngs as well as a base
49    /// seed. Note that rngs are created lazily when `get_rng` is called.
50    fn init_random(&mut self, base_seed: u64) {
51        trace!("initializing random module");
52        let data_container = self.get_data_mut(RngPlugin);
53        data_container.base_seed = base_seed;
54
55        // Clear any existing Rngs to ensure they get re-seeded when `get_rng` is called
56        let mut rng_map = data_container.rng_holders.try_borrow_mut().unwrap();
57        rng_map.clear();
58    }
59
60    /// Gets a random sample from the random number generator associated with the given
61    /// [`RngId`] by applying the specified sampler function. If the Rng has not been used
62    /// before, one will be created with the base seed you defined in `set_base_random_seed`.
63    /// Note that this will panic if `set_base_random_seed` was not called yet.
64    #[must_use]
65    fn sample<R: RngId + 'static, T>(
66        &self,
67        _rng_type: R,
68        sampler: impl FnOnce(&mut R::RngType) -> T,
69    ) -> T {
70        let mut rng = get_rng::<R>(self);
71        sampler(&mut rng)
72    }
73
74    /// Gets a random sample from the specified distribution using a random number generator
75    /// associated with the given [`RngId`]. If the Rng has not been used before, one will be
76    /// created with the base seed you defined in `set_base_random_seed`.
77    /// Note that this will panic if `set_base_random_seed` was not called yet.
78    #[must_use]
79    fn sample_distr<R: RngId + 'static, T>(
80        &self,
81        _rng_type: R,
82        distribution: impl Distribution<T>,
83    ) -> T
84    where
85        R::RngType: Rng,
86    {
87        let mut rng = get_rng::<R>(self);
88        distribution.sample::<R::RngType>(&mut rng)
89    }
90
91    /// Gets a random sample within the range provided by `range`
92    /// using the generator associated with the given [`RngId`].
93    /// Note that this will panic if `set_base_random_seed` was not called yet.
94    #[must_use]
95    fn sample_range<R: RngId + 'static, S, T>(&self, rng_id: R, range: S) -> T
96    where
97        R::RngType: Rng,
98        S: SampleRange<T>,
99        T: SampleUniform,
100    {
101        self.sample(rng_id, |rng| rng.random_range(range))
102    }
103
104    /// Gets a random boolean value which is true with probability `p`
105    /// using the generator associated with the given [`RngId`].
106    /// Note that this will panic if `set_base_random_seed` was not called yet.
107    #[must_use]
108    fn sample_bool<R: RngId + 'static>(&self, rng_id: R, p: f64) -> bool
109    where
110        R::RngType: Rng,
111    {
112        self.sample(rng_id, |rng| rng.random_bool(p))
113    }
114
115    /// Draws a random entry out of the list provided in `weights`
116    /// with the given weights using the generator associated with the
117    /// given [`RngId`].  Note that this will panic if
118    /// `set_base_random_seed` was not called yet.
119    #[must_use]
120    fn sample_weighted<R: RngId + 'static, T>(&self, _rng_id: R, weights: &[T]) -> usize
121    where
122        R::RngType: Rng,
123        T: Clone
124            + Default
125            + SampleUniform
126            + for<'a> std::ops::AddAssign<&'a T>
127            + PartialOrd
128            + Weight,
129    {
130        let index = WeightedIndex::new(weights).unwrap();
131        let mut rng = get_rng::<R>(self);
132        index.sample(&mut *rng)
133    }
134}
135
136impl ContextRandomExt for Context {}
137
138#[cfg(test)]
139mod test {
140    use crate::context::Context;
141    use crate::rand::distr::weighted::WeightedIndex;
142    use crate::rand::distr::Distribution;
143    use crate::rand::RngCore;
144    use crate::random::context_ext::ContextRandomExt;
145    use crate::{define_data_plugin, define_rng};
146
147    define_rng!(FooRng);
148    define_rng!(BarRng);
149
150    #[test]
151    fn get_rng_basic() {
152        let mut context = Context::new();
153        context.init_random(42);
154
155        assert_ne!(
156            context.sample(FooRng, RngCore::next_u64),
157            context.sample(FooRng, RngCore::next_u64)
158        );
159    }
160
161    #[test]
162    fn multiple_rng_types() {
163        let mut context = Context::new();
164        context.init_random(42);
165
166        assert_ne!(
167            context.sample(FooRng, RngCore::next_u64),
168            context.sample(BarRng, RngCore::next_u64)
169        );
170    }
171
172    #[test]
173    fn reset_seed() {
174        let mut context = Context::new();
175        context.init_random(42);
176
177        let run_0 = context.sample(FooRng, RngCore::next_u64);
178        let run_1 = context.sample(FooRng, RngCore::next_u64);
179
180        // Reset with same seed, ensure we get the same values
181        context.init_random(42);
182        assert_eq!(run_0, context.sample(FooRng, RngCore::next_u64));
183        assert_eq!(run_1, context.sample(FooRng, RngCore::next_u64));
184
185        // Reset with different seed, ensure we get different values
186        context.init_random(88);
187        assert_ne!(run_0, context.sample(FooRng, RngCore::next_u64));
188        assert_ne!(run_1, context.sample(FooRng, RngCore::next_u64));
189    }
190
191    define_data_plugin!(
192        SamplerData,
193        WeightedIndex<f64>,
194        WeightedIndex::new(vec![1.0]).unwrap()
195    );
196
197    #[test]
198    fn sampler_function_closure_capture() {
199        let mut context = Context::new();
200        context.init_random(42);
201
202        // Initialize weighted sampler. Zero is selected with probability 1/3, one with a
203        // probability of 2/3.
204        *context.get_data_mut(SamplerData) = WeightedIndex::new(vec![1.0, 2.0]).unwrap();
205
206        let parameters = context.get_data(SamplerData);
207        let n_samples = 3000;
208        let mut zero_counter = 0;
209        for _ in 0..n_samples {
210            let sample = context.sample(FooRng, |rng| parameters.sample(rng));
211            if sample == 0 {
212                zero_counter += 1;
213            }
214        }
215        // The expected value of `zero_counter` is 1000.
216        assert!((zero_counter - 1000_i32).abs() < 100);
217    }
218
219    #[test]
220    fn sample_distribution() {
221        let mut context = Context::new();
222        context.init_random(42);
223
224        // Initialize weighted sampler. Zero is selected with probability 1/3, one with a
225        // probability of 2/3.
226        *context.get_data_mut(SamplerData) = WeightedIndex::new(vec![1.0, 2.0]).unwrap();
227
228        let parameters = context.get_data(SamplerData);
229        let n_samples = 3000;
230        let mut zero_counter = 0;
231        for _ in 0..n_samples {
232            let sample = context.sample_distr(FooRng, parameters);
233            if sample == 0 {
234                zero_counter += 1;
235            }
236        }
237        // The expected value of `zero_counter` is 1000.
238        assert!((zero_counter - 1000_i32).abs() < 100);
239    }
240
241    #[test]
242    fn sample_range() {
243        let mut context = Context::new();
244        context.init_random(42);
245        let result = context.sample_range(FooRng, 0..10);
246        assert!((0..10).contains(&result));
247    }
248
249    #[test]
250    fn sample_bool() {
251        let mut context = Context::new();
252        context.init_random(42);
253        let _r: bool = context.sample_bool(FooRng, 0.5);
254    }
255
256    #[test]
257    fn sample_weighted() {
258        let mut context = Context::new();
259        context.init_random(42);
260        let r: usize = context.sample_weighted(FooRng, &[0.1, 0.3, 0.4]);
261        assert!(r < 3);
262    }
263}
ixa/random/context_ext.rs

ixa/random/
context_ext.rs
