feat: add config structs for each solver (#181)

* feat: add config structs for each solver

* finalise it

* add some tests

Author: martinjrobins

diffsol/src/lib.rs

@@ -196,11 +196,26 @@ use ode_solver::jacobian_update::JacobianUpdate;
 pub use ode_solver::sde::SdeSolverMethod;
 pub use ode_solver::state::{StateRef, StateRefMut};
 pub use ode_solver::{
-    adjoint::AdjointOdeSolverMethod, bdf::Bdf, bdf_state::BdfState, builder::OdeBuilder,
-    checkpointing::Checkpointing, checkpointing::HermiteInterpolator, explicit_rk::ExplicitRk,
-    method::AugmentedOdeSolverMethod, method::OdeSolverMethod, method::OdeSolverStopReason,
-    problem::OdeSolverProblem, sdirk::Sdirk, sdirk_state::RkState,
-    sensitivities::SensitivitiesOdeSolverMethod, state::OdeSolverState, tableau::Tableau,
+    adjoint::AdjointOdeSolverMethod,
+    bdf::Bdf,
+    bdf_state::BdfState,
+    builder::OdeBuilder,
+    checkpointing::Checkpointing,
+    checkpointing::HermiteInterpolator,
+    config::{
+        BdfConfig, ExplicitRkConfig, OdeSolverConfig, OdeSolverConfigMut, OdeSolverConfigRef,
+        SdirkConfig,
+    },
+    explicit_rk::ExplicitRk,
+    method::AugmentedOdeSolverMethod,
+    method::OdeSolverMethod,
+    method::OdeSolverStopReason,
+    problem::OdeSolverProblem,
+    sdirk::Sdirk,
+    sdirk_state::RkState,
+    sensitivities::SensitivitiesOdeSolverMethod,
+    state::OdeSolverState,
+    tableau::Tableau,
 };
 pub use op::constant_op::{ConstantOp, ConstantOpSens, ConstantOpSensAdjoint};
 pub use op::linear_op::{LinearOp, LinearOpSens, LinearOpTranspose};

diffsol/src/ode_solver/bdf.rs

@@ -13,11 +13,12 @@ use serde::Serialize;
 use crate::ode_solver_error;
 use crate::{
     matrix::MatrixRef, nonlinear_solver::root::RootFinder, op::bdf::BdfCallable, scalar::scale,
-    AugmentedOdeEquations, BdfState, DenseMatrix, IndexType, JacobianUpdate, MatrixViewMut,
-    NonLinearOp, NonLinearSolver, OdeEquationsImplicit, OdeSolverMethod, OdeSolverProblem,
-    OdeSolverState, OdeSolverStopReason, Op, Scalar, Vector, VectorRef, VectorView, VectorViewMut,
+    AugmentedOdeEquations, BdfState, DenseMatrix, JacobianUpdate, MatrixViewMut, NonLinearOp,
+    NonLinearSolver, OdeEquationsImplicit, OdeSolverMethod, OdeSolverProblem, OdeSolverState,
+    OdeSolverStopReason, Op, Scalar, Vector, VectorRef, VectorView, VectorViewMut,
 };
 
+use super::config::BdfConfig;
 use super::jacobian_update::SolverState;
 use super::method::AugmentedOdeSolverMethod;
 
@@ -110,6 +111,7 @@ pub struct Bdf<
     root_finder: Option<RootFinder<Eqn::V>>,
     is_state_modified: bool,
     jacobian_update: JacobianUpdate<Eqn::T>,
+    config: BdfConfig<Eqn::T>,
 }
 
 impl<M, Eqn, Nls, AugmentedEqn> Clone for Bdf<'_, Eqn, Nls, M, AugmentedEqn>
@@ -162,6 +164,7 @@ where
             root_finder: self.root_finder.clone(),
             is_state_modified: self.is_state_modified,
             jacobian_update: self.jacobian_update.clone(),
+            config: self.config.clone(),
         }
     }
 }
