Merge pull request #20 from martinjrobins/docs

#19 document OdeBuilder

Cargo.toml

@@ -52,7 +52,7 @@ diffsl16-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffs
 diffsl17-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm17-0"], optional = true }
 petgraph = "0.6.4"
 faer = "0.18.2"
-sundials-sys = { version = "0.4.0", features = ["ida"], optional = true }
+sundials-sys = { version = "0.4.0", features = ["ida", "static_libraries"], optional = true }
 
 
 [dev-dependencies]

src/jacobian/mod.rs

@@ -166,7 +166,7 @@ mod tests {
 
     #[test]
     fn build_coloring() {
-        let test_triplets = vec![
+        let test_triplets = [
             vec![(0, 0, 1.0), (1, 1, 1.0)],
             vec![(0, 0, 1.0), (0, 1, 1.0), (1, 1, 1.0)],
             vec![(1, 1, 1.0)],

src/lib.rs

@@ -56,7 +56,8 @@ use matrix::{DenseMatrix, Matrix, MatrixViewMut};
 pub use nonlinear_solver::newton::NewtonNonlinearSolver;
 use nonlinear_solver::NonLinearSolver;
 pub use ode_solver::{
-    bdf::Bdf, equations::OdeEquations, OdeSolverMethod, OdeSolverProblem, OdeSolverState,
+    bdf::Bdf, builder::OdeBuilder, equations::OdeEquations, OdeSolverMethod, OdeSolverProblem,
+    OdeSolverState,
 };
 use op::NonLinearOp;
 use scalar::{IndexType, Scalar, Scale};
@@ -68,7 +69,9 @@ pub use scalar::scale;
 #[cfg(test)]
 mod tests {
 
-    use crate::{ode_solver::OdeBuilder, vector::Vector, Bdf, OdeSolverMethod, OdeSolverState};
+    use crate::{
+        ode_solver::builder::OdeBuilder, vector::Vector, Bdf, OdeSolverMethod, OdeSolverState,
+    };
 
     // WARNING: if this test fails and you make a change to the code, you should update the README.md file as well!!!
     #[test]

src/ode_solver/builder.rs

@@ -0,0 +1,268 @@
+use crate::{op::Op, Matrix, OdeSolverProblem, Vector};
+use anyhow::Result;
+
+use super::equations::{OdeSolverEquations, OdeSolverEquationsMassI};
+
+/// Builder for ODE problems. Use methods to set parameters and then call one of the build methods when done.
+pub struct OdeBuilder {
+    t0: f64,
+    h0: f64,
+    rtol: f64,
+    atol: Vec<f64>,
+    p: Vec<f64>,
+    use_coloring: bool,
+}
+
+impl Default for OdeBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl OdeBuilder {
+    /// Create a new builder with default parameters:
+    /// - t0 = 0.0
+    /// - h0 = 1.0
+    /// - rtol = 1e-6
+    /// - atol = [1e-6]
+    /// - p = []
+    /// - use_coloring = false
+    pub fn new() -> Self {
+        Self {
+            t0: 0.0,
+            h0: 1.0,
+            rtol: 1e-6,
+            atol: vec![1e-6],
+            p: vec![],
+            use_coloring: false,
+        }
+    }
+
+    /// Set the initial time.
+    pub fn t0(mut self, t0: f64) -> Self {
+        self.t0 = t0;
+        self
+    }
+
+    /// Set the initial step size.
+    pub fn h0(mut self, h0: f64) -> Self {
+        self.h0 = h0;
+        self
+    }
+
+    /// Set the relative tolerance.
+    pub fn rtol(mut self, rtol: f64) -> Self {
+        self.rtol = rtol;
+        self
+    }
+
+    /// Set the absolute tolerance.
+    pub fn atol<V, T>(mut self, atol: V) -> Self
+    where
+        V: IntoIterator<Item = T>,
+        f64: From<T>,
+    {
+        self.atol = atol.into_iter().map(|x| f64::from(x)).collect();
+        self
+    }
+
+    /// Set the parameters.
+    pub fn p<V, T>(mut self, p: V) -> Self
+    where
+        V: IntoIterator<Item = T>,
+        f64: From<T>,
+    {
+        self.p = p.into_iter().map(|x| f64::from(x)).collect();
+        self
+    }
+
+    /// Set whether to use coloring when computing the Jacobian.
+    /// This can speed up the computation of the Jacobian for large sparse systems.
+    /// However, it relys on the sparsity of the Jacobian being constant,
+    /// and for certain systems it may detect the wrong sparsity pattern.
+    pub fn use_coloring(mut self, use_coloring: bool) -> Self {
+        self.use_coloring = use_coloring;
+        self
+    }
+
+    fn build_atol<V: Vector>(atol: Vec<f64>, nstates: usize) -> Result<V> {
+        if atol.len() == 1 {
+            Ok(V::from_element(nstates, V::T::from(atol[0])))
+        } else if atol.len() != nstates {
+            Err(anyhow::anyhow!(
+                "atol must have length 1 or equal to the number of states"
+            ))
+        } else {
+            let mut v = V::zeros(nstates);
+            for (i, &a) in atol.iter().enumerate() {
+                v[i] = V::T::from(a);
+            }
+            Ok(v)
+        }
+    }
+
+    fn build_p<V: Vector>(p: Vec<f64>) -> V {
+        let mut v = V::zeros(p.len());
+        for (i, &p) in p.iter().enumerate() {
+            v[i] = V::T::from(p);
+        }
+        v
+    }
+
+    /// Build an ODE problem with a mass matrix.
+    ///
+    /// # Arguments
+    ///
+    /// - `rhs`: Function of type Fn(x: &V, p: &V, t: S, y: &mut V) that computes the right-hand side of the ODE.
+    /// - `rhs_jac`: Function of type Fn(x: &V, p: &V, t: S, v: &V, y: &mut V) that computes the multiplication of the Jacobian of the right-hand side with the vector v.
+    /// - `mass`: Function of type Fn(v: &V, p: &V, t: S, y: &mut V) that computes the multiplication of the mass matrix with the vector v.
+    /// - `init`: Function of type Fn(p: &V, t: S) -> V that computes the initial state.
+    ///
+    /// # Generic Arguments
+    ///
+    /// - `M`: Type that implements the `Matrix` trait. Often this must be provided explicitly (i.e. `type M = DMatrix<f64>; builder.build_ode::<M, _, _, _>`).
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use diffsol::OdeBuilder;
+    /// use nalgebra::DVector;
+    /// type M = nalgebra::DMatrix<f64>;
+    ///
+    /// // dy/dt = y
+    /// // 0 = z - y
+    /// // y(0) = 0.1
+    /// // z(0) = 0.1
+    /// let problem = OdeBuilder::new()
+    ///   .build_ode_with_mass::<M, _, _, _, _>(
+    ///       |x, _p, _t, y| {
+    ///           y[0] = x[0];
+    ///           y[1] = x[1] - x[0];
+    ///       },
+    ///       |x, _p, _t, v, y|  {
+    ///           y[0] = v[0];
+    ///           y[1] = v[1] - v[0];
+    ///       },
+    ///       |v, _p, _t, y| {
+    ///           y[0] = v[0];
+    ///           y[1] = 0.0;
+    ///       },
+    ///       |p, _t| DVector::from_element(2, 0.1),
+    /// );
+    /// ```
+    #[allow(clippy::type_complexity)]
+    pub fn build_ode_with_mass<M, F, G, H, I>(
+        self,
+        rhs: F,
+        rhs_jac: G,
+        mass: H,
+        init: I,
+    ) -> Result<OdeSolverProblem<OdeSolverEquations<M, F, G, H, I>>>
+    where
+        M: Matrix,
+        F: Fn(&M::V, &M::V, M::T, &mut M::V),
+        G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
+        H: Fn(&M::V, &M::V, M::T, &mut M::V),
+        I: Fn(&M::V, M::T) -> M::V,
+    {
+        let p = Self::build_p(self.p);
+        let eqn = OdeSolverEquations::new_ode_with_mass(
+            rhs,
+            rhs_jac,
+            mass,
+            init,
+            p,
+            M::T::from(self.t0),
+            self.use_coloring,
+        );
+        let atol = Self::build_atol(self.atol, eqn.nstates())?;
+        Ok(OdeSolverProblem::new(
+            eqn,
+            M::T::from(self.rtol),
+            atol,
+            M::T::from(self.t0),
+            M::T::from(self.h0),
+        ))
+    }
+
+    /// Build an ODE problem with a mass matrix that is the identity matrix.
+    ///
+    /// # Arguments
+    ///
+    /// - `rhs`: Function of type Fn(x: &V, p: &V, t: S, y: &mut V) that computes the right-hand side of the ODE.
+    /// - `rhs_jac`: Function of type Fn(x: &V, p: &V, t: S, v: &V, y: &mut V) that computes the multiplication of the Jacobian of the right-hand side with the vector v.
+    /// - `init`: Function of type Fn(p: &V, t: S) -> V that computes the initial state.
+    ///
+    /// # Generic Arguments
+    ///
+    /// - `M`: Type that implements the `Matrix` trait. Often this must be provided explicitly (i.e. `type M = DMatrix<f64>; builder.build_ode::<M, _, _, _>`).
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use diffsol::OdeBuilder;
+    /// use nalgebra::DVector;
+    /// type M = nalgebra::DMatrix<f64>;
+    ///
+    ///
+    /// // dy/dt = y
+    /// // y(0) = 1
+    /// let problem = OdeBuilder::new()
+    ///    .build_ode::<M, _, _, _>(
+    ///        |x, _p, _t, y| y.copy_from(x),
+    ///        |x, _p, _t, v , y| y.copy_from(v),
+    ///        |p, _t| DVector::from_element(1, 0.1),
+    ///    );
+    /// ```
+    pub fn build_ode<M, F, G, I>(
+        self,
+        rhs: F,
+        rhs_jac: G,
+        init: I,
+    ) -> Result<OdeSolverProblem<OdeSolverEquationsMassI<M, F, G, I>>>
+    where
+        M: Matrix,
+        F: Fn(&M::V, &M::V, M::T, &mut M::V),
+        G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
+        I: Fn(&M::V, M::T) -> M::V,
+    {
+        let p = Self::build_p(self.p);
+        let eqn = OdeSolverEquationsMassI::new_ode(
+            rhs,
+            rhs_jac,
+            init,
+            p,
+            M::T::from(self.t0),
+            self.use_coloring,
+        );
+        let atol = Self::build_atol(self.atol, eqn.nstates())?;
+        Ok(OdeSolverProblem::new(
+            eqn,
+            M::T::from(self.rtol),
+            atol,
+            M::T::from(self.t0),
+            M::T::from(self.h0),
+        ))
+    }
+
+    /// Build an ODE problem using the DiffSL language (requires the `diffsl` feature).
+    /// The source code is provided as a string, please see the [DiffSL documentation](https://martinjrobins.github.io/diffsl/) for more information.
+    #[cfg(feature = "diffsl")]
+    pub fn build_diffsl(
+        self,
+        source: &str,
+    ) -> Result<OdeSolverProblem<crate::ode_solver::diffsl::DiffSl>> {
+        type V = crate::ode_solver::diffsl::V;
+        type T = crate::ode_solver::diffsl::T;
+        let p = Self::build_p::<V>(self.p);
+        let eqn = crate::ode_solver::diffsl::DiffSl::new(source, p, self.use_coloring)?;
+        let atol = Self::build_atol::<V>(self.atol, eqn.nstates())?;
+        Ok(OdeSolverProblem::new(
+            eqn,
+            T::from(self.rtol),
+            atol,
+            T::from(self.t0),
+            T::from(self.h0),
+        ))
+    }
+}

src/ode_solver/diffsl.rs

@@ -214,8 +214,8 @@ mod tests {
     use nalgebra::DVector;
 
     use crate::{
-        linear_solver::lu::LU, nonlinear_solver::newton::NewtonNonlinearSolver,
-        ode_solver::OdeBuilder, Bdf, OdeEquations, OdeSolverMethod, Vector,
+        linear_solver::lu::LU, nonlinear_solver::newton::NewtonNonlinearSolver, Bdf, OdeBuilder,
+        OdeEquations, OdeSolverMethod, Vector,
     };
 
     use super::DiffSl;