Initial commit

examples/epidemics/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "epidemics"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+nalgebra = "0.32.3"
+rand = "0.8.5"
+rustmodel = { path = "../..", features = ["plot"] }

examples/epidemics/src/main.rs

@@ -0,0 +1,107 @@
+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],
+			}
+		}
+	}
+}