feat: sensitivity dense out now supports output ops (#184)

* update ode equations trait and builder

* finish loop, need to fix the last state

* works

* remove odeequationssens as unnecessary

Author: martinjrobins

diffsol/src/lib.rs

@@ -189,8 +189,8 @@ pub use ode_equations::{
     adjoint_equations::AdjointInit, adjoint_equations::AdjointRhs, sens_equations::SensEquations,
     sens_equations::SensInit, sens_equations::SensRhs, AugmentedOdeEquations,
     AugmentedOdeEquationsImplicit, NoAug, OdeEquations, OdeEquationsAdjoint, OdeEquationsImplicit,
-    OdeEquationsImplicitAdjoint, OdeEquationsImplicitSens, OdeEquationsRef, OdeEquationsSens,
-    OdeEquationsStoch, OdeSolverEquations,
+    OdeEquationsImplicitAdjoint, OdeEquationsImplicitSens, OdeEquationsRef, OdeEquationsStoch,
+    OdeSolverEquations,
 };
 use ode_solver::jacobian_update::JacobianUpdate;
 pub use ode_solver::sde::SdeSolverMethod;

diffsol/src/ode_equations/mod.rs

@@ -347,25 +347,11 @@ impl<T> OdeEquationsStoch for T where
 {
 }
 
-pub trait OdeEquationsSens:
-    OdeEquations<
-    Rhs: NonLinearOpSens<M = Self::M, V = Self::V, T = Self::T, C = Self::C>,
-    Init: ConstantOpSens<M = Self::M, V = Self::V, T = Self::T, C = Self::C>,
->
-{
-}
-
-impl<T> OdeEquationsSens for T where
-    T: OdeEquations<
-        Rhs: NonLinearOpSens<M = T::M, V = T::V, T = T::T, C = T::C>,
-        Init: ConstantOpSens<M = T::M, V = T::V, T = T::T, C = T::C>,
-    >
-{
-}
-
 pub trait OdeEquationsImplicitSens:
     OdeEquationsImplicit<
     Rhs: NonLinearOpSens<M = Self::M, V = Self::V, T = Self::T, C = Self::C>,
+    Out: NonLinearOpSens<M = Self::M, V = Self::V, T = Self::T, C = Self::C>
+             + NonLinearOpJacobian<M = Self::M, V = Self::V, T = Self::T, C = Self::C>,
     Init: ConstantOpSens<M = Self::M, V = Self::V, T = Self::T, C = Self::C>,
 >
 {
@@ -374,6 +360,8 @@ pub trait OdeEquationsImplicitSens:
 impl<T> OdeEquationsImplicitSens for T where
     T: OdeEquationsImplicit<
         Rhs: NonLinearOpSens<M = T::M, V = T::V, T = T::T, C = T::C>,
+        Out: NonLinearOpSens<M = T::M, V = T::V, T = T::T, C = T::C>
+                 + NonLinearOpJacobian<M = T::M, V = T::V, T = T::T, C = T::C>,
         Init: ConstantOpSens<M = T::M, V = T::V, T = T::T, C = T::C>,
     >
 {

diffsol/src/ode_equations/sens_equations.rs

@@ -4,7 +4,7 @@ use std::cell::RefCell;
 use crate::{
     op::nonlinear_op::NonLinearOpJacobian, AugmentedOdeEquations, ConstantOp, ConstantOpSens,
     Matrix, NonLinearOp, NonLinearOpSens, OdeEquations, OdeEquationsImplicitSens, OdeEquationsRef,
-    OdeEquationsSens, OdeSolverProblem, Op, Vector,
+    OdeSolverProblem, Op, Vector,
 };
 
 pub struct SensInit<'a, Eqn>
@@ -64,7 +64,7 @@ where
 
 impl<Eqn> ConstantOp for SensInit<'_, Eqn>
 where
-    Eqn: OdeEquationsSens,
+    Eqn: OdeEquationsImplicitSens,
 {
     fn call_inplace(&self, _t: Self::T, y: &mut Self::V) {
         self.eqn.init().sens_mul_inplace(self.t0, &self.tmp, y);

diffsol/src/ode_equations/test_models/exponential_decay.rs

@@ -123,6 +123,18 @@ fn exponential_decay_out_adj_mul<M: MatrixHost>(
     y[1] = -M::T::from(2.0) * v[0] - M::T::from(4.0) * v[1];
 }
 
+/// J = |0 0|
+///     |0 0|
+fn exponential_decay_out_sens<M: MatrixHost>(
+    _x: &M::V,
+    _p: &M::V,
+    _t: M::T,
+    _v: &M::V,
+    y: &mut M::V,
+) {
+    y.fill(M::T::zero());
+}
+
 /// J = |0 0|
 ///     |0 0|
 fn exponential_decay_out_sens_adj<M: MatrixHost>(
@@ -422,6 +434,78 @@ pub fn exponential_decay_problem_sens<M: MatrixHost + 'static>(
     (problem, soln)
 }
 
+#[allow(clippy::type_complexity)]
+pub fn exponential_decay_problem_sens_with_out<M: MatrixHost + 'static>(
+    use_coloring: bool,
+) -> (
+    OdeSolverProblem<impl OdeEquationsImplicitSens<M = M, V = M::V, T = M::T, C = M::C>>,
+    OdeSolverSolution<M::V>,
+) {
+    let k = 0.1;
+    let y0 = 1.0;
+    let problem = OdeBuilder::<M>::new()
+        .p([k, y0])
+        .sens_rtol(1e-6)
+        .sens_atol([1e-6, 1e-6])
+        .use_coloring(use_coloring)
+        .rhs_sens_implicit(
+            exponential_decay::<M>,
+            exponential_decay_jacobian::<M>,
+            exponential_decay_sens::<M>,
+        )
+        .init_sens(
+            exponential_decay_init::<M>,
+            exponential_decay_init_sens::<M>,
+            2,
+        )
+        // g_1 = 1 * x_1  +  2 * x_2
+        // g_2 = 3 * x_1  +  4 * x_2
+        .out_sens_implicit(
+            exponential_decay_out::<M>,
+            exponential_decay_out_jac_mul::<M>,
+            exponential_decay_out_sens::<M>,
+            2,
+        )
+        .build()
+        .unwrap();
+    let p = [M::T::from(k), M::T::from(y0)];
+    let mut soln = OdeSolverSolution::default();
+
+    for i in 0..10 {
+        let t = M::T::from(i as f64);
+        let y0: M::V = problem.eqn.init().call(M::T::zero());
+        let y = y0.clone() * scale(M::T::exp(-p[0] * t));
+        let y_out = M::V::from_vec(
+            vec![
+                M::T::from(1.0) * y[0] + M::T::from(2.0) * y[1],
+                M::T::from(3.0) * y[0] + M::T::from(4.0) * y[1],
+            ],
+            y.context().clone(),
+        );
+        let ypk = y0 * scale(-t * M::T::exp(-p[0] * t));
+        let ypk_out = M::V::from_vec(
+            vec![
+                M::T::from(1.0) * ypk[0] + M::T::from(2.0) * ypk[1],
+                M::T::from(3.0) * ypk[0] + M::T::from(4.0) * ypk[1],
+            ],
+            y.context().clone(),
+        );
+        let ypy0 = M::V::from_vec(
+            vec![(-p[0] * t).exp(), (-p[0] * t).exp()],
+            y.context().clone(),
+        );
+        let ypy0_out = M::V::from_vec(
+            vec![
+                M::T::from(1.0) * ypy0[0] + M::T::from(2.0) * ypy0[1],
+                M::T::from(3.0) * ypy0[0] + M::T::from(4.0) * ypy0[1],
+            ],
+            y.context().clone(),
+        );
+        soln.push_sens(y_out, t, &[ypk_out, ypy0_out]);
+    }
+    (problem, soln)
+}
+
 #[cfg(test)]
 mod tests {
     #[cfg(feature = "diffsl-llvm")]

diffsol/src/ode_solver/builder.rs

@@ -607,6 +607,50 @@ where
         }
     }
 
+    pub fn out_sens_implicit<F, G, H>(
+        self,
+        out: F,
+        out_jac: G,
+        out_sens: H,
+        nout: usize,
+    ) -> OdeBuilder<M, Rhs, Init, Mass, Root, ClosureWithSens<M, F, G, H>>
+    where
+        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, &M::V, &mut M::V),
+    {
+        let nstates = 0;
+        OdeBuilder::<M, Rhs, Init, Mass, Root, ClosureWithSens<M, F, G, H>> {
+            rhs: self.rhs,
+            init: self.init,
+            mass: self.mass,
+            root: self.root,
+            out: Some(ClosureWithSens::new(
+                out,
+                out_jac,
+                out_sens,
+                nstates,
+                nstates,
+                nout,
+                self.ctx.clone(),
+            )),
+            t0: self.t0,
+            h0: self.h0,
+            rtol: self.rtol,
+            atol: self.atol,
+            sens_atol: self.sens_atol,
+            sens_rtol: self.sens_rtol,
+            out_rtol: self.out_rtol,
+            out_atol: self.out_atol,
+            param_rtol: self.param_rtol,
+            param_atol: self.param_atol,
+            p: self.p,
+            use_coloring: self.use_coloring,
+            integrate_out: self.integrate_out,
+            ctx: self.ctx,
+        }
+    }
+
     #[allow(clippy::type_complexity)]
     pub fn out_adjoint_implicit<F, G, H, I>(
         self,

diffsol/src/ode_solver/method.rs

@@ -406,7 +406,7 @@ mod test {
         matrix::dense_nalgebra_serial::NalgebraMat,
         ode_equations::test_models::exponential_decay::{
             exponential_decay_problem, exponential_decay_problem_adjoint,
-            exponential_decay_problem_sens,
+            exponential_decay_problem_sens, exponential_decay_problem_sens_with_out,
         },
         scale, AdjointOdeSolverMethod, DenseMatrix, NalgebraLU, NalgebraVec, OdeEquations,
         OdeSolverMethod, Op, SensitivitiesOdeSolverMethod, Vector, VectorView,
@@ -501,6 +501,30 @@ mod test {
         }
     }
 
+    #[test]
+    fn test_dense_solve_sensitivities_with_out() {
+        let (problem, soln) = exponential_decay_problem_sens_with_out::<NalgebraMat<f64>>(false);
+        let mut s = problem.bdf_sens::<NalgebraLU<f64>>().unwrap();
+
+        let t_eval = soln.solution_points.iter().map(|p| p.t).collect::<Vec<_>>();
+        let (y, sens) = s.solve_dense_sensitivities(t_eval.as_slice()).unwrap();
+        for (i, soln_pt) in soln.solution_points.iter().enumerate() {
+            let y_i = y.column(i).into_owned();
+            y_i.assert_eq_norm(&soln_pt.state, &problem.atol, problem.rtol, 15.0);
+        }
+        for (j, soln_pts) in soln.sens_solution_points.unwrap().iter().enumerate() {
+            for (i, soln_pt) in soln_pts.iter().enumerate() {
+                let sens_i = sens[j].column(i).into_owned();
+                sens_i.assert_eq_norm(
+                    &soln_pt.state,
+                    problem.sens_atol.as_ref().unwrap(),
+                    problem.sens_rtol.unwrap(),
+                    15.0,
+                );
+            }
+        }
+    }
+
     #[test]
     fn test_solve_adjoint() {
         let (problem, soln) = exponential_decay_problem_adjoint::<NalgebraMat<f64>>(true);

diffsol/src/ode_solver/sensitivities.rs

@@ -1,8 +1,11 @@
 use crate::{
     error::DiffsolError, error::OdeSolverError, ode_solver_error, AugmentedOdeSolverMethod,
-    Context, DefaultDenseMatrix, DefaultSolver, DenseMatrix, OdeEquationsImplicitSens,
-    OdeSolverStopReason, Op, SensEquations, VectorViewMut,
+    Context, DefaultDenseMatrix, DefaultSolver, DenseMatrix, NonLinearOp, NonLinearOpJacobian,
+    NonLinearOpSens, OdeEquationsImplicitSens, OdeSolverStopReason, Op, SensEquations, Vector,
+    VectorViewMut,
 };
+use num_traits::{One, Zero};
+use std::ops::AddAssign;
 
 pub trait SensitivitiesOdeSolverMethod<'a, Eqn>:
     AugmentedOdeSolverMethod<'a, Eqn, SensEquations<'a, Eqn>>
@@ -36,7 +39,19 @@ where
                 "Cannot integrate out when solving for sensitivities"
             ));
         }
-        let nrows = self.problem().eqn.rhs().nstates();
+        let (mut tmp_nout, mut tmp_nparms, nrows) = if let Some(out) = self.problem().eqn.out() {
+            (
+                Some(Eqn::V::zeros(out.nout(), self.problem().context().clone())),
+                Some(Eqn::V::zeros(
+                    out.nparams(),
+                    self.problem().context().clone(),
+                )),
+                out.nout(),
+            )
+        } else {
+            (None, None, self.problem().eqn.rhs().nout())
+        };
+
         let mut ret = self
             .problem()
             .context()
@@ -62,10 +77,27 @@ where
                 step_reason = self.step()?;
             }
             let y = self.interpolate(*t)?;
