refactor: tidy up output in sens_dense_sensitivities (#186)

* refactor: tidy up output in sens_dense_sensitivities

* add test for invalid t_eval

Author: martinjrobins

diffsol/src/ode_solver/method.rs

@@ -9,79 +9,6 @@ use crate::{
     OdeSolverState, Op, StateRef, StateRefMut, Vector, VectorViewMut,
 };
 
-/// Utility function to write out the solution at a given timepoint
-/// This function is used by the `solve_dense` method to write out the solution at a given timepoint.
-fn dense_write_out<'a, Eqn: OdeEquations + 'a, S: OdeSolverMethod<'a, Eqn>>(
-    s: &S,
-    y_out: &mut <Eqn::V as DefaultDenseMatrix>::M,
-    t_eval: &[Eqn::T],
-    i: usize,
-) -> Result<(), DiffsolError>
-where
-    Eqn::V: DefaultDenseMatrix,
-{
-    let mut y_out = y_out.column_mut(i);
-    let t = t_eval[i];
-    if s.problem().integrate_out {
-        let g = s.interpolate_out(t)?;
-        y_out.copy_from(&g);
-    } else {
-        let y = s.interpolate(t)?;
-        match s.problem().eqn.out() {
-            Some(out) => y_out.copy_from(&out.call(&y, t_eval[i])),
-            None => y_out.copy_from(&y),
-        }
-    }
-    Ok(())
-}
-
-/// utility function to write out the solution at a given timepoint
-/// This function is used by the `solve` method to write out the solution at a given timepoint.
-fn write_out<'a, Eqn: OdeEquations + 'a, S: OdeSolverMethod<'a, Eqn>>(
-    s: &S,
-    ret_y: &mut Vec<Eqn::V>,
-    ret_t: &mut Vec<Eqn::T>,
-) {
-    let t = s.state().t;
-    let y = s.state().y;
-    ret_t.push(t);
-    match s.problem().eqn.out() {
-        Some(out) => {
-            if s.problem().integrate_out {
-                ret_y.push(s.state().g.clone());
-            } else {
-                ret_y.push(out.call(y, t));
-            }
-        }
-        None => ret_y.push(y.clone()),
-    }
-}
-
-fn dense_allocate_return<'a, Eqn: OdeEquations + 'a, S: OdeSolverMethod<'a, Eqn>>(
-    s: &S,
-    t_eval: &[Eqn::T],
-) -> Result<<Eqn::V as DefaultDenseMatrix>::M, DiffsolError>
-where
-    Eqn::V: DefaultDenseMatrix,
-{
-    let nrows = if s.problem().eqn.out().is_some() {
-        s.problem().eqn.out().unwrap().nout()
-    } else {
-        s.problem().eqn.rhs().nstates()
-    };
-    let ret = s
-        .problem()
-        .context()
-        .dense_mat_zeros::<Eqn::V>(nrows, t_eval.len());
-
-    // check t_eval is increasing and all values are greater than or equal to the current time
-    let t0 = s.state().t;
-    if t_eval.windows(2).any(|w| w[0] > w[1] || w[0] < t0) {
-        return Err(ode_solver_error!(InvalidTEval));
-    }
-    Ok(ret)
-}
-
 #[derive(Debug, PartialEq)]
 pub enum OdeSolverStopReason<T: Scalar> {
     InternalTimestep,
@@ -400,9 +327,85 @@ where
     fn augmented_eqn(&self) -> Option<&AugmentedEqn>;
 }
 
