Finished dopri5, interpolation, and callbacks

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> {
-        let mut times: Vec::<f64> = Vec::new();
-        let mut states: Vec::<SVector<f64,D>> = Vec::new();
+    pub fn solve(&mut self) -> Solution<S, D> {
+        let mut times: Vec::<f64> = vec![self.ode.t];
+        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);
-        states.push(self.ode.y);
-        let (mut new_y, mut err_option) = self.integrator.step(&self.ode, 0.0);
+        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 {
-                Some(mut err) => {
-                    // Adaptive Step Size
-                    let mut accepted: bool = false;
-                    while !accepted {
-                        (accepted, step) = <PIController as Controller<D>>::determine_step(&mut self.controller, step, err);
-                        (new_y, err_option) = self.integrator.step(&self.ode, step);
-                        err = err_option.unwrap();
+            let mut dense_option: Option<Vec<SVector<f64,D>>> = None;
+            if S::ADAPTIVE {
+                let mut err = err_option.unwrap();
+                let mut accepted: bool = false;
+                while !accepted {
+                    // Try a step and if that isn't acceptable, then change the step until it is
+                    (accepted, step) = <PIController as Controller<D>>::determine_step(&mut self.controller, step, err);
+                    (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);
+    }
 }