Finished dopri5, interpolation, and callbacks

Cargo.toml

@@ -9,6 +9,7 @@ edition = "2021"
 serde = { version = "1.0", features = ["derive"] }
 nalgebra = { version = "0.32", features = ["serde-serialize"] }
 num-traits = "0.2.15"
+roots = "0.0.8"
 
 [dev-dependencies]
 approx = "0.5"

src/callback.rs

@@ -0,0 +1,34 @@
+use nalgebra::SVector;
+use super::ode::ODE;
+
+/// A function that takes in a time and a state and outputs a single float value
+///
+/// The integration solver will check this function for zero crossings
+#[derive(Clone, Copy)]
+pub struct Callback<'a, const D: usize> {
+    /// The function to check for zero crossings
+    pub event: &'a dyn Fn(f64, SVector<f64,D>) -> f64,
+
+    /// The function to change the ODE
+    pub effect: &'a dyn Fn(ODE<D>) -> ODE<D>,
+}
+
+/// A convenience function for stopping the integration
+pub fn stop<const D: usize>(ode: ODE<D>) -> ODE<D> {
+    let mut new_ode = ode.clone();
+    new_ode.t_end = new_ode.t;
+    new_ode
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_basic_callbacks() {
+        let _value_too_high = Callback {
+            event: &|_: f64, y: SVector<f64,3>| { 10.0 - y[0] },
+            effect: &stop,
+        };
+    }
+}

src/controller.rs

