108 lines
2.1 KiB
Rust
108 lines
2.1 KiB
Rust
use nalgebra::vector;
|
|
use rustmodel::prelude::*;
|
|
|
|
fn main() {
|
|
// Start with 1% infected
|
|
let x0 = vector![0.99, 0.01];
|
|
let dt = 0.1;
|
|
let nsamples: usize = 400;
|
|
|
|
let settings = sir::SirSettings {
|
|
beta: 0.6,
|
|
gamma: 0.1,
|
|
pop: 1.0,
|
|
};
|
|
|
|
let model = sir::Sir {
|
|
s: settings.clone(),
|
|
};
|
|
let solver = ImplicitEuler {
|
|
dt,
|
|
tol: 0.000001,
|
|
niters: 100,
|
|
};
|
|
let mut xlist = Vec::with_capacity(nsamples);
|
|
xlist.push(x0);
|
|
let mut x = x0;
|
|
for _ in 1..nsamples {
|
|
x = solver.f(&model, x);
|
|
xlist.push(x);
|
|
}
|
|
|
|
xlist.iter().for_each(|x| println!("{}\t{}", x[0], x[1]));
|
|
|
|
PlotBuilder::default()
|
|
.dt(dt)
|
|
.title(Some(String::from("Epidemics Compartmental Model")))
|
|
.y_min(Some(0.0))
|
|
.y_max(Some(1.0))
|
|
.x_label(Some(String::from("Time")))
|
|
.build()
|
|
.unwrap()
|
|
.plot(
|
|
"target/sir.png",
|
|
&xlist,
|
|
&[("Susceptible", colors::BLUE), ("Infected", colors::RED)],
|
|
);
|
|
}
|
|
|
|
/// SIR model without vital dynamics
|
|
pub mod sir {
|
|
use nalgebra::{base::*, matrix, vector};
|
|
use rustmodel::prelude::*;
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct SirSettings<T> {
|
|
/// Transmission probability
|
|
pub beta: T,
|
|
/// Removal probability
|
|
pub gamma: T,
|
|
/// Total population
|
|
pub pop: T,
|
|
}
|
|
|
|
impl<T> Settings for SirSettings<T> {}
|
|
|
|
#[derive(Clone)]
|
|
pub struct Sir<T> {
|
|
pub s: SirSettings<T>,
|
|
}
|
|
|
|
impl Model<f64, SirSettings<f64>, 2> for Sir<f64> {
|
|
fn f(&self, x: Vector2<f64>) -> Vector2<f64> {
|
|
vector![
|
|
-self.s.beta * x[0] * x[1] / self.s.pop,
|
|
self.s.beta * x[0] * x[1] / self.s.pop - self.s.gamma * x[1]
|
|
]
|
|
}
|
|
fn df(&self, x: Vector2<f64>) -> Matrix2<f64> {
|
|
matrix![
|
|
-self.s.beta*x[1]/self.s.pop, -self.s.beta*x[0]/self.s.pop;
|
|
self.s.beta*x[1]/self.s.pop, self.s.beta*x[0]/self.s.pop-self.s.gamma
|
|
]
|
|
}
|
|
fn get_settings(&self) -> &SirSettings<f64> {
|
|
&self.s
|
|
}
|
|
fn get_settings_mut(&mut self) -> &mut SirSettings<f64> {
|
|
&mut self.s
|
|
}
|
|
}
|
|
|
|
impl<T> From<SirSettings<T>> for Vect<T, 3> {
|
|
fn from(s: SirSettings<T>) -> Self {
|
|
vector![s.beta, s.gamma, s.pop]
|
|
}
|
|
}
|
|
|
|
impl<T: Scalar + Copy> From<Vect<T, 3>> for SirSettings<T> {
|
|
fn from(v: Vect<T, 3>) -> Self {
|
|
Self {
|
|
beta: v[0],
|
|
gamma: v[1],
|
|
pop: v[2],
|
|
}
|
|
}
|
|
}
|
|
}
|