new closures to keep up with diffsol

Cargo.toml

@@ -16,6 +16,7 @@ exclude = [".github/*", ".vscode/*"]
 
 [dependencies]
 diffsol = { git = "https://github.com/nabla-rs/diffsol.git", branch = "pmcore" }
+# diffsol = "0.1.5"
 dashmap = "5.5.3"
 lazy_static = "1.4.0"
 csv = "1.2.1"
@@ -43,6 +44,7 @@ faer = "0.18.2"
 faer-ext = {version = "0.1.0", features = ["nalgebra", "ndarray"]}
 serde_with = "3.8.1"
 
+
 [profile.release]
 codegen-units = 1
 opt-level = 3

examples/bimodal_ke/config.toml

@@ -5,7 +5,7 @@ log = "bimodal_ke.log"
 
 [config]
 cycles = 1024
-engine = "NPOD"
+engine = "NPAG"
 init_points = 2129
 seed = 22
 tui = false

src/simulator/ode/closure.rs

@@ -0,0 +1,122 @@
+use std::{cell::RefCell, rc::Rc};
+
+use diffsol::{
+    jacobian::{find_non_zeros_nonlinear, JacobianColoring},
+    matrix::{Matrix, MatrixSparsity},
+    op::{NonLinearOp, Op, OpStatistics},
+    vector::Vector,
+};
+
+use crate::prelude::data::{Covariates, Infusion};
+pub struct PMClosure<M, F>
+where
+    M: Matrix,
+    F: Fn(&M::V, &M::V, M::T, &mut M::V, M::V, &Covariates),
+{
+    func: F,
+    nstates: usize,
+    nout: usize,
+    nparams: usize,
+    p: Rc<M::V>,
+    coloring: Option<JacobianColoring<M>>,
+    sparsity: Option<M::Sparsity>,
+    statistics: RefCell<OpStatistics>,
+    covariates: Covariates,
+    infusions: Vec<Infusion>,
+}
+
+impl<M, F> PMClosure<M, F>
+where
+    M: Matrix,
+    F: Fn(&M::V, &M::V, M::T, &mut M::V, M::V, &Covariates),
+{
+    pub fn new(
+        func: F,
+        nstates: usize,
+        nout: usize,
+        p: Rc<M::V>,
+        covariates: Covariates,
+        infusions: Vec<Infusion>,
+    ) -> Self {
+        let nparams = p.len();
+        Self {
+            func,
+            nstates,
+            nout,
+            nparams,
+            p,
+            statistics: RefCell::new(OpStatistics::default()),
+            coloring: None,
+            sparsity: None,
+            covariates: covariates,
+            infusions: infusions,
+        }
+    }
+    pub fn calculate_sparsity(&mut self, y0: &M::V, t0: M::T) {
+        let non_zeros = find_non_zeros_nonlinear(self, y0, t0);
+        self.sparsity = Some(
+            MatrixSparsity::try_from_indices(self.nout(), self.nstates(), non_zeros.clone())
+                .expect("invalid sparsity pattern"),
+        );
+        self.coloring = Some(JacobianColoring::new_from_non_zeros(self, non_zeros));
+    }
+}
+
+impl<M, F> Op for PMClosure<M, F>
+where
+    M: Matrix,
+    F: Fn(&M::V, &M::V, M::T, &mut M::V, M::V, &Covariates),
+{
+    type V = M::V;
+    type T = M::T;
+    type M = M;
+    fn nstates(&self) -> usize {
+        self.nstates
+    }
+    fn nout(&self) -> usize {
+        self.nout
+    }
+    fn nparams(&self) -> usize {
+        self.nparams
+    }
+    fn sparsity(&self) -> Option<&<Self::M as Matrix>::Sparsity> {
+        self.sparsity.as_ref()
+    }
+    fn statistics(&self) -> OpStatistics {
+        self.statistics.borrow().clone()
+    }
+}
+
+impl<M, F> NonLinearOp for PMClosure<M, F>
+where
+    M: Matrix,
+    F: Fn(&M::V, &M::V, M::T, &mut M::V, M::V, &Covariates),
+{
+    fn call_inplace(&self, x: &M::V, t: M::T, y: &mut M::V) {
+        let mut rateiv = Self::V::zeros(self.nstates);
+        //TODO: This should be pre-calculated
+        for infusion in &self.infusions {
+            if t >= Self::T::from(infusion.time)
+                && t <= Self::T::from(infusion.duration + infusion.time)
+            {
+                rateiv[infusion.input] = Self::T::from(infusion.amount / infusion.duration);
+            }
+        }
+        self.statistics.borrow_mut().increment_call();
+        (self.func)(x, self.p.as_ref(), t, y, rateiv, &self.covariates)
+    }
+    fn jac_mul_inplace(&self, _x: &M::V, t: M::T, v: &M::V, y: &mut M::V) {
+        let rateiv = Self::V::zeros(self.nstates);
+        self.statistics.borrow_mut().increment_jac_mul();
+        (self.func)(v, self.p.as_ref(), t, y, rateiv, &self.covariates);
+        // (self.jacobian_action)(x, self.p.as_ref(), t, v, y)
+    }
+    fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
+        self.statistics.borrow_mut().increment_matrix();
+        if let Some(coloring) = self.coloring.as_ref() {
+            coloring.jacobian_inplace(self, x, t, y);
+        } else {
+            self._default_jacobian_inplace(x, t, y);
+        }
+    }
+}

