proof_engine/ecology/
disease.rs1#[derive(Debug, Clone)]
5pub struct SirState {
6 pub susceptible: f64,
7 pub infected: f64,
8 pub recovered: f64,
9 pub beta: f64, pub gamma: f64, }
12
13impl SirState {
14 pub fn new(pop: f64, initial_infected: f64, beta: f64, gamma: f64) -> Self {
15 Self {
16 susceptible: pop - initial_infected,
17 infected: initial_infected,
18 recovered: 0.0,
19 beta, gamma,
20 }
21 }
22
23 pub fn total(&self) -> f64 { self.susceptible + self.infected + self.recovered }
24
25 pub fn r0(&self) -> f64 { self.beta / self.gamma }
27
28 pub fn step(&mut self, dt: f64) {
30 let n = self.total();
31 if n < 1.0 { return; }
32 let new_infections = self.beta * self.susceptible * self.infected / n * dt;
33 let new_recoveries = self.gamma * self.infected * dt;
34
35 self.susceptible -= new_infections;
36 self.infected += new_infections - new_recoveries;
37 self.recovered += new_recoveries;
38
39 self.susceptible = self.susceptible.max(0.0);
40 self.infected = self.infected.max(0.0);
41 self.recovered = self.recovered.max(0.0);
42 }
43
44 pub fn is_over(&self) -> bool { self.infected < 0.5 }
46}
47
48#[derive(Debug, Clone)]
50pub struct SisState {
51 pub susceptible: f64,
52 pub infected: f64,
53 pub beta: f64,
54 pub gamma: f64,
55}
56
57impl SisState {
58 pub fn new(pop: f64, initial_infected: f64, beta: f64, gamma: f64) -> Self {
59 Self { susceptible: pop - initial_infected, infected: initial_infected, beta, gamma }
60 }
61
62 pub fn step(&mut self, dt: f64) {
63 let n = self.susceptible + self.infected;
64 if n < 1.0 { return; }
65 let new_infections = self.beta * self.susceptible * self.infected / n * dt;
66 let new_recoveries = self.gamma * self.infected * dt;
67 self.susceptible += new_recoveries - new_infections;
68 self.infected += new_infections - new_recoveries;
69 self.susceptible = self.susceptible.max(0.0);
70 self.infected = self.infected.max(0.0);
71 }
72
73 pub fn endemic_fraction(&self) -> f64 {
75 if self.beta <= self.gamma { 0.0 }
76 else { 1.0 - self.gamma / self.beta }
77 }
78}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83
84 #[test]
85 fn test_sir_epidemic() {
86 let mut sir = SirState::new(1000.0, 10.0, 0.3, 0.1);
87 assert!(sir.r0() > 1.0, "R0 should be > 1 for epidemic");
88 for _ in 0..10000 { sir.step(0.1); }
89 assert!(sir.is_over(), "epidemic should resolve");
90 assert!(sir.recovered > 500.0, "most should have been infected");
91 }
92
93 #[test]
94 fn test_sir_no_epidemic() {
95 let mut sir = SirState::new(1000.0, 10.0, 0.05, 0.1);
96 assert!(sir.r0() < 1.0, "R0 should be < 1");
97 for _ in 0..10000 { sir.step(0.1); }
98 assert!(sir.recovered < 50.0, "disease should die out quickly");
99 }
100
101 #[test]
102 fn test_sis_endemic() {
103 let mut sis = SisState::new(1000.0, 10.0, 0.3, 0.1);
104 for _ in 0..10000 { sis.step(0.1); }
105 assert!(sis.infected > 10.0, "SIS should reach endemic equilibrium");
106 }
107}