+/// Utility function to write out the solution at a given timepoint
+/// This function is used by the `solve_dense` method to write out the solution at a given timepoint.
+fn dense_write_out<'a, Eqn: OdeEquations + 'a, S: OdeSolverMethod<'a, Eqn>>(
+    s: &S,
+    y_out: &mut <Eqn::V as DefaultDenseMatrix>::M,
+    t_eval: &[Eqn::T],
+    i: usize,
+) -> Result<(), DiffsolError>
+where
+    Eqn::V: DefaultDenseMatrix,
+{
+    let mut y_out = y_out.column_mut(i);
+    let t = t_eval[i];
+    if s.problem().integrate_out {
+        let g = s.interpolate_out(t)?;
+        y_out.copy_from(&g);
+    } else {
+        let y = s.interpolate(t)?;
+        match s.problem().eqn.out() {
+            Some(out) => y_out.copy_from(&out.call(&y, t_eval[i])),
+            None => y_out.copy_from(&y),
+        }
+    }
+    Ok(())
+}
+
+/// utility function to write out the solution at a given timepoint
+/// This function is used by the `solve` method to write out the solution at a given timepoint.
+fn write_out<'a, Eqn: OdeEquations + 'a, S: OdeSolverMethod<'a, Eqn>>(
+    s: &S,
+    ret_y: &mut Vec<Eqn::V>,
+    ret_t: &mut Vec<Eqn::T>,
+) {
+    let t = s.state().t;
+    let y = s.state().y;
+    ret_t.push(t);
+    match s.problem().eqn.out() {
+        Some(out) => {
+            if s.problem().integrate_out {
+                ret_y.push(s.state().g.clone());
+            } else {
+                ret_y.push(out.call(y, t));
+            }
+        }
+        None => ret_y.push(y.clone()),
+    }
+}
+
+/// Utility function to allocate the return matrix for the `solve_dense`
+/// and `solve_dense_sensitivities` methods.
+fn dense_allocate_return<'a, Eqn: OdeEquations + 'a, S: OdeSolverMethod<'a, Eqn>>(
+    s: &S,
+    t_eval: &[Eqn::T],
+) -> Result<<Eqn::V as DefaultDenseMatrix>::M, DiffsolError>
+where
+    Eqn::V: DefaultDenseMatrix,
+{
+    let nrows = if s.problem().eqn.out().is_some() {
+        s.problem().eqn.out().unwrap().nout()
+    } else {
+        s.problem().eqn.rhs().nstates()
+    };
+    let ret = s
+        .problem()
+        .context()
+        .dense_mat_zeros::<Eqn::V>(nrows, t_eval.len());
+
+    // check t_eval is increasing and all values are greater than or equal to the current time
+    let t0 = s.state().t;
+    if t_eval.windows(2).any(|w| w[0] > w[1] || w[0] < t0) {
+        return Err(ode_solver_error!(InvalidTEval));
+    }
+    Ok(ret)
+}
+
 #[cfg(test)]
 mod test {
     use crate::{
+        error::{DiffsolError, OdeSolverError},
         matrix::dense_nalgebra_serial::NalgebraMat,
         ode_equations::test_models::exponential_decay::{
             exponential_decay_problem, exponential_decay_problem_adjoint,
@@ -477,6 +480,18 @@ mod test {
         }
     }
 
+    #[test]
+    fn test_t_eval_errors() {
+        let (problem, _soln) = exponential_decay_problem::<NalgebraMat<f64>>(false);
+        let mut s = problem.bdf::<NalgebraLU<f64>>().unwrap();
+        let t_eval = vec![0.0, 1.0, 0.5, 2.0];
+        let err = s.solve_dense(t_eval.as_slice()).unwrap_err();
+        assert!(matches!(
+            err,
+            DiffsolError::OdeSolverError(OdeSolverError::InvalidTEval)
+        ));
+    }
+
     #[test]
     fn test_dense_solve_sensitivities() {
         let (problem, soln) = exponential_decay_problem_sens::<NalgebraMat<f64>>(false);

diffsol/src/ode_solver/sensitivities.rs

@@ -85,12 +85,12 @@ where
                 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
+                    let mut ret_sens = ret_sens[j].column_mut(i);
                     tmp_nparams.set_index(j, Eqn::T::one());
                     out.jac_mul_inplace(&y, *t, s_j, tmp_nout);
-                    s_j.copy_from(tmp_nout);
+                    ret_sens.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);
+                    ret_sens.add_assign(&*tmp_nout);
                     tmp_nparams.set_index(j, Eqn::T::zero());
                 }
             } else {
@@ -107,23 +107,21 @@ where
         }
         let y = self.state().y;
         let s = self.state().s;
-        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
+                let mut ret_sens = ret_sens[j].column_mut(i);
                 tmp_nparams.set_index(j, Eqn::T::one());
                 out.jac_mul_inplace(y, *t, s_j, tmp_nout);
-                s_tmp.copy_from(tmp_nout);
+                ret_sens.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);
+                ret_sens.add_assign(&*tmp_nout);
                 tmp_nparams.set_index(j, Eqn::T::zero());
             }
         } else {