@@ -176,27 +179,20 @@ where
     for<'b> &'b Eqn::M: MatrixRef<Eqn::M>,
     Nls: NonLinearSolver<Eqn::M>,
 {
-    const NEWTON_MAXITER: IndexType = 4;
-    const MIN_FACTOR: f64 = 0.5;
-    const MAX_FACTOR: f64 = 2.1;
-    const MAX_THRESHOLD: f64 = 2.0;
-    const MIN_THRESHOLD: f64 = 0.9;
-    const MIN_TIMESTEP: f64 = 1e-32;
-    const MAX_ERROR_TEST_FAILS: usize = 40;
-
     pub fn new(
         problem: &'a OdeSolverProblem<Eqn>,
         state: BdfState<Eqn::V, M>,
         nonlinear_solver: Nls,
     ) -> Result<Self, DiffsolError> {
-        Self::_new(problem, state, nonlinear_solver, true)
+        Self::_new(problem, state, nonlinear_solver, true, BdfConfig::default())
     }
 
     fn _new(
         problem: &'a OdeSolverProblem<Eqn>,
         mut state: BdfState<Eqn::V, M>,
         mut nonlinear_solver: Nls,
         integrate_main_eqn: bool,
+        config: BdfConfig<Eqn::T>,
     ) -> Result<Self, DiffsolError> {
         // kappa values for difference orders, taken from Table 1 of [1]
         let kappa = [
@@ -226,7 +222,7 @@ where
         state.check_consistent_with_problem(problem)?;
 
         let mut convergence = Convergence::new(problem.rtol, &problem.atol);
-        convergence.set_max_iter(Self::NEWTON_MAXITER);
+        convergence.set_max_iter(config.maximum_newton_iterations);
 
         let op = if integrate_main_eqn {
             // setup linear solver for first step
@@ -297,6 +293,7 @@ where
             root_finder,
             is_state_modified,
             jacobian_update: JacobianUpdate::default(),
+            config,
         })
     }
 
@@ -305,6 +302,22 @@ where
         problem: &'a OdeSolverProblem<Eqn>,
         augmented_eqn: AugmentedEqn,
         nonlinear_solver: Nls,
+    ) -> Result<Self, DiffsolError> {
+        Self::new_augmented_with_config(
+            state,
+            problem,
+            augmented_eqn,
+            nonlinear_solver,
+            BdfConfig::default(),
+        )
+    }
+
+    pub fn new_augmented_with_config(
+        state: BdfState<Eqn::V, M>,
+        problem: &'a OdeSolverProblem<Eqn>,
+        augmented_eqn: AugmentedEqn,
+        nonlinear_solver: Nls,
+        config: BdfConfig<Eqn::T>,
     ) -> Result<Self, DiffsolError> {
         state.check_sens_consistent_with_problem(problem, &augmented_eqn)?;
 
@@ -313,6 +326,7 @@ where
             state,
             nonlinear_solver,
             augmented_eqn.integrate_main_eqn(),
+            config,
         )?;
 
         ret.state.set_augmented_problem(problem, &augmented_eqn)?;
@@ -456,7 +470,7 @@ where
         self.state.h = new_h;
 
         // if step size too small, then fail
