1use nalgebra::{DVector, DMatrix};
4use serde::{Serialize, Deserialize};
5use rand::Rng;
6
7use crate::brownian::BrownianConfig;
8use crate::fokker_planck::DriftDiffusionDecomposition;
9use crate::langevin::{HMCConfig, HMCResult, PotentialEnergy};
10use crate::fractional::LevyFlight;
11use crate::spectral_diffusion::{SpectralDecomposition, graph_laplacian};
12use crate::transport_diffusion::DiscreteDistribution;
13use crate::anisotropic::DiffusionTensor;
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct AgentDiffusionConfig {
17 pub n_agents: usize,
18 pub dimension: usize,
19 pub dt: f64,
20 pub total_time: f64,
21 pub diffusion_coeff: f64,
22 pub drift_coeff: f64,
23}
24
25impl Default for AgentDiffusionConfig {
26 fn default() -> Self {
27 Self { n_agents: 100, dimension: 2, dt: 0.01, total_time: 1.0, diffusion_coeff: 1.0, drift_coeff: 0.0 }
28 }
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct AgentState {
33 pub positions: Vec<Vec<f64>>,
34 pub time: f64,
35}
36
37pub struct AgentDiffusion {
38 pub config: AgentDiffusionConfig,
39 pub state: AgentState,
40}
41
42impl AgentDiffusion {
43 pub fn new(config: AgentDiffusionConfig) -> Self {
44 let positions = vec![vec![0.0; config.dimension]; config.n_agents];
45 Self { config, state: AgentState { positions, time: 0.0 } }
46 }
47
48 pub fn new_random(config: AgentDiffusionConfig, spread: f64) -> Self {
49 let mut rng = rand::thread_rng();
50 let positions = (0..config.n_agents)
51 .map(|_| (0..config.dimension).map(|_| rng.gen_range(-spread..spread)).collect())
52 .collect();
53 Self { config, state: AgentState { positions, time: 0.0 } }
54 }
55
56 pub fn brownian_step(&mut self) {
57 let mut rng = rand::thread_rng();
58 let sqrt_dt = self.config.dt.sqrt();
59 for pos in &mut self.state.positions {
60 for j in 0..self.config.dimension {
61 let z: f64 = rng.gen_range(-1.0..1.0);
62 pos[j] += self.config.drift_coeff * self.config.dt
63 + self.config.diffusion_coeff * sqrt_dt * z * std::f64::consts::SQRT_2;
64 }
65 }
66 self.state.time += self.config.dt;
67 }
68
69 pub fn run_brownian(&mut self) {
70 let n_steps = (self.config.total_time / self.config.dt) as usize;
71 for _ in 0..n_steps { self.brownian_step(); }
72 }
73
74 pub fn heat_kernel(&self, i: usize, j: usize, t: f64) -> f64 {
75 let dim = self.config.dimension;
76 let dist_sq: f64 = self.state.positions[i].iter()
77 .zip(self.state.positions[j].iter())
78 .map(|(a, b)| (a - b).powi(2)).sum();
79 let norm = (4.0 * std::f64::consts::PI * t).powf(-(dim as f64) / 2.0);
80 norm * (-dist_sq / (4.0 * t)).exp()
81 }
82
83 pub fn estimate_drift_diffusion(&self, trajectories: &[Vec<Vec<f64>>]) -> Vec<DriftDiffusionDecomposition> {
84 trajectories.iter().map(|traj| {
85 let obs: Vec<f64> = traj.iter().map(|p| p[0]).collect();
86 DriftDiffusionDecomposition::estimate(&obs, self.config.dt)
87 }).collect()
88 }
89
90 pub fn hmc_sample(&self, hmc_config: &HMCConfig, potential: &dyn PotentialEnergy) -> HMCResult {
91 let x0 = DVector::zeros(self.config.dimension);
92 crate::langevin::hmc_sample(hmc_config, &x0, potential)
93 }
94
95 pub fn levy_flight(&self, _agent_idx: usize, alpha: f64, n_steps: usize) -> LevyFlight {
96 LevyFlight::simulate(self.config.dimension, alpha, n_steps, self.config.diffusion_coeff)
97 }
98
99 pub fn spectral_decomposition(&self, adjacency: &DMatrix<f64>) -> Result<SpectralDecomposition, String> {
100 let lap = graph_laplacian(adjacency);
101 SpectralDecomposition::decompose(&lap)
102 }
103
104 pub fn position_distribution(&self, dim: usize, n_bins: usize, range: (f64, f64)) -> DiscreteDistribution {
105 let bin_width = (range.1 - range.0) / n_bins as f64;
106 let mut counts = vec![0.0; n_bins];
107 for pos in &self.state.positions {
108 if dim < pos.len() {
109 let idx = ((pos[dim] - range.0) / bin_width).floor() as usize;
110 if idx < n_bins { counts[idx] += 1.0; }
111 }
112 }
113 let support: Vec<f64> = (0..n_bins).map(|i| range.0 + (i as f64 + 0.5) * bin_width).collect();
114 DiscreteDistribution::normalized(support, counts)
115 }
116
117 pub fn covariance_tensor(&self) -> DiffusionTensor {
118 let n = self.state.positions.len() as f64;
119 let dim = self.config.dimension.min(3);
120 let mut mean = nalgebra::Vector3::zeros();
121 for pos in &self.state.positions {
122 for d in 0..dim { mean[d] += pos[d]; }
123 }
124 mean /= n;
125 let mut cov = nalgebra::Matrix3::zeros();
126 for pos in &self.state.positions {
127 for i in 0..dim {
128 for j in 0..dim {
129 cov[(i, j)] += (pos[i] - mean[i]) * (pos[j] - mean[j]);
130 }
131 }
132 }
133 cov /= (n - 1.0).max(1.0);
134 DiffusionTensor { tensor: cov }
135 }
136
137 pub fn mean_position(&self) -> Vec<f64> {
138 let n = self.state.positions.len();
139 if n == 0 { return vec![0.0; self.config.dimension]; }
140 let mut mean = vec![0.0; self.config.dimension];
141 for pos in &self.state.positions {
142 for (i, &v) in pos.iter().enumerate() { mean[i] += v; }
143 }
144 for m in &mut mean { *m /= n as f64; }
145 mean
146 }
147
148 pub fn mean_squared_displacement(&self) -> f64 {
149 self.state.positions.iter()
150 .map(|p| p.iter().map(|x| x * x).sum::<f64>())
151 .sum::<f64>() / self.state.positions.len() as f64
152 }
153
154 pub fn reset(&mut self) {
155 for pos in &mut self.state.positions { pos.fill(0.0); }
156 self.state.time = 0.0;
157 }
158}
159
160#[cfg(test)]
161mod tests {
162 use super::*;
163 use approx::assert_relative_eq;
164
165 #[test]
166 fn test_agent_diffusion_new() {
167 let config = AgentDiffusionConfig::default();
168 let ad = AgentDiffusion::new(config);
169 assert_eq!(ad.state.positions.len(), 100);
170 assert_relative_eq!(ad.state.time, 0.0);
171 }
172
173 #[test]
174 fn test_agent_diffusion_default_config() {
175 let config = AgentDiffusionConfig::default();
176 assert_eq!(config.n_agents, 100);
177 assert_eq!(config.dimension, 2);
178 }
179
180 #[test]
181 fn test_brownian_step() {
182 let config = AgentDiffusionConfig { n_agents: 10, ..Default::default() };
183 let mut ad = AgentDiffusion::new(config);
184 ad.brownian_step();
185 assert_relative_eq!(ad.state.time, 0.01);
186 assert!(ad.state.positions.iter().any(|p| p.iter().any(|&x| x != 0.0)));
187 }
188
189 #[test]
190 fn test_run_brownian() {
191 let config = AgentDiffusionConfig { n_agents: 10, total_time: 0.1, dt: 0.01, ..Default::default() };
192 let mut ad = AgentDiffusion::new(config);
193 ad.run_brownian();
194 assert!(ad.state.time > 0.09);
195 assert!(ad.mean_squared_displacement() > 0.0);
196 }
197
198 #[test]
199 fn test_heat_kernel_agents() {
200 let config = AgentDiffusionConfig { n_agents: 2, dimension: 1, ..Default::default() };
201 let mut ad = AgentDiffusion::new(config);
202 ad.state.positions[0] = vec![0.0];
203 ad.state.positions[1] = vec![1.0];
204 let k = ad.heat_kernel(0, 1, 1.0);
205 assert!(k > 0.0);
206 }
207
208 #[test]
209 fn test_position_distribution() {
210 let config = AgentDiffusionConfig { n_agents: 100, dimension: 1, ..Default::default() };
211 let ad = AgentDiffusion::new_random(config, 1.0);
212 let dist = ad.position_distribution(0, 10, (-5.0, 5.0));
213 let sum: f64 = dist.probabilities.iter().sum();
214 assert_relative_eq!(sum, 1.0, epsilon = 0.01);
215 }
216
217 #[test]
218 fn test_mean_position() {
219 let config = AgentDiffusionConfig { n_agents: 3, dimension: 2, ..Default::default() };
220 let mut ad = AgentDiffusion::new(config);
221 ad.state.positions[0] = vec![1.0, 0.0];
222 ad.state.positions[1] = vec![0.0, 1.0];
223 ad.state.positions[2] = vec![-1.0, -1.0];
224 let mean = ad.mean_position();
225 assert_relative_eq!(mean[0], 0.0, epsilon = 1e-10);
226 }
227
228 #[test]
229 fn test_reset() {
230 let config = AgentDiffusionConfig { n_agents: 5, ..Default::default() };
231 let mut ad = AgentDiffusion::new(config);
232 ad.run_brownian();
233 ad.reset();
234 for pos in &ad.state.positions { assert!(pos.iter().all(|&x| x == 0.0)); }
235 }
236
237 #[test]
238 fn test_covariance_tensor() {
239 let config = AgentDiffusionConfig { n_agents: 50, dimension: 2, ..Default::default() };
240 let ad = AgentDiffusion::new_random(config, 2.0);
241 let tensor = ad.covariance_tensor();
242 assert!(tensor.mean_diffusivity() > 0.0);
243 }
244
245 #[test]
246 fn test_spectral_decomposition() {
247 let config = AgentDiffusionConfig { n_agents: 5, dimension: 1, ..Default::default() };
248 let ad = AgentDiffusion::new(config);
249 let adj = DMatrix::from_element(5, 5, 0.2);
250 let decomp = ad.spectral_decomposition(&adj).unwrap();
251 assert_eq!(decomp.eigenvalues.len(), 5);
252 }
253
254 #[test]
255 fn test_levy_flight() {
256 let config = AgentDiffusionConfig { n_agents: 5, dimension: 2, ..Default::default() };
257 let ad = AgentDiffusion::new(config);
258 let flight = ad.levy_flight(0, 1.5, 50);
259 assert_eq!(flight.positions.len(), 51);
260 }
261
262 #[test]
263 fn test_mean_squared_displacement() {
264 let config = AgentDiffusionConfig { n_agents: 100, dimension: 2, total_time: 0.5, ..Default::default() };
265 let mut ad = AgentDiffusion::new(config);
266 ad.run_brownian();
267 assert!(ad.mean_squared_displacement() > 0.0);
268 }
269}