Skip to main content

proof_engine/ecology/
disease.rs

1//! Disease propagation — SIR/SIS models with spatial spread.
2
3/// SIR model state.
4#[derive(Debug, Clone)]
5pub struct SirState {
6    pub susceptible: f64,
7    pub infected: f64,
8    pub recovered: f64,
9    pub beta: f64,    // transmission rate
10    pub gamma: f64,   // recovery rate
11}
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    /// Basic reproduction number.
26    pub fn r0(&self) -> f64 { self.beta / self.gamma }
27
28    /// Step the SIR model.
29    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    /// Is the epidemic over?
45    pub fn is_over(&self) -> bool { self.infected < 0.5 }
46}
47
48/// SIS model (no immunity — recovered become susceptible again).
49#[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    /// Endemic equilibrium infected fraction.
74    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}