-        if self.state.h.abs() < Eqn::T::from(Self::MIN_TIMESTEP) {
+        if self.state.h.abs() < self.config.minimum_timestep {
             return Err(DiffsolError::from(OdeSolverError::StepSizeTooSmall {
                 time: self.state.t.into(),
             }));
@@ -860,6 +874,15 @@ where
     for<'b> &'b Eqn::M: MatrixRef<Eqn::M>,
 {
     type State = BdfState<Eqn::V, M>;
+    type Config = BdfConfig<Eqn::T>;
+
+    fn config(&self) -> &BdfConfig<Eqn::T> {
+        &self.config
+    }
+
+    fn config_mut(&mut self) -> &mut BdfConfig<Eqn::T> {
+        &mut self.config
+    }
 
     fn order(&self) -> usize {
         self.state.order
@@ -1126,8 +1149,8 @@ where
                 // calculate optimal step size factor as per eq 2.46 of [2]
                 // and reduce step size and try again
                 let mut factor = safety * error_norm.pow(Eqn::T::from(-0.5 / (order as f64 + 1.0)));
-                if factor < Eqn::T::from(Self::MIN_FACTOR) {
-                    factor = Eqn::T::from(Self::MIN_FACTOR);
+                if factor < self.config.minimum_timestep_shrink {
+                    factor = self.config.minimum_timestep_shrink;
                 }
                 let new_h = self._update_step_size(factor)?;
                 self._jacobian_updates(new_h * self.alpha[order], SolverState::ErrorTestFail);
@@ -1138,7 +1161,7 @@ where
                 // update statistics
                 self.statistics.number_of_error_test_failures += 1;
                 if self.statistics.number_of_error_test_failures - old_num_error_test_failures
-                    >= Self::MAX_ERROR_TEST_FAILS
+                    >= self.config.maximum_error_test_failures
                 {
                     return Err(DiffsolError::from(
                         OdeSolverError::TooManyErrorTestFailures {
@@ -1230,14 +1253,14 @@ where
             };
 
             let mut factor = safety * factors[max_index];
-            if factor > Eqn::T::from(Self::MAX_FACTOR) {
-                factor = Eqn::T::from(Self::MAX_FACTOR);
+            if factor > self.config.maximum_timestep_growth {
+                factor = self.config.maximum_timestep_growth;
             }
-            if factor < Eqn::T::from(Self::MIN_FACTOR) {
-                factor = Eqn::T::from(Self::MIN_FACTOR);
+            if factor < self.config.minimum_timestep_shrink {
+                factor = self.config.minimum_timestep_shrink;
             }
-            if factor >= Eqn::T::from(Self::MAX_THRESHOLD)
-                || factor < Eqn::T::from(Self::MIN_THRESHOLD)
+            if factor >= self.config.minimum_timestep_growth
+                || factor < self.config.maximum_timestep_shrink
                 || max_index == 0
                 || max_index == 2
             {
@@ -1305,8 +1328,8 @@ mod test {
         },
         ode_solver::tests::{
             setup_test_adjoint, setup_test_adjoint_sum_squares, test_adjoint,
-            test_adjoint_sum_squares, test_checkpointing, test_interpolate, test_ode_solver,
-            test_problem, test_state_mut, test_state_mut_on_problem,
+            test_adjoint_sum_squares, test_checkpointing, test_config, test_interpolate,
+            test_ode_solver, test_problem, test_state_mut, test_state_mut_on_problem,
         },
         Context, DenseMatrix, FaerLU, FaerMat, FaerSparseLU, FaerSparseMat, MatrixCommon,
         OdeEquations, OdeSolverMethod, Op, Vector, VectorView,
@@ -1321,6 +1344,11 @@ mod test {
         test_state_mut(test_problem::<M>().bdf::<LS>().unwrap());
     }
 
+    #[test]
+    fn bdf_config() {
+        test_config(robertson_ode::<M>(false, 1).0.bdf::<LS>().unwrap());
+    }
+
     #[test]
     fn bdf_test_interpolate() {
         test_interpolate(test_problem::<M>().bdf::<LS>().unwrap());

diffsol/src/ode_solver/config.rs

@@ -0,0 +1,144 @@
+use crate::Scalar;
+
+pub trait OdeSolverConfig<T> {
+    fn as_base_ref(&self) -> OdeSolverConfigRef<'_, T>;
+    fn as_base_mut(&mut self) -> OdeSolverConfigMut<'_, T>;
+}
+pub struct OdeSolverConfigRef<'a, T> {
+    pub minimum_timestep: &'a T,
+    pub maximum_error_test_failures: &'a usize,
+    pub maximum_timestep_growth: &'a T,
+    pub minimum_timestep_shrink: &'a T,
+}
+
+pub struct OdeSolverConfigMut<'a, T> {
+    pub minimum_timestep: &'a mut T,
+    pub maximum_error_test_failures: &'a mut usize,
+    pub maximum_timestep_growth: &'a mut T,
+    pub minimum_timestep_shrink: &'a mut T,
+}
+
+#[derive(Debug, Clone)]
+pub struct BdfConfig<T> {
+    pub minimum_timestep: T,
+    pub maximum_error_test_failures: usize,
+    pub maximum_timestep_growth: T,
+    pub minimum_timestep_growth: T,
+    pub maximum_timestep_shrink: T,
+    pub minimum_timestep_shrink: T,
+    pub maximum_newton_iterations: usize,
+}
+
+impl<T: Scalar> OdeSolverConfig<T> for BdfConfig<T> {
+    fn as_base_ref(&self) -> OdeSolverConfigRef<'_, T> {
+        OdeSolverConfigRef {
+            minimum_timestep: &self.minimum_timestep,
+            maximum_error_test_failures: &self.maximum_error_test_failures,
+            maximum_timestep_growth: &self.maximum_timestep_growth,
+            minimum_timestep_shrink: &self.minimum_timestep_shrink,
+        }
+    }
+
+    fn as_base_mut(&mut self) -> OdeSolverConfigMut<'_, T> {
+        OdeSolverConfigMut {
+            minimum_timestep: &mut self.minimum_timestep,
+            maximum_error_test_failures: &mut self.maximum_error_test_failures,
+            maximum_timestep_growth: &mut self.maximum_timestep_growth,
+            minimum_timestep_shrink: &mut self.minimum_timestep_shrink,
+        }
+    }
+}
+
+impl<T: Scalar> Default for BdfConfig<T> {
+    fn default() -> Self {
+        Self {
+            minimum_timestep: T::from(1e-32),
+            maximum_error_test_failures: 40,
+            maximum_timestep_growth: T::from(2.1),
+            minimum_timestep_growth: T::from(2.0),
+            maximum_timestep_shrink: T::from(0.9),
+            minimum_timestep_shrink: T::from(0.5),
+            maximum_newton_iterations: 4,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct SdirkConfig<T> {
+    pub minimum_timestep: T,
+    pub maximum_error_test_failures: usize,
+    pub maximum_timestep_growth: T,
+    pub minimum_timestep_shrink: T,
+    pub maximum_newton_iterations: usize,
+}
+
+impl<T: Scalar> Default for SdirkConfig<T> {
+    fn default() -> Self {
+        Self {
+            minimum_timestep: T::from(1e-13),
+            maximum_error_test_failures: 40,
+            maximum_timestep_growth: T::from(10.0),
+            minimum_timestep_shrink: T::from(0.2),
+            maximum_newton_iterations: 10,
+        }
+    }
+}
+
+impl<T: Scalar> OdeSolverConfig<T> for SdirkConfig<T> {
+    fn as_base_ref(&self) -> OdeSolverConfigRef<'_, T> {
+        OdeSolverConfigRef {
+            minimum_timestep: &self.minimum_timestep,
+            maximum_error_test_failures: &self.maximum_error_test_failures,
+            maximum_timestep_growth: &self.maximum_timestep_growth,
+            minimum_timestep_shrink: &self.minimum_timestep_shrink,
+        }
+    }
+
+    fn as_base_mut(&mut self) -> OdeSolverConfigMut<'_, T> {
+        OdeSolverConfigMut {
+            minimum_timestep: &mut self.minimum_timestep,
+            maximum_error_test_failures: &mut self.maximum_error_test_failures,
+            maximum_timestep_growth: &mut self.maximum_timestep_growth,
+            minimum_timestep_shrink: &mut self.minimum_timestep_shrink,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ExplicitRkConfig<T> {
+    pub minimum_timestep: T,
+    pub maximum_error_test_failures: usize,
+    pub maximum_timestep_growth: T,
+    pub minimum_timestep_shrink: T,
+}
+
+impl<T: Scalar> Default for ExplicitRkConfig<T> {
+    fn default() -> Self {
+        Self {
+            minimum_timestep: T::from(1e-13),
+            maximum_error_test_failures: 40,
+            maximum_timestep_growth: T::from(10.0),
+            minimum_timestep_shrink: T::from(0.2),
+        }
+    }
+}
+
+impl<T: Scalar> OdeSolverConfig<T> for ExplicitRkConfig<T> {
+    fn as_base_ref(&self) -> OdeSolverConfigRef<'_, T> {
+        OdeSolverConfigRef {
+            minimum_timestep: &self.minimum_timestep,
+            maximum_error_test_failures: &self.maximum_error_test_failures,
+            maximum_timestep_growth: &self.maximum_timestep_growth,
+            minimum_timestep_shrink: &self.minimum_timestep_shrink,
+        }
+    }
+
+    fn as_base_mut(&mut self) -> OdeSolverConfigMut<'_, T> {
+        OdeSolverConfigMut {
+            minimum_timestep: &mut self.minimum_timestep,
+            maximum_error_test_failures: &mut self.maximum_error_test_failures,
+            maximum_timestep_growth: &mut self.maximum_timestep_growth,
+            minimum_timestep_shrink: &mut self.minimum_timestep_shrink,
+        }
+    }
+}