-            ret.column_mut(i).copy_from(&y);
-            let s = self.interpolate_sens(*t)?;
-            for (j, s_j) in s.iter().enumerate() {
-                ret_sens[j].column_mut(i).copy_from(s_j);
+            let mut s = self.interpolate_sens(*t)?;
+            if let Some(out) = self.problem().eqn.out() {
+                let tmp_nout = tmp_nout.as_mut().unwrap();
+                let tmp_nparams = tmp_nparms.as_mut().unwrap();
+                out.call_inplace(&y, *t, tmp_nout);
+                ret.column_mut(i).copy_from(tmp_nout);
+                for (j, s_j) in s.iter_mut().enumerate() {
+                    // compute J * s_j + dF/dp * e_j where e_j is the jth basis vector
+                    tmp_nparams.set_index(j, Eqn::T::one());
+                    out.jac_mul_inplace(&y, *t, s_j, tmp_nout);
+                    s_j.copy_from(tmp_nout);
+                    out.sens_mul_inplace(&y, *t, tmp_nparams, tmp_nout);
+                    s_j.add_assign(&*tmp_nout);
+                    ret_sens[j].column_mut(i).copy_from(s_j);
+                    tmp_nparams.set_index(j, Eqn::T::zero());
+                }
+            } else {
+                ret.column_mut(i).copy_from(&y);
+                for (j, s_j) in s.iter().enumerate() {
+                    ret_sens[j].column_mut(i).copy_from(s_j);
+                }
             }
         }
 
