Skip to main content

des_sim/modeling/sampler/combinator/
aggregate.rs

1//! The `aggregate` module provides the `AggregateSampler` and `AggregateBuilder`.
2//!
3//! `AggregateSampler` combines the outputs of multiple `DurationSampler` instances
4//! using a custom aggregation function, allowing for complex statistical modeling.
5//! `AggregateBuilder` offers a fluent API for constructing these aggregated samplers.
6
7use crate::modeling::sampler::{DurationSampler, PendingDuration};
8use crate::primitive::time::SimTime;
9use rand::Rng;
10
11/// A builder for constructing an `AggregateSampler` in a fluent manner.
12pub struct AggregateBuilder {
13    samplers: Vec<Box<dyn DurationSampler>>,
14}
15
16impl AggregateBuilder {
17    /// Initializes the builder with the first sampler.
18    pub fn from_sampler(sampler: impl DurationSampler + 'static) -> Self {
19        AggregateBuilder {
20            samplers: vec![Box::new(sampler)],
21        }
22    }
23
24    /// Adds another sampler to the aggregation list.
25    pub fn add_sampler(mut self, sampler: Box<dyn DurationSampler>) -> Self {
26        self.samplers.push(sampler);
27        self
28    }
29
30    /// Consumes the builder to produce an `AggregateSampler` using the provided function `f`.
31    pub fn build<F>(self, f: F) -> AggregateSampler<F>
32    where
33        F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
34    {
35        AggregateSampler {
36            samplers: self.samplers,
37            f,
38        }
39    }
40}
41
42/// A sampler that aggregates results from multiple sub-samplers using a closure.
43pub struct AggregateSampler<F>
44where
45    F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
46{
47    samplers: Vec<Box<dyn DurationSampler>>,
48    f: F,
49}
50
51impl<F> DurationSampler for AggregateSampler<F>
52where
53    F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
54{
55    fn sample(&mut self, rng: &mut dyn Rng, current_tick: SimTime) -> PendingDuration {
56        let mut sampled_list = Vec::with_capacity(self.samplers.len());
57        for sampler in &mut self.samplers {
58            let sampled = sampler.sample(rng, current_tick);
59            sampled_list.push(sampled.raw_value());
60        }
61
62        PendingDuration::new((self.f)(rng, current_tick, sampled_list))
63    }
64}
65
66impl<F> AggregateSampler<F>
67where
68    F: FnMut(&mut dyn Rng, SimTime, Vec<f64>) -> f64,
69{
70    /// Creates a new `AggregateSampler`. Panics if `samplers` is empty.
71    pub fn new(samplers: impl IntoIterator<Item = Box<dyn DurationSampler>>, f: F) -> Self {
72        let samplers: Vec<_> = samplers.into_iter().collect();
73
74        assert!(
75            !samplers.is_empty(),
76            "AggregateSampler requires at least one sampler."
77        );
78
79        AggregateSampler { samplers, f }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::modeling::sampler::CombinatorExt;
87    use crate::modeling::sampler::instance::ConstantSampler;
88    use rand::SeedableRng;
89    use rand::rngs::SmallRng;
90
91    #[test]
92    fn test_aggregate_sampler_sum() {
93        let mut rng = SmallRng::seed_from_u64(2);
94        let s1 = ConstantSampler::new(10.0);
95        let s2 = ConstantSampler::new(20.0);
96        let s3 = ConstantSampler::new(30.0);
97
98        let mut sampler =
99            AggregateSampler::new(vec![s1.boxed(), s2.boxed(), s3.boxed()], |_, _, samples| {
100                samples.iter().sum()
101            });
102
103        let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
104        assert_eq!(sample.raw_value(), 60.0);
105    }
106
107    #[test]
108    fn test_aggregate_sampler_builder() {
109        let mut rng = SmallRng::seed_from_u64(2);
110        let s1 = ConstantSampler::new(5.0);
111        let s2 = ConstantSampler::new(15.0);
112
113        let mut sampler = AggregateBuilder::from_sampler(s1)
114            .add_sampler(s2.boxed())
115            .build(|_, _, samples| samples[0] * samples[1]);
116
117        let sample = sampler.sample(&mut rng, SimTime::from_ticks(0));
118        assert_eq!(sample.raw_value(), 75.0); // 5.0 * 15.0
119    }
120
121    #[test]
122    #[should_panic]
123    fn test_aggregate_sampler_empty_samplers() {
124        let _sampler = AggregateSampler::new(vec![], |_, _, _| 0.0);
125    }
126}