src/simulator/ode/diffsol_traits.rs

@@ -1,29 +1,25 @@
 use crate::routines::data::{Covariates, Infusion};
 use diffsol::{
-    jacobian::{find_non_zero_entries, JacobianColoring},
     matrix::Matrix,
-    ode_solver::equations::OdeEquationsStatistics,
     ode_solver::problem::OdeSolverProblem,
-    op::{closure::Closure, Op},
-    vector::{Vector, VectorIndex},
+    op::{unit::UnitCallable, Op},
+    vector::Vector,
     OdeBuilder, OdeEquations, Result, Zero,
 };
-use std::{cell::RefCell, rc::Rc};
+use std::rc::Rc;
+
+use super::closure::PMClosure;
 
 pub struct OdePmSolverEquationsMassI<M, F, I>
 where
     M: Matrix,
     F: Fn(&M::V, &M::V, M::T, &mut M::V, M::V, &Covariates),
     I: Fn(&M::V, M::T) -> M::V,
 {
-    rhs: F,
+    rhs: Rc<PMClosure<M, F>>,
+    mass: Rc<UnitCallable<M>>,
     init: I,
     p: Rc<M::V>,
-    nstates: usize,
-    coloring: Option<JacobianColoring>,
-    statistics: RefCell<OdeEquationsStatistics>,
-    covariates: Covariates,
-    infusions: Vec<Infusion>,
 }
 
 impl<M, F, I> OdePmSolverEquationsMassI<M, F, I>
@@ -37,95 +33,50 @@ where
         init: I,
         p: M::V,
         t0: M::T,
-        use_coloring: bool,
+        calculate_sparsity: bool,
         covariates: Covariates,
         infusions: Vec<Infusion>,
     ) -> Self {
         let y0 = init(&p, M::T::zero());
         let nstates = y0.len();
         let p = Rc::new(p);
-
-        let statistics = RefCell::default();
-        let mut ret = Self {
+        let mut rhs = PMClosure::new(rhs, nstates, nstates, p.clone(), covariates, infusions);
+        if calculate_sparsity {
+            rhs.calculate_sparsity(&y0, t0);
+        }
+        let mass = UnitCallable::<M>::new(nstates);
+        let rhs = Rc::new(rhs);
+        let mass = Rc::new(mass);
+        Self {
             rhs,
+            mass,
             init,
             p: p.clone(),
-            nstates,
-            coloring: None,
-            statistics,
-            covariates,
-            infusions,
-        };
-        let coloring = if use_coloring {
-            let rhs_inplace = |x: &M::V, _p: &M::V, t: M::T, y_rhs: &mut M::V| {
-                ret.rhs_inplace(t, x, y_rhs);
-            };
-            let rhs_jac_inplace = |x: &M::V, _p: &M::V, t: M::T, v: &M::V, y: &mut M::V| {
-                ret.rhs_jac_inplace(t, x, v, y);
-            };
-            let op =
-                Closure::<M, _, _>::new(rhs_inplace, rhs_jac_inplace, nstates, nstates, p.clone());
-            Some(JacobianColoring::new(&op, &y0, t0))
-        } else {
-            None
-        };
-        ret.coloring = coloring;
-        ret
+        }
     }
 }
 
