des_sim/modeling/sampler/instance/
choice.rs

1//! The `choice` module provides the `ChoiceSampler`, a `DurationSampler`
2//! that selects one of several sub-samplers based on a weighted distribution.
3//!
4//! This allows for modeling scenarios where the next event's duration
5//! can come from different underlying processes, each with a specific
6//! probability or weight. It supports both custom weighted distributions
7//! and uniform selection.
8
9use crate::modeling::sampler::{DurationSampler, PendingDuration};
10use crate::primitive::time::SimTime;
11use rand::{Rng, RngExt};
12
13/// A sampler that selects one of several sub-samplers based on a weighted distribution.
14pub struct ChoiceSampler {
15    cdf: Vec<u64>, // Cumulative distribution function (boundary values)
16    values: Vec<Box<dyn DurationSampler>>, // Corresponding samplers
17    total_weight: u64, // Sum of all weights
18}
19
20impl DurationSampler for ChoiceSampler {
21    fn sample(&mut self, rng: &mut dyn Rng, current_tick: SimTime) -> PendingDuration {
22        // Generate a random number in the range [0, total_weight)
23        let r = rng.random_range(0..self.total_weight);
24
25        // Find the index of the interval to which r belongs using binary search.
26        // `partition_point` returns the index of the first element that satisfies `!predicate`.
27        // Since we want the last index where `x <= r`, we search for `x <= r` and subtract 1.
28        let idx = self.cdf.partition_point(|&x| x <= r) - 1;
29
30        self.values[idx].sample(rng, current_tick)
31    }
32}
33
34impl ChoiceSampler {
35    /// Constructs a `ChoiceSampler` from a collection of `(DurationSampler, Weight)` pairs.
36    pub fn new(histogram: impl IntoIterator<Item = (Box<dyn DurationSampler>, u64)>) -> Self {
37        // Converted in advance so that it can be sampled using Inverse Transform Sampling
38        // from the empirical distribution
39        let mut cdf = Vec::new();
40        let mut values = Vec::new();
41        let mut current_sum = 0;
42
43        // The first boundary is 0
44        cdf.push(0);
45        for (duration, weight) in histogram.into_iter() {
46            current_sum += weight;
47            cdf.push(current_sum);
48            values.push(duration);
49        }
50
51        assert!(
52            !values.is_empty(),
53            "ChoiceSampler must have at least one sampler with positive weight"
54        );
55        assert!(current_sum > 0, "Total weight must be greater than 0");
56
57        Self {
58            cdf,
59            values,
60            total_weight: current_sum,
61        }
62    }
63
64    /// Constructs a `ChoiceSampler` with uniform weight (1) for all provided samplers.
65    pub fn new_as_uniform(histogram: impl IntoIterator<Item = Box<dyn DurationSampler>>) -> Self {
66        ChoiceSampler::new(histogram.into_iter().map(|s| (s, 1)))
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use crate::modeling::sampler::CombinatorExt;
74    use crate::modeling::sampler::instance::ConstantSampler;
75    use rand::SeedableRng;
76    use rand::rngs::SmallRng;
77    use std::collections::HashMap;
78
79    #[test]
80    fn test_choice_sampler_weighted() {
81        let mut rng = SmallRng::seed_from_u64(2);
82        let s1 = ConstantSampler::new(10.0);
83        let s2 = ConstantSampler::new(20.0);
84        let s3 = ConstantSampler::new(30.0);
85
86        // Weights: 1, 2, 1. Total = 4.
87        let mut sampler =
88            ChoiceSampler::new(vec![(s1.boxed(), 1), (s2.boxed(), 2), (s3.boxed(), 1)]);
89
90        let mut results: HashMap<String, usize> = HashMap::new();
91        for _ in 0..1000 {
92            let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
93            let entry = results
94                .entry(format!("{:<.2}", sample.raw_value()))
95                .or_insert(0);
96            *entry += 1;
97        }
98
99        // Check if the distribution is roughly as expected
100        // With 1000 samples, and weights 1:2:1, we expect roughly 250:500:250
101        let count_10 = *results.get(&format!("{:<.2}", 10.0)).unwrap_or(&0);
102        let count_20 = *results.get(&format!("{:<.2}", 20.0)).unwrap_or(&0);
103        let count_30 = *results.get(&format!("{:<.2}", 30.0)).unwrap_or(&0);
104
105        // Allow for some deviation due to randomness
106        assert!(count_10 > 200 && count_10 < 300);
107        assert!(count_20 > 450 && count_20 < 550);
108        assert!(count_30 > 200 && count_30 < 300);
109    }
110
111    #[test]
112    fn test_choice_sampler_uniform() {
113        let mut rng = SmallRng::seed_from_u64(2);
114        let s1 = ConstantSampler::new(10.0);
115        let s2 = ConstantSampler::new(20.0);
116        let s3 = ConstantSampler::new(30.0);
117
118        let mut sampler = ChoiceSampler::new_as_uniform(vec![s1.boxed(), s2.boxed(), s3.boxed()]);
119
120        let mut results: HashMap<String, usize> = HashMap::new();
121        for _ in 0..1000 {
122            let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
123            let entry = results
124                .entry(format!("{:<.2}", sample.raw_value()))
125                .or_insert(0);
126            *entry += 1_usize;
127        }
128
129        // With 1000 samples and 3 choices, we expect roughly 333 for each
130        let count_10 = *results.get(&format!("{:<.2}", 10.0)).unwrap_or(&0);
131        let count_20 = *results.get(&format!("{:<.2}", 20.0)).unwrap_or(&0);
132        let count_30 = *results.get(&format!("{:<.2}", 30.0)).unwrap_or(&0);
133
134        assert!(count_10 > 280 && count_10 < 380);
135        assert!(count_20 > 280 && count_20 < 380);
136        assert!(count_30 > 280 && count_30 < 380);
137    }
138
139    #[test]
140    #[should_panic(expected = "ChoiceSampler must have at least one sampler")]
141    fn test_choice_sampler_empty_histogram() {
142        let _sampler = ChoiceSampler::new(vec![]);
143    }
144
145    #[test]
146    #[should_panic(expected = "Total weight must be greater than 0")]
147    fn test_choice_sampler_zero_total_weight() {
148        let s1 = ConstantSampler::new(10.0);
149        let _sampler = ChoiceSampler::new(vec![(s1.boxed(), 0)]);
150    }
151}
des_sim/modeling/sampler/instance/choice.rs

des_sim/modeling/sampler/instance/
choice.rs
