1use dsfb::TrustStats;
2use serde::{Deserialize, Serialize};
3
4use crate::disturbances::{build_disturbance, DisturbanceKind};
5use crate::envelope::{ResidualEnvelope, TrustWeight};
6
7#[derive(Clone, Debug, Serialize, Deserialize)]
8pub struct SimulationConfig {
9 pub n_steps: usize,
10 pub rho: f64,
11 pub beta: f64,
12 pub disturbance_kind: DisturbanceKind,
13 pub epsilon_bound: f64,
14}
15
16#[derive(Clone, Debug, Default, Serialize, Deserialize)]
17pub struct SimulationResult {
18 pub s: Vec<f64>,
19 pub w: Vec<f64>,
20 pub r: Vec<f64>,
21 pub d: Vec<f64>,
22}
23
24impl SimulationResult {
25 pub fn len(&self) -> usize {
26 self.s.len()
27 }
28
29 pub fn is_empty(&self) -> bool {
30 self.s.is_empty()
31 }
32
33 pub fn final_trust_stats(&self) -> TrustStats {
34 TrustStats {
35 residual_ema: *self.s.last().unwrap_or(&0.0),
36 weight: *self.w.last().unwrap_or(&1.0),
37 }
38 }
39}
40
41pub fn run_simulation(config: &SimulationConfig) -> SimulationResult {
42 run_simulation_with_s0(config, 0.0)
43}
44
45pub fn run_simulation_with_s0(config: &SimulationConfig, s0: f64) -> SimulationResult {
46 simulate_channel(config, s0, 0, &config.disturbance_kind)
47}
48
49pub fn run_multichannel_simulation(
50 config: &SimulationConfig,
51 n_channels: usize,
52 group_assignments: Option<&[usize]>,
53 correlated_groups: bool,
54) -> Vec<SimulationResult> {
55 assert!(n_channels > 0, "n_channels must be > 0");
56
57 if let Some(groups) = group_assignments {
58 assert_eq!(
59 groups.len(),
60 n_channels,
61 "group_assignments length must match n_channels",
62 );
63 }
64
65 let default_groups: Vec<usize> = (0..n_channels).collect();
66 let groups = group_assignments.unwrap_or(&default_groups);
67
68 (0..n_channels)
69 .map(|channel_idx| {
70 let key = if correlated_groups {
71 groups[channel_idx]
72 } else {
73 channel_idx
74 };
75 let kind = config.disturbance_kind.channelized(key);
76 let s0 = 0.02 * key as f64;
77 simulate_channel(config, s0, key, &kind)
78 })
79 .collect()
80}
81
82fn simulate_channel(
83 config: &SimulationConfig,
84 s0: f64,
85 channel_key: usize,
86 disturbance_kind: &DisturbanceKind,
87) -> SimulationResult {
88 assert!(config.n_steps > 0, "n_steps must be > 0");
89 assert!(
90 config.rho > 0.0 && config.rho < 1.0,
91 "rho must be in (0, 1)"
92 );
93 assert!(config.beta > 0.0, "beta must be > 0");
94 assert!(
95 config.epsilon_bound.is_finite() && config.epsilon_bound >= 0.0,
96 "epsilon_bound must be finite and >= 0",
97 );
98
99 let mut envelope = ResidualEnvelope::new(config.rho, s0);
100 let mut disturbance = build_disturbance(disturbance_kind);
101 disturbance.reset();
102
103 let mut result = SimulationResult {
104 s: Vec::with_capacity(config.n_steps),
105 w: Vec::with_capacity(config.n_steps),
106 r: Vec::with_capacity(config.n_steps),
107 d: Vec::with_capacity(config.n_steps),
108 };
109
110 for n in 0..config.n_steps {
111 let d = disturbance.next(n);
112 let epsilon = epsilon_at(n, config.epsilon_bound, channel_key);
113 let r = epsilon + d;
114 let s = envelope.update(r);
115 let w = TrustWeight::weight(config.beta, s);
116
117 result.d.push(d);
118 result.r.push(r);
119 result.s.push(s);
120 result.w.push(w);
121 }
122
123 result
124}
125
126fn epsilon_at(n: usize, epsilon_bound: f64, channel_key: usize) -> f64 {
127 if epsilon_bound == 0.0 {
128 return 0.0;
129 }
130
131 let phase = channel_key as f64 * 0.37;
132 let a = (0.17 * n as f64 + phase).sin();
133 let b = (0.043 * n as f64 + 0.5 * phase).cos();
134 epsilon_bound * (0.6 * a + 0.4 * b)
135}
136
137#[cfg(test)]
138mod tests {
139 use super::{run_multichannel_simulation, run_simulation, SimulationConfig};
140 use crate::disturbances::DisturbanceKind;
141
142 #[test]
143 fn pointwise_simulation_reaches_plateau() {
144 let config = SimulationConfig {
145 n_steps: 64,
146 rho: 0.95,
147 beta: 2.0,
148 disturbance_kind: DisturbanceKind::PointwiseBounded { d: 0.4 },
149 epsilon_bound: 0.0,
150 };
151
152 let result = run_simulation(&config);
153 let final_s = *result.s.last().expect("result should be non-empty");
154 assert!(final_s > 0.35 && final_s < 0.41);
155 }
156
157 #[test]
158 fn multichannel_group_correlation_reuses_disturbance() {
159 let config = SimulationConfig {
160 n_steps: 12,
161 rho: 0.9,
162 beta: 3.0,
163 disturbance_kind: DisturbanceKind::PersistentElevated {
164 r_nom: 0.1,
165 r_high: 0.5,
166 step_time: 4,
167 },
168 epsilon_bound: 0.0,
169 };
170
171 let results = run_multichannel_simulation(&config, 3, Some(&[0, 0, 1]), true);
172 assert_eq!(results[0].d, results[1].d);
173 assert_ne!(results[0].d, results[2].d);
174 }
175}