-// impl Op
-impl<M, F, I> Op for OdePmSolverEquationsMassI<M, F, I>
+impl<M, F, I> OdeEquations for OdePmSolverEquationsMassI<M, F, I>
 where
     M: Matrix,
     F: Fn(&M::V, &M::V, M::T, &mut M::V, M::V, &Covariates),
     I: Fn(&M::V, M::T) -> M::V,
 {
-    type M = M;
-    type V = M::V;
     type T = M::T;
-
-    fn nout(&self) -> usize {
-        self.nstates
-    }
-    fn nparams(&self) -> usize {
-        self.p.len()
-    }
-    fn nstates(&self) -> usize {
-        self.nstates
+    type V = M::V;
+    type M = M;
+    type Rhs = PMClosure<M, F>;
+    type Mass = UnitCallable<M>;
+    fn mass(&self) -> &Rc<Self::Mass> {
+        &self.mass
     }
-}
 
-impl<M, F, I> OdeEquations for OdePmSolverEquationsMassI<M, F, I>
-where
-    M: Matrix,
-    F: Fn(&M::V, &M::V, M::T, &mut M::V, M::V, &Covariates),
-    I: Fn(&M::V, M::T) -> M::V,
-{
-    fn rhs_inplace(&self, t: Self::T, y: &Self::V, rhs_y: &mut Self::V) {
-        let p = self.p.as_ref();
-        let mut rateiv = Self::V::zeros(self.nstates);
-        //TODO: This should be pre-calculated
-        for infusion in &self.infusions {
-            if t >= Self::T::from(infusion.time)
-                && t <= Self::T::from(infusion.duration + infusion.time)
-            {
-                rateiv[infusion.input] = Self::T::from(infusion.amount / infusion.duration);
-            }
-        }
-        // (self.sec_eq)(&mut p, &self.covariates);
-        (self.rhs)(y, &p, t, rhs_y, rateiv, &self.covariates);
-        self.statistics.borrow_mut().number_of_rhs_evals += 1;
+    fn rhs(&self) -> &Rc<Self::Rhs> {
+        &self.rhs
     }
 
-    fn rhs_jac_inplace(&self, t: Self::T, _x: &Self::V, v: &Self::V, y: &mut Self::V) {
-        //TODO: Instead of this code, we should call rhs_inplace
-        let p = self.p.as_ref();
-        let rateiv = Self::V::zeros(self.nstates);
-        //TODO: This should be pre-calculated
-        (self.rhs)(v, &p, t, y, rateiv, &self.covariates);
-        // (self.rhs_jac)(x, p, t, v, y);
-        self.statistics.borrow_mut().number_of_jac_mul_evals += 1;
+    fn is_mass_constant(&self) -> bool {
+        true
     }
 
     fn init(&self, t: Self::T) -> Self::V {
@@ -136,37 +87,6 @@ where
     fn set_params(&mut self, p: Self::V) {
         self.p = Rc::new(p);
     }
-
-    fn algebraic_indices(&self) -> <Self::V as Vector>::Index {
-        <Self::V as Vector>::Index::zeros(0)
-    }
-
-    fn jacobian_matrix(&self, x: &Self::V, t: Self::T) -> Self::M {
-        self.statistics.borrow_mut().number_of_jacobian_matrix_evals += 1;
-        let rhs_inplace = |x: &Self::V, _p: &Self::V, t: Self::T, y_rhs: &mut Self::V| {
-            self.rhs_inplace(t, x, y_rhs);
-        };
-        let rhs_jac_inplace =
-            |x: &Self::V, _p: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V| {
-                self.rhs_jac_inplace(t, x, v, y);
-            };
-        let op = Closure::<M, _, _>::new(
-            rhs_inplace,
-            rhs_jac_inplace,
-            self.nstates,
-            self.nstates,
-            self.p.clone(),
-        );
-        let triplets = if let Some(coloring) = &self.coloring {
-            coloring.find_non_zero_entries(&op, x, t)
-        } else {
-            find_non_zero_entries(&op, x, t)
-        };
-        Self::M::try_from_triplets(self.nstates(), self.nout(), triplets).unwrap()
-    }
-    fn get_statistics(&self) -> OdeEquationsStatistics {
-        self.statistics.borrow().clone()
-    }
 }
 
 pub trait BuildPmOde {
@@ -206,7 +126,8 @@ impl BuildPmOde for OdeBuilder {
             cov,
             infusions,
         );
-        let atol = Self::build_atol(self.atol, eqn.nstates())?;
+
+        let atol = Self::build_atol(self.atol, eqn.rhs().nstates())?;
         Ok(OdeSolverProblem::new(
             eqn,
             M::T::from(self.rtol),

src/simulator/ode/mod.rs

@@ -1,3 +1,4 @@
+pub mod closure;
 pub mod diffsol_traits;
 use crate::{
     routines::data::{Covariates, Infusion},