@@ -74,11 +106,33 @@ where
             step_reason = self.step()?;
         }
         let y = self.state().y;
-        ret.column_mut(t_eval.len() - 1).copy_from(y);
         let s = self.state().s;
-        for (j, s_j) in s.iter().enumerate() {
-            ret_sens[j].column_mut(t_eval.len() - 1).copy_from(s_j);
+        let mut s_tmp = tmp_nout.clone();
+        let i = t_eval.len() - 1;
+        let t = t_eval.last().unwrap();
+        if let Some(out) = self.problem().eqn.out() {
+            let tmp_nout = tmp_nout.as_mut().unwrap();
+            let tmp_nparams = tmp_nparms.as_mut().unwrap();
+            let s_tmp = s_tmp.as_mut().unwrap();
+            out.call_inplace(y, *t, tmp_nout);
+            ret.column_mut(i).copy_from(tmp_nout);
+            for (j, s_j) in s.iter().enumerate() {
+                // compute J * s_j + dF/dp * e_j where e_j is the jth basis vector
+                tmp_nparams.set_index(j, Eqn::T::one());
+                out.jac_mul_inplace(y, *t, s_j, tmp_nout);
+                s_tmp.copy_from(tmp_nout);
+                out.sens_mul_inplace(y, *t, tmp_nparams, tmp_nout);
+                s_tmp.add_assign(&*tmp_nout);
+                ret_sens[j].column_mut(i).copy_from(s_tmp);
+                tmp_nparams.set_index(j, Eqn::T::zero());
+            }
+        } else {
+            ret.column_mut(i).copy_from(y);
+            for (j, s_j) in s.iter().enumerate() {
+                ret_sens[j].column_mut(i).copy_from(s_j);
+            }
         }
+
         Ok((ret, ret_sens))
     }
 }