@@ -21,8 +21,6 @@ impl<const D:usize> Controller<D> for PIController {
         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;
-        // let mut h_new = 0.9 * h * err.powf(-1.0 / 5.0);
-        println!("err: {}\th_new: {}", err, h_new);
         if err <= 1.0 {
             // Accept the stepsize
             self.factor_old = err.max(1.0e-4);

src/integrator/dormand_prince.rs

@@ -8,10 +8,11 @@ pub trait DormandPrinceIntegrator {
     const A: &'static [f64];
     const B: &'static [f64];
     const C: &'static [f64];
+    const D: &'static [f64];
 }
 
+#[derive(Debug, Clone, Copy)]
 pub struct DormandPrince45<const D: usize> {
-    k: Vec<SVector<f64,D>>,
     a_tol: f64,
     r_tol: f64,
 }
@@ -19,7 +20,6 @@ pub struct DormandPrince45<const D: usize> {
 impl<const D: usize> DormandPrince45<D> where DormandPrince45<D>: Integrator<D> {
     pub fn new(a_tol: f64, r_tol: f64) -> Self {
         Self {
-            k: vec![SVector::<f64,D>::zeros(); Self::STAGES],
             a_tol: a_tol,
             r_tol: r_tol,
         }
@@ -75,35 +75,61 @@ impl<const D: usize> DormandPrinceIntegrator for DormandPrince45<D> {
         1.0,
         1.0,
     ];
+    const D: &'static [f64] = &[
+        -12715105075.0 / 11282082432.0,
+        0.0,
+        87487479700.0 / 32700410799.0,
+        -10690763975.0 / 1880347072.0,
+        701980252875.0 / 199316789632.0,
+        -1453857185.0 / 822651844.0,
+        69997945.0 / 29380423.0,
+    ];
 }
 
 impl<const D: usize> Integrator<D> for DormandPrince45<D>
 where
     DormandPrince45<D>: DormandPrinceIntegrator,
 {
+    const ORDER: usize = 5;
     const STAGES: usize = 7;
+    const ADAPTIVE: bool = true;
+    const DENSE: bool = true;
 
-    fn step(&mut self, ode: &ODE<D>, h: f64) -> (SVector<f64,D>, Option<f64>) {
+    fn step(&self, ode: &ODE<D>, h: f64) -> (SVector<f64,D>, Option<f64>, Option<Vec<SVector<f64, D>>>) {
+        let mut k: Vec<SVector::<f64,D>> = vec![SVector::<f64,D>::zeros(); Self::STAGES];
         let mut next_y = ode.y.clone();
         let mut err = SVector::<f64, D>::zeros();
+        let mut rcont5 = SVector::<f64, D>::zeros();
         // Do the first of the summations
-        self.k[0] = (ode.f)(ode.t, ode.y);
+        k[0] = (ode.f)(ode.t, ode.y);
-        next_y += self.k[0] * Self::B[0] * h;
+        next_y += k[0] * Self::B[0] * h;
-        err += self.k[0] * (Self::B[0] - Self::B[Self::STAGES]) * h;
+        err += k[0] * (Self::B[0] - Self::B[Self::STAGES]) * h;
+        let rcont1 = ode.y;
+        rcont5 += k[0] * h * Self::D[0];
         // Then the rest
         for i in 1..Self::STAGES {
             // Compute the ks
             let mut y_term = SVector::<f64,D>::zeros();
             for j in 0..i {
-                y_term += self.k[i-j-1] * Self::A[( i * (i - 1) ) / 2 + j];
+                y_term += k[j] * Self::A[( i * (i - 1) ) / 2 + j];
             }
-            self.k[i] = (ode.f)(ode.t + Self::C[i] * h, ode.y + y_term * h);
+            k[i] = (ode.f)(ode.t + Self::C[i] * h, ode.y + y_term * h);
 
             // Use that and bis to calculate the y and error terms
-            next_y += self.k[i] * h * Self::B[i];
+            next_y += k[i] * h * Self::B[i];
-            err += self.k[i] * (Self::B[i] - Self::B[i + Self::STAGES]) * h;
+            err += k[i] * (Self::B[i] - Self::B[i + Self::STAGES]) * h;
+            rcont5 += k[i] * h * Self::D[i];
         }
+        let rcont2 = next_y - ode.y;
+        let rcont3 = h * k[0] - rcont2;
+        let rcont4 = rcont2 - k[Self::STAGES - 1] * h - rcont3;
         let tol = SVector::<f64,D>::repeat(self.a_tol) + ode.y * self.r_tol;
-        (next_y, Some((err.component_div(&tol)).norm()))
+        let rcont = vec![ rcont1, rcont2, rcont3, rcont4, rcont5, ];
+        (next_y, Some((err.component_div(&tol)).norm()), Some(rcont))
+    }
+    fn interpolate(&self, t_start: f64, t_end: f64, dense: &Vec<SVector<f64,D>>, t: f64) -> SVector<f64,D> {
+        let s = (t - t_start)/(t_end - t_start);
+        let s1 = 1.0 - s;
+        dense[0] + (dense[1] + (dense[2] + (dense[3] + dense[4] * s1) * s) * s1) * s
     }
 }

src/integrator/mod.rs

@@ -3,12 +3,18 @@ use nalgebra::SVector;
 use super::ode::ODE;
 
 pub mod dormand_prince;
-pub mod rosenbrock;
+// pub mod rosenbrock;
 
 /// Integrator Trait
 pub trait Integrator<const D: usize> {
+    const ORDER: usize;
     const STAGES: usize;
-    fn step(&mut self, ode: &ODE<D>, h: f64) -> (SVector<f64,D>, Option<f64>);
+    const ADAPTIVE: bool;
+    const DENSE: bool;
+    /// Returns a new y value, then possibly an error value, and possibly a dense output
+    /// coefficient set
+    fn step(&self, ode: &ODE<D>, h: f64) -> (SVector<f64,D>, Option<f64>, Option<Vec<SVector<f64, D>>>);
+    fn interpolate(&self, t_start: f64, t_end: f64, dense: &Vec<SVector<f64,D>>, t: f64) -> SVector<f64,D>;
 }
 
 
@@ -27,13 +33,13 @@ mod tests {
         let y0 = Vector3::new(1.0, 1.0, 1.0);
         let mut ode = ODE::new(&derivative, 0.0, 4.0, y0);
 
-        let mut dp45 = DormandPrince45::new(1e-12_f64, 1e-4_f64);
+        let dp45 = DormandPrince45::new(1e-12_f64, 1e-4_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.0005;
+        let step = 0.001;
         while ode.t < ode.t_end {
-            let (new_y, err) = dp45.step(&ode, step);
+            let (new_y, err, _) = dp45.step(&ode, step);
             ode.y = new_y;
             ode.t += step;
             assert_relative_eq!(ode.y[0], ode.t.exp(), max_relative=0.01);

src/integrator/rosenbrock.rs

@@ -91,11 +91,12 @@ where
     Rodas4<D>: RosenbrockIntegrator,
 {
     const STAGES: usize = 6;
+    const ADAPTIVE: bool = true;
 
     // TODO: Finish this
-    fn step(&mut self, ode: &ODE<D>, h: f64) -> (SVector<f64,D>, Option<f64>) {
+    fn step(&self, ode: &ODE<D>, _h: f64) -> (SVector<f64,D>, Option<f64>) {
-        let mut next_y = ode.y.clone();
+        let next_y = ode.y.clone();
-        let mut err = SVector::<f64, D>::zeros();
+        let err = SVector::<f64, D>::zeros();
         (next_y, Some(err.norm()))
     }
 }

src/lib.rs

@@ -3,7 +3,7 @@
 pub mod ode;
 pub mod integrator;
 pub mod controller;
-// pub mod callback;
+pub mod callback;
 pub mod problem;
 
 
@@ -23,7 +23,6 @@ mod tests {
         // Calculate one period
         let a = 6.7781363e6_f64;
         let period = 2.0 * PI * (a.powi(3)/3.98600441500000e14).sqrt();
-        println!("{}", period);
 
         // Set up the system
         fn derivative(_t: f64, state: Vector6<f64>) -> Vector6<f64> {
@@ -31,32 +30,29 @@ mod tests {
             Vector6::new(state[3], state[4], state[5], acc[0], acc[1], acc[2])
         }
         let y0 = Vector6::new(
-            4.26387250e+06,
+            4.263868426884883e6,
-            5.14619397e+06,
+            5.146189057155391e6,
-            1.13102192e+06,
+            1.1310208421331816e6,
-            -5.92345023e+03,
+            -5923.454461876975,
-            4.49679662e+03,
+            4496.802639690076,
-            1.87038714e+03,
+            1870.3893008991558,
         );
 
         // Integrate
-        let ode = ODE::new(&derivative, 0.0, period, y0);
+        let ode = ODE::new(&derivative, 0.0, 10.0*period, y0);
-        let dp45 = DormandPrince45::new(1e-12_f64, 1e-8_f64);
+        let dp45 = DormandPrince45::new(1e-12_f64, 1e-12_f64);
-        let controller = PIController::new(0.37, 0.04, 10.0, 0.2, 10.0, 0.9, 1e-4);
+        let controller = PIController::new(0.37, 0.04, 10.0, 0.2, 1000.0, 0.9, 0.01);
 
         let mut problem = Problem::new(ode, dp45, controller);
 
         let solution = problem.solve();
 
-        println!("{}", solution.times.len());
+        assert_relative_eq!(solution.times[solution.states.len()-1], 10.0 * period, max_relative=1e-12);
-        // panic!();
+        assert_relative_eq!(solution.states[solution.states.len()-1][0], y0[0], max_relative=1e-9);
-        // TODO: Something still isn't right with these tolerances I think...
+        assert_relative_eq!(solution.states[solution.states.len()-1][1], y0[1], max_relative=1e-9);
-        assert_relative_eq!(solution.times[solution.states.len()-1], period, max_relative=1e-7);
+        assert_relative_eq!(solution.states[solution.states.len()-1][2], y0[2], max_relative=1e-9);
-        assert_relative_eq!(solution.states[solution.states.len()-1][0], y0[0], max_relative=1e-3);
+        assert_relative_eq!(solution.states[solution.states.len()-1][3], y0[3], max_relative=1e-9);
-        assert_relative_eq!(solution.states[solution.states.len()-1][1], y0[1], max_relative=1e-3);
+        assert_relative_eq!(solution.states[solution.states.len()-1][4], y0[4], max_relative=1e-9);
-        assert_relative_eq!(solution.states[solution.states.len()-1][2], y0[2], max_relative=1e-3);
+        assert_relative_eq!(solution.states[solution.states.len()-1][5], y0[5], max_relative=1e-9);
-        assert_relative_eq!(solution.states[solution.states.len()-1][3], y0[3], max_relative=1e-3);
-        assert_relative_eq!(solution.states[solution.states.len()-1][4], y0[4], max_relative=1e-3);
-        assert_relative_eq!(solution.states[solution.states.len()-1][5], y0[5], max_relative=1e-3);
     }
 }

src/ode.rs

@@ -1,13 +1,5 @@
 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<T, const D: usize> {
-    fn derivative(&self, t: T, y: SVector<T,D>) -> SVector<T,D>;
-}
-
 /// 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(Clone, Copy)]

src/problem.rs

@@ -1,9 +1,12 @@
 use nalgebra::SVector;
+use roots::find_root_regula_falsi;
 
 use super::ode::ODE;
 use super::controller::{Controller, PIController};
 use super::integrator::Integrator;
+use super::callback::Callback;
 
+#[derive(Clone)]
 pub struct Problem<'a, const D: usize, S>
 where
     S: Integrator<D>,
@@ -11,56 +14,119 @@ where
     ode: ODE<'a, D>,
     integrator: S,
     controller: PIController,
+    callbacks: Vec<Callback<'a, D>>,
 }
 
 impl<'a, const D: usize, S> Problem<'a,D,S>
 where
-    S: Integrator<D>,
+    S: Integrator<D> + Copy,
 {
     pub fn new(ode: ODE<'a,D>, integrator: S, controller: PIController) -> Self {
         Problem {
             ode: ode,
             integrator: integrator,
             controller: controller,
+            callbacks: Vec::new(),
         }
     }
-    pub fn solve(&mut self) -> Solution<D> {
+    pub fn solve(&mut self) -> Solution<S, D> {
-        let mut times: Vec::<f64> = Vec::new();
+        let mut times: Vec::<f64> = vec![self.ode.t];
-        let mut states: Vec::<SVector<f64,D>> = Vec::new();
+        let mut states: Vec::<SVector<f64,D>> = vec![self.ode.y];
+        let mut dense_coefficients: Vec::<Vec<SVector<f64,D>>> = Vec::new();
         let mut step: f64 = self.controller.old_h;
-        times.push(self.ode.t);
+        let (mut new_y, mut err_option, _) = self.integrator.step(&self.ode, 0.0);
-        states.push(self.ode.y);
-        let (mut new_y, mut err_option) = self.integrator.step(&self.ode, 0.0);
         while self.ode.t < self.ode.t_end {
-            match err_option {
+            let mut dense_option: Option<Vec<SVector<f64,D>>> = None;
-                Some(mut err) => {
+            if S::ADAPTIVE {
-                    // Adaptive Step Size
+                let mut err = err_option.unwrap();
-                    let mut accepted: bool = false;
+                let mut accepted: bool = false;
-                    while !accepted {
+                while !accepted {
-                        (accepted, step) = <PIController as Controller<D>>::determine_step(&mut self.controller, step, err);
+                    // Try a step and if that isn't acceptable, then change the step until it is
-                        (new_y, err_option) = self.integrator.step(&self.ode, step);
+                    (accepted, step) = <PIController as Controller<D>>::determine_step(&mut self.controller, step, err);
-                        err = err_option.unwrap();
+                    (new_y, err_option, dense_option) = self.integrator.step(&self.ode, step);
+                    err = err_option.unwrap();
+                }
+                self.controller.old_h = step;
+                self.controller.h_max = self.controller.h_max.min(self.ode.t_end - self.ode.t - step);
+            } else {
+                // If fixed time step just step forward one step
+                (new_y, _, dense_option) = self.integrator.step(&self.ode, step);
+            }
+            if self.callbacks.len() > 0 {
+                // Check for events occurring
+                for callback in &self.callbacks {
+                    println!("{}", (callback.event)(self.ode.t, self.ode.y) * (callback.event)(self.ode.t + step, new_y));
+                    if (callback.event)(self.ode.t, self.ode.y) * (callback.event)(self.ode.t + step, new_y) < 0.0 {
+                        // If the event crossed zero, then find the root
+                        let f = |test_t| {
+                            let test_y = self.integrator.step(&self.ode, test_t).0;
+                            (callback.event)(self.ode.t + test_t, test_y)
+                        };
+                        let root = find_root_regula_falsi(0.0, step, &f, &mut 1e-12).unwrap();
+                        step = root;
+                        (new_y, _, dense_option) = self.integrator.step(&self.ode, step);
+                        self.ode = (callback.effect)(self.ode);
                     }
-                    self.controller.old_h = step;
+                }
-                    self.controller.h_max = self.controller.h_max.min(self.ode.t_end - self.ode.t - step);
+            }
-                },
-                None => {},
-            };
             self.ode.y = new_y;
             self.ode.t += step;
             times.push(self.ode.t);
             states.push(self.ode.y);
+            // TODO: Implement third order interpolation for non-dense algorithms
+            dense_coefficients.push(dense_option.unwrap());
         }
         Solution {
+            integrator: self.integrator,
             times: times,
             states: states,
+            dense: dense_coefficients,
+        }
+    }
+
+    pub fn with_callback(mut self, callback: Callback<'a, D>) -> Self {
+        self.callbacks.push(callback);
+        Self {
+            ode: self.ode,
+            integrator: self.integrator,
+            controller: self.controller,
+            callbacks: self.callbacks,
         }
     }
 }
 
-pub struct Solution<const D: usize> {
+pub struct Solution<S, const D: usize> where S: Integrator<D> {
+    pub integrator: S,
     pub times: Vec<f64>,
     pub states: Vec<SVector<f64,D>>,
+    pub dense: Vec::<Vec<SVector<f64,D>>>,
+}
+
+impl<S, const D: usize> Solution<S,D> where S: Integrator<D> {
+    pub fn interpolate(&self, t: f64) -> SVector<f64, D> {
+        // First check that the t is within bounds
+        let last = self.times.last().unwrap();
+        let first = self.times.first().unwrap();
+
+        // TODO: Improve these errors
+        let mut times = self.times.clone();
+        if *first > *last { times.reverse(); }
+        if t < *first || t > *last { panic!(); }
+
+        // Then find the two t values closest to the desired t
+        let mut end_index: usize = 0;
+        for (i, time) in self.times.iter().enumerate() {
+            if time > &t {
+                end_index = i;
+                break;
+            }
+        }
+
+        // Then send that to the integrator
+        let t_start = times[end_index - 1];
+        let t_end = times[end_index];
+        self.integrator.interpolate(t_start, t_end, &self.dense[end_index - 1], t)
+    }
 }
 
 #[cfg(test)]
@@ -70,6 +136,7 @@ mod tests {
     use approx::assert_relative_eq;
     use crate::integrator::dormand_prince::DormandPrince45;
     use crate::controller::PIController;
+    use crate::callback::stop;
 
     #[test]
     fn test_problem() {
@@ -83,10 +150,44 @@ mod tests {
         let mut problem = Problem::new(ode, dp45, controller);
 
         let solution = problem.solve();
-        // println!("{}", solution.times.len());
-        // panic!();
         solution.times.iter().zip(solution.states.iter()).for_each(|(time, state)| {
             assert_relative_eq!(state[0], time.exp(), max_relative=1e-2);
         })
     }
+
+    #[test]
+    fn test_with_callback() {
+        fn derivative(_t: f64, y: Vector3<f64>) -> Vector3<f64> { y }
+        let y0 = Vector3::new(1.0, 1.0, 1.0);
+
+        let ode = ODE::new(&derivative, 0.0, 5.0, y0);
+        let dp45 = DormandPrince45::new(1e-12_f64, 1e-5_f64);
+        let controller = PIController::new(0.17, 0.04, 10.0, 0.2, 0.1, 0.9, 1e-8);
+
+        let value_too_high = Callback {
+            event: &|_: f64, y: SVector<f64,3>| { 10.0 - y[0] },
+            effect: &stop,
+        };
+
+        let mut problem = Problem::new(ode, dp45, controller).with_callback(value_too_high);
+        let solution = problem.solve();
+
+        println!("{}", solution.states.last().unwrap()[0]);
+        assert!(solution.states.last().unwrap()[0] == 10.0);
+    }
+
+    #[test]
+    fn test_with_interpolation() {
+        fn derivative(_t: f64, y: Vector3<f64>) -> Vector3<f64> { y }
+        let y0 = Vector3::new(1.0, 1.0, 1.0);
+
+        let ode = ODE::new(&derivative, 0.0, 10.0, y0);
+        let dp45 = DormandPrince45::new(1e-12_f64, 1e-6_f64);
+        let controller = PIController::new(0.17, 0.04, 10.0, 0.2, 0.1, 0.9, 1e-8);
+
+        let mut problem = Problem::new(ode, dp45, controller);
+        let solution = problem.solve();
+
+        assert_relative_eq!(solution.interpolate(8.8)[0], 8.8_f64.exp(), max_relative=1e-6);
+    }
 }