diff --git a/Cargo.toml b/Cargo.toml index 5f304bc..4ca52a6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,3 +6,9 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +serde = { version = "1.0", features = ["derive"] } +nalgebra = { version = "0.32", features = ["serde-serialize"] } +num-traits = "0.2.15" + +[dev-dependencies] +approx = "0.5" diff --git a/src/controller.rs b/src/controller.rs new file mode 100644 index 0000000..7877458 --- /dev/null +++ b/src/controller.rs @@ -0,0 +1,90 @@ +use super::ode::SystemTrait; +use num_traits::Float; + +pub trait Controller +where + f64: From, + F: SystemTrait, + T: Copy + From, +{ + fn determine_step(&mut self, h: T, err: T) -> (bool, T); +} + +#[derive(Debug, Clone, Copy)] +pub struct PIController where f64: From { + pub alpha: T, + pub beta: T, + pub factor_c1: T, + pub factor_c2: T, + pub factor_old: T, + pub h_max: T, + pub safety_factor: T, + pub old_h: T, +} + +impl Controller for PIController +where + f64: From, + F: SystemTrait, + T: Copy + From + Float, +{ + /// Determines if the previously run step size and error were valid or not. Either way, it also + /// returns what the next step size should be + fn determine_step(&mut self, h: T, err: T) -> (bool, T) { + let factor_11 = err.powf(self.alpha); + let factor = self.factor_c2.max(self.factor_c1.min(factor_11 * self.factor_old.powf(-self.beta) / self.safety_factor)); + let mut h_new = h / factor; + if err <= 1.0.into() { + // Accept the stepsize + self.factor_old = err.max(1.0e-4.into()); + if h_new.abs() > self.h_max { + // If the step is too big + h_new = self.h_max.copysign(h_new); + } + (true, h_new) + } else { + // Reject the stepsize and propose a smaller one + h_new = h / (self.factor_c1.min(factor_11 / self.safety_factor)); + (false, h_new) + } + } +} + +impl PIController +where + f64: From, + T: Copy + From + Float, +{ + pub fn new(alpha:T, beta:T, max_factor: T, min_factor: T, h_max: T, safety_factor: T, initial_h: T) -> Self { + Self { + alpha: alpha, + beta: beta, + factor_c1: >::from(1.0.into()) / min_factor, + factor_c2: >::from(1.0.into()) / max_factor, + factor_old: 1.0e-4.into(), + h_max: h_max.abs(), + safety_factor: safety_factor, + old_h: initial_h, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_controller_creation() { + let controller = PIController::new(0.17, 0.04, 10.0, 0.2, 10.0, 0.9, 1e-4); + + assert!(controller.alpha == 0.17); + assert!(controller.beta == 0.04); + assert!(controller.factor_c1 == 1.0/0.2); + assert!(controller.factor_c2 == 1.0/10.0); + assert!(controller.factor_old == 1.0e-4); + assert!(controller.h_max == 10.0); + assert!(controller.safety_factor == 0.9); + assert!(controller.old_h == 1e-4); + // assert!(controller.determine_step(0.1, 0.1) == (true, 0.01)); + } +} diff --git a/src/integrator.rs b/src/integrator.rs new file mode 100644 index 0000000..0c8f0dc --- /dev/null +++ b/src/integrator.rs @@ -0,0 +1,392 @@ +use num_traits::identities; +use num_traits::Float; +use std::cmp::PartialEq; +use std::fmt::Debug; +use std::ops::{Add, AddAssign, Sub, SubAssign, Mul, MulAssign}; +use nalgebra::{SVector, SMatrix}; +use nalgebra as na; + +use super::ode::{ODE, SystemTrait}; + +#[derive(Debug, Clone, Copy)] +pub struct Tableau { + a: SMatrix, + b: SVector, + b_star: Option>, + c: SVector, +} + + +/// Integrator Trait +/// +/// T = The value type, usually f64, but must at least implement to/from f64 +/// F = A type that implements System +/// D1 = The dimension of the system +/// D2 = The dimension of the butcher tableau (the order of the integrator) +pub trait Integrator +where + f64: From, + F: SystemTrait, + T: Copy + From + identities::One + identities::Zero + PartialEq + Debug + Add + AddAssign + Sub + SubAssign + Mul + MulAssign + 'static, +{ + /// Returns the tableau + fn tableau(&self) -> Tableau; + + fn _step1(&self, ode: &ODE, h: T) -> (SVector, Option>) { + let k0 = ode.f.derivative(ode.t, ode.y); + let sum = k0 * self.tableau().b[0]; + (ode.y + (sum * h), None) + } + + fn _step2(&self, ode: &ODE, h: T) -> (SVector, Option>) { + let k0 = ode.f.derivative(ode.t, ode.y); + let k1 = ode.f.derivative( + ode.t + self.tableau().c[1] * h, + ode.y + k0 * self.tableau().a[(1,0)] * h, + ); + let sum = + k0 * self.tableau().b[0] + + k1 * self.tableau().b[1]; + let err = match self.tableau().b_star { + Some(b_star) => Some(sum + (k0 * b_star[0] + + k1 * b_star[1]) * T::from(-1.0_f64)), + None => None + }; + (ode.y + (sum * h), err) + } + + fn _step3(&self, ode: &ODE, h: T) -> (SVector, Option>) { + let k0 = ode.f.derivative(ode.t, ode.y); + let k1 = ode.f.derivative( + ode.t + self.tableau().c[1] * h, + ode.y + k0 * self.tableau().a[(1,0)] * h, + ); + let k2 = ode.f.derivative( + ode.t + self.tableau().c[2] * h, + ode.y + ( + k0 * self.tableau().a[(2,0)] + + k1 * self.tableau().a[(2,1)] + ) * h, + ); + let sum = + k0 * self.tableau().b[0] + + k1 * self.tableau().b[1] + + k2 * self.tableau().b[2]; + let err = match self.tableau().b_star { + Some(b_star) => Some(sum + (k0 * b_star[0] + + k1 * b_star[1] + + k2 * b_star[2]) * T::from(-1.0_f64)), + None => None + }; + (ode.y + (sum * h), err) + } + + fn _step4(&self, ode: &ODE, h: T) -> (SVector, Option>) { + let k0 = ode.f.derivative(ode.t, ode.y); + let k1 = ode.f.derivative( + ode.t + self.tableau().c[1] * h, + ode.y + k0 * self.tableau().a[(1,0)] * h, + ); + let k2 = ode.f.derivative( + ode.t + self.tableau().c[2] * h, + ode.y + ( + k0 * self.tableau().a[(2,0)] + + k1 * self.tableau().a[(2,1)] + ) * h, + ); + let k3 = ode.f.derivative( + ode.t + self.tableau().c[3] * h, + ode.y + ( + k0 * self.tableau().a[(3,0)] + + k1 * self.tableau().a[(3,1)] + + k2 * self.tableau().a[(3,2)] + ) * h, + ); + let sum = + k0 * self.tableau().b[0] + + k1 * self.tableau().b[1] + + k2 * self.tableau().b[2] + + k3 * self.tableau().b[3]; + let err = match self.tableau().b_star { + Some(b_star) => Some(sum + (k0 * b_star[0] + + k1 * b_star[1] + + k2 * b_star[2] + + k3 * b_star[3]) * T::from(-1.0_f64)), + None => None + }; + (ode.y + (sum * h), err) + } + + fn _step5(&self, ode: &ODE, h: T) -> (SVector, Option>) { + let k0 = ode.f.derivative(ode.t, ode.y); + let k1 = ode.f.derivative( + ode.t + self.tableau().c[1] * h, + ode.y + k0 * self.tableau().a[(1,0)] * h, + ); + let k2 = ode.f.derivative( + ode.t + self.tableau().c[2] * h, + ode.y + ( + k0 * self.tableau().a[(2,0)] + + k1 * self.tableau().a[(2,1)] + ) * h, + ); + let k3 = ode.f.derivative( + ode.t + self.tableau().c[3] * h, + ode.y + ( + k0 * self.tableau().a[(3,0)] + + k1 * self.tableau().a[(3,1)] + + k2 * self.tableau().a[(3,2)] + ) * h, + ); + let k4 = ode.f.derivative( + ode.t + self.tableau().c[4] * h, + ode.y + ( + k0 * self.tableau().a[(4,0)] + + k1 * self.tableau().a[(4,1)] + + k2 * self.tableau().a[(4,2)] + + k3 * self.tableau().a[(4,3)] + ) * h, + ); + let sum = + k0 * self.tableau().b[0] + + k1 * self.tableau().b[1] + + k2 * self.tableau().b[2] + + k3 * self.tableau().b[3] + + k4 * self.tableau().b[4]; + let err = match self.tableau().b_star { + Some(b_star) => Some(sum + (k0 * b_star[0] + + k1 * b_star[1] + + k2 * b_star[2] + + k3 * b_star[3] + + k4 * b_star[4]) * T::from(-1.0_f64)), + None => None + }; + (ode.y + (sum * h), err) + } + + fn _step6(&self, ode: &ODE, h: T) -> (SVector, Option>) { + let k0 = ode.f.derivative(ode.t, ode.y); + let k1 = ode.f.derivative( + ode.t + self.tableau().c[1] * h, + ode.y + k0 * self.tableau().a[(1,0)] * h, + ); + let k2 = ode.f.derivative( + ode.t + self.tableau().c[2] * h, + ode.y + ( + k0 * self.tableau().a[(2,0)] + + k1 * self.tableau().a[(2,1)] + ) * h, + ); + let k3 = ode.f.derivative( + ode.t + self.tableau().c[3] * h, + ode.y + ( + k0 * self.tableau().a[(3,0)] + + k1 * self.tableau().a[(3,1)] + + k2 * self.tableau().a[(3,2)] + ) * h, + ); + let k4 = ode.f.derivative( + ode.t + self.tableau().c[4] * h, + ode.y + ( + k0 * self.tableau().a[(4,0)] + + k1 * self.tableau().a[(4,1)] + + k2 * self.tableau().a[(4,2)] + + k3 * self.tableau().a[(4,3)] + ) * h, + ); + let k5 = ode.f.derivative( + ode.t + self.tableau().c[5] * h, + ode.y + ( + k0 * self.tableau().a[(5,0)] + + k1 * self.tableau().a[(5,1)] + + k2 * self.tableau().a[(5,2)] + + k3 * self.tableau().a[(5,3)] + + k3 * self.tableau().a[(5,4)] + ) * h, + ); + let sum = + k0 * self.tableau().b[0] + + k1 * self.tableau().b[1] + + k2 * self.tableau().b[2] + + k3 * self.tableau().b[3] + + k4 * self.tableau().b[4] + + k5 * self.tableau().b[5]; + let err = match self.tableau().b_star { + Some(b_star) => Some(sum + (k0 * b_star[0] + + k1 * b_star[1] + + k2 * b_star[2] + + k3 * b_star[3] + + k4 * b_star[4] + + k5 * b_star[5]) * T::from(-1.0_f64)), + None => None + }; + (ode.y + (sum * h), err) + } + + fn step(&self, ode: &ODE, h: T) -> (SVector, Option); + +} + +pub struct FixedFourthOrder +where + f64: From, + T: Copy + From + identities::One + identities::Zero + PartialEq + Debug + Add + AddAssign + Sub + SubAssign + Mul + MulAssign + 'static, +{ + pub tableau: Tableau, +} + +impl Integrator for FixedFourthOrder +where + f64: From, + F: SystemTrait, + T: Copy + From + identities::One + identities::Zero + PartialEq + Debug + Add + AddAssign + Sub + SubAssign + Mul + MulAssign + 'static, +{ + fn tableau(&self) -> Tableau { self.tableau } + + fn step(&self, ode: &ODE, h: T) -> (SVector, Option) { (self._step4(ode, h).0, None) } +} + +pub struct AdaptiveSixthOrder +where + f64: From, + T: Copy + From + identities::One + identities::Zero + PartialEq + Debug + Add + AddAssign + Sub + SubAssign + Mul + MulAssign + 'static, +{ + pub tableau: Tableau, + pub atol: T, + pub rtol: T, +} + +impl Integrator for AdaptiveSixthOrder +where + f64: From, + F: SystemTrait, + T: Copy + + From + + identities::One + + identities::Zero + + PartialEq + + Debug + + Add + + AddAssign + + Sub + + SubAssign + + Mul + + MulAssign + + Float + + 'static, +{ + fn tableau(&self) -> Tableau { self.tableau } + + fn step(&self, ode: &ODE, h: T) -> (SVector, Option) { + let (next_y, err_option) = self._step6(ode, h); + let mut err: T = 0.0.into(); + let err_array = err_option.unwrap(); + ode.y.iter().zip(next_y.iter()).enumerate().for_each(|(i, (yi, next_yi))| { + let tol = self.atol + yi.abs().max(next_yi.abs()) * self.rtol; + err += (err_array[i] / tol).powi(2); + }); + (next_y, Some(err)) + } +} + +pub const RUNGE_KUTTA_4_TABLEAU: Tableau = Tableau { + a: na::Matrix4::new( + 0.0, 0.0, 0.0, 0.0, + 0.5, 0.0, 0.0, 0.0, + 0.0, 0.5, 0.0, 0.0, + 0.0, 0.0, 1.0, 0.0, + ), + b: na::Vector4::new(1.0/6.0, 1.0/3.0, 1.0/3.0, 1.0/6.0), + b_star: None, + c: na::Vector4::new(0.0, 0.5, 0.5, 1.0), +}; + +pub const RUNGE_KUTTA_FEHLBERG_54_TABLEAU: Tableau = Tableau { + a: na::Matrix6::new( + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.25, 0.0, 0.0, 0.0, 0.0, 0.0, + 3.0/32.0, 9.0/32.0, 0.0, 0.0, 0.0, 0.0, + 1932.0/2197.0, -7200.0/2197.0, 7296.0/2197.0, 0.0, 0.0, 0.0, + 439.0/216.0, -8.0, 3680.0/513.0, -845.0/4104.0, 0.0, 0.0, + -8.0/27.0, 2.0, -3544.0/2565.0, 1859.0/4104.0, -11.0/40.0, 0.0, + ), + b: na::Vector6::new( + 16.0/135.0, + 0.0, + 6656.0/12825.0, + 28561.0/56430.0, + -9.0/50.0, + 2.0/55.0, + ), + b_star: Some( + na::Vector6::new( + 25.0/216.0, + 0.0, + 1408.0/2565.0, + 2197.0/4104.0, + -1.0/5.0, + 0.0, + ) + ), + c: na::Vector6::new(0.0, 0.25, 3.0/8.0, 12.0/13.0, 1.0, 0.5), +}; + +#[cfg(test)] +mod tests { + use super::*; + use nalgebra::Vector3; + use approx::assert_relative_eq; + + #[test] + fn test_rk4() { + struct System {} + + impl SystemTrait for System { + fn derivative(&self, _t: f64, y: Vector3) -> Vector3 { y } + } + + let system = System {}; + let y0 = Vector3::new(1.0, 1.0, 1.0); + let mut ode = ODE::new(system, 0.0, 10.0, y0); + + let rk4 = FixedFourthOrder { tableau: RUNGE_KUTTA_4_TABLEAU }; + + // Test that y'(t) = y(t) solves to y(t) = e^t for rk4 + let step = 0.001; + while ode.t < ode.t_end { + ode.y = rk4.step(&ode, step).0; + ode.t += step; + assert_relative_eq!(ode.y[0], ode.t.exp(), max_relative=0.000001); + } + } + + #[test] + fn test_rkf54() { + struct System {} + + impl SystemTrait for System { + fn derivative(&self, _t: f64, y: Vector3) -> Vector3 { y } + } + + let system = System {}; + let y0 = Vector3::new(1.0, 1.0, 1.0); + let mut ode = ODE::new(system, 0.0, 10.0, y0); + + let rkf54 = AdaptiveSixthOrder { + tableau: RUNGE_KUTTA_FEHLBERG_54_TABLEAU, + atol: 1e-3_f64, + rtol: 1e-3_f64, + }; + + // Test that y'(t) = y(t) solves to y(t) = e^t for rkf54 + // and also that the error seems reasonable + let step = 0.001; + while ode.t < ode.t_end { + let (new_y, err) = rkf54.step(&ode, step); + ode.y = new_y; + ode.t += step; + assert_relative_eq!(ode.y[0], ode.t.exp(), max_relative=0.000001); + assert!(err.unwrap() < 0.0001); + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 7d12d9a..f01ed98 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,14 +1,48 @@ -pub fn add(left: usize, right: usize) -> usize { - left + right -} +#![allow(dead_code)] + +pub mod ode; +pub mod integrator; +pub mod controller; +// pub mod callback; +pub mod problem; + #[cfg(test)] mod tests { - use super::*; + use nalgebra::Vector6; + use approx::assert_relative_eq; + use crate::integrator::{AdaptiveSixthOrder, RUNGE_KUTTA_FEHLBERG_54_TABLEAU}; + use crate::controller::{PIController}; + use crate::ode::{SystemTrait, ODE}; + use crate::problem::Problem; #[test] - fn it_works() { - let result = add(2, 2); - assert_eq!(result, 4); + fn test_ode_creation() { + struct System { + mu: f64, + } + + impl SystemTrait for System { + fn derivative(&self, _t: f64, state: Vector6) -> Vector6 { + let acc = -(self.mu * state.fixed_rows<3>(0)) / (state.fixed_rows<3>(0).norm.powi(3)); + Vector6::new(state[3], state[4], state[5], acc[0], acc[1], acc[2]) + } + } + + let system = System { mu: 3.98600441500000e14 }; + let y0 = Vector3::new(1.0, 0.0, 0.0); + let ode = ODE::new(system, 0.0, 10.0, y0); + let controller = PIController::new(0.17, 0.04, 10.0, 0.2, 10.0, 0.9, 1e-4); + let rkf54 = AdaptiveSixthOrder { + tableau: RUNGE_KUTTA_FEHLBERG_54_TABLEAU, + atol: 1e-8_f64, + rtol: 1e-8_f64, + }; + let mut problem = Problem::new(ode, rkf54, controller); + + let solution = problem.solve(); + solution.times.iter().zip(solution.states.iter()).for_each(|(time, state)| { + assert_relative_eq!(state[0], time.exp(), max_relative=1e-7); + }) } } diff --git a/src/ode.rs b/src/ode.rs new file mode 100644 index 0000000..9adf5a4 --- /dev/null +++ b/src/ode.rs @@ -0,0 +1,67 @@ +use nalgebra::SVector; + +/// A System trait. +/// +/// The user will have to define their own system. They are free to add params to their system +/// definition and use those in the derivative function +pub trait SystemTrait { + fn derivative(&self, t: T, y: SVector) -> SVector; +} + +/// The basic ODE object that will be passed around. The type (T) and the size (D) will be +/// determined upon creation of the object +#[derive(Debug, Clone, Copy)] +pub struct ODE where F: SystemTrait{ + pub f: F, + pub y: SVector, + pub t: T, + pub t0: T, + pub t_end: T, + pub h: T, + pub finished: bool, +} + +impl ODE +where + F: SystemTrait, + T: Copy + From, + f64: From, +{ + pub fn new(f: F, t0: T, t_end: T, y0: SVector) -> Self { + Self { + f: f, + y: y0, + t: t0, + t0: t0, + t_end: t_end, + h: 0.001.into(), + finished: false, + } + } +} + + +#[cfg(test)] +mod tests { + use super::*; + use nalgebra::Vector3; + + #[test] + fn test_ode_creation() { + struct System {} + + impl SystemTrait for System { + fn derivative(&self, _t: f64, y: Vector3) -> Vector3 { -y } + } + + let system = System {}; + let y0 = Vector3::new(1.0, 0.0, 0.0); + let ode = ODE::new(system, 0.0, 10.0, y0); + + assert!(ode.f.derivative(0.0, y0) == Vector3::new(-1.0, 0.0, 0.0)); + assert!(ode.y == Vector3::new(1.0, 0.0, 0.0)); + assert!(ode.t == 0.0); + assert!(!ode.finished); + assert!(ode.t_end == 10.0); + } +} diff --git a/src/problem.rs b/src/problem.rs new file mode 100644 index 0000000..d84602d --- /dev/null +++ b/src/problem.rs @@ -0,0 +1,138 @@ +use num_traits::identities; +use num_traits::Float; +use std::fmt::Debug; +use std::cmp::PartialOrd; +use std::ops::{Add, AddAssign, Sub, SubAssign, Mul, MulAssign}; +use nalgebra::SVector; + +use super::ode::{ODE, SystemTrait}; +use super::controller::{Controller, PIController}; +use super::integrator::Integrator; + +pub struct Problem +where + f64: From, + F: SystemTrait, + T: Copy + + From + + identities::One + + identities::Zero + + PartialEq + + Debug + + Add + + AddAssign + + Sub + + SubAssign + + Mul + + MulAssign + + Float + + PartialOrd + + 'static, + S: Integrator, +{ + ode: ODE, + integrator: S, + controller: PIController, +} + +impl Problem +where + f64: From, + F: SystemTrait, + T: Copy + + From + + identities::One + + identities::Zero + + PartialEq + + Debug + + Add + + AddAssign + + Sub + + SubAssign + + Mul + + MulAssign + + Float + + PartialOrd + + 'static, + S: Integrator, +{ + pub fn new(ode: ODE, integrator: S, controller: PIController) -> Self { + Problem { + ode: ode, + integrator: integrator, + controller: controller, + } + } + pub fn solve(&mut self) -> Solution { + let mut times: Vec:: = Vec::new(); + let mut states: Vec::> = Vec::new(); + let mut step: T = self.controller.old_h; + while self.ode.t < self.ode.t_end { + let (mut new_y, mut err_option) = self.integrator.step(&self.ode, step); + times.push(self.ode.t); + states.push(self.ode.y); + match err_option { + Some(mut err) => { + // Adaptive Step Size + let mut accepted: bool = false; + while !accepted { + (accepted, step) = as Controller>::determine_step(&mut self.controller, step, err); + (new_y, err_option) = self.integrator.step(&self.ode, step); + err = err_option.unwrap(); + } + self.controller.old_h = step; + self.ode.y = new_y; + self.ode.t += step; + }, + None => { + // Fixed Step Size + self.ode.y = new_y; + self.ode.t += self.controller.old_h; + }, + }; + } + Solution { + times: times, + states: states, + } + } +} + +pub struct Solution { + pub times: Vec, + pub states: Vec>, +} + +#[cfg(test)] +mod tests { + use super::*; + use nalgebra::Vector3; + use approx::assert_relative_eq; + use crate::integrator::{AdaptiveSixthOrder, RUNGE_KUTTA_FEHLBERG_54_TABLEAU}; + use crate::controller::{PIController}; + + #[test] + fn test_ode_creation() { + struct System {} + + impl SystemTrait for System { + fn derivative(&self, _t: f64, y: Vector3) -> Vector3 { y } + } + + let system = System {}; + let y0 = Vector3::new(1.0, 0.0, 0.0); + let ode = ODE::new(system, 0.0, 10.0, y0); + let controller = PIController::new(0.17, 0.04, 10.0, 0.2, 10.0, 0.9, 1e-4); + let rkf54 = AdaptiveSixthOrder { + tableau: RUNGE_KUTTA_FEHLBERG_54_TABLEAU, + atol: 1e-8_f64, + rtol: 1e-8_f64, + }; + let mut problem = Problem::new(ode, rkf54, controller); + + let solution = problem.solve(); + solution.times.iter().zip(solution.states.iter()).for_each(|(time, state)| { + assert_relative_eq!(state[0], time.exp(), max_relative=1e-7); + }) + } +}