perf: avoid Jacobian resets for start/endpoint checkpoints (#237)

* Avoid Jacobian reset for initial and final checkpoints

* update snapshots

Author: martinjrobins

diffsol/src/ode_solver/bdf.rs

@@ -140,7 +140,6 @@ where
         let op = if let Some(op) = self.op.as_ref() {
             let op = op.clone_state(&self.ode_problem.eqn);
             nonlinear_solver.set_problem(&op);
-            nonlinear_solver.reset_jacobian(&op, &self.state.y, self.state.t);
             Some(op)
         } else {
             None
@@ -1138,6 +1137,10 @@ where
         self.state.clone()
     }
 
+    fn state_clone(&self) -> Self::State {
+        self.state.clone()
+    }
+
     fn step(&mut self) -> Result<OdeSolverStopReason<Eqn::T>, DiffsolError> {
         debug!(
             "Taking BDF step at time {} with step size {} and order {}",
@@ -1663,8 +1666,8 @@ mod test {
         test_adjoint(adjoint_solver, dgdu);
         insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
         number_of_calls: 160
-        number_of_jac_muls: 6
-        number_of_matrix_evals: 3
+        number_of_jac_muls: 2
+        number_of_matrix_evals: 1
         number_of_jac_adj_muls: 220
         "###);
     }
@@ -1685,8 +1688,8 @@ mod test {
         test_adjoint_sum_squares(adjoint_solver, dgdp, soln, data, times.as_slice());
         insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
         number_of_calls: 588
-        number_of_jac_muls: 10
-        number_of_matrix_evals: 5
+        number_of_jac_muls: 6
+        number_of_matrix_evals: 3
         number_of_jac_adj_muls: 1054
         "###);
     }
@@ -1737,8 +1740,8 @@ mod test {
         test_adjoint(adjoint_solver, dgdu);
         insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
         number_of_calls: 163
-        number_of_jac_muls: 18
-        number_of_matrix_evals: 6
+        number_of_jac_muls: 12
+        number_of_matrix_evals: 4
         number_of_jac_adj_muls: 106
         "###);
     }
@@ -1759,8 +1762,8 @@ mod test {
         test_adjoint_sum_squares(adjoint_solver, dgdp, soln, data, times.as_slice());
         insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
         number_of_calls: 220
-        number_of_jac_muls: 18
-        number_of_matrix_evals: 6
+        number_of_jac_muls: 12
+        number_of_matrix_evals: 4
         number_of_jac_adj_muls: 404
         "###);
     }

diffsol/src/ode_solver/explicit_rk.rs

@@ -159,6 +159,10 @@ where
         self.rk.checkpoint()
     }
 
+    fn state_clone(&self) -> Self::State {
+        self.rk.state().clone()
+    }
+
     fn step(&mut self) -> Result<OdeSolverStopReason<Eqn::T>, DiffsolError> {
         let mut h = self.rk.start_step()?;
         debug!(

diffsol/src/ode_solver/method.rs

@@ -56,6 +56,9 @@ where
     /// Note that this will force a reinitialisation of the internal Jacobian for the solver, if it has one.
     fn checkpoint(&mut self) -> Self::State;
 
+    /// Clone the current state of the solver without triggering any internal Jacobian reset.
+    fn state_clone(&self) -> Self::State;
+
     /// Replace the current state of the solver with a new state.
     fn set_state(&mut self, state: Self::State);
 
@@ -329,7 +332,7 @@ where
         // allocate checkpoint info
         let mut nsteps = 0;
         let t0 = self.state().t;
-        let mut checkpoints = vec![self.checkpoint()];
+        let mut checkpoints = vec![self.state_clone()];
         let mut ts = vec![t0];
         let mut ys = vec![self.state().y.clone()];
         let mut ydots = vec![self.state().dy.clone()];
@@ -361,7 +364,7 @@ where
         ts.push(self.state().t);
         ys.push(self.state().y.clone());
         ydots.push(self.state().dy.clone());
-        checkpoints.push(self.checkpoint());
+        checkpoints.push(self.state_clone());
 
         // construct checkpointing
         let last_segment = HermiteInterpolator::new(ys, ydots, ts);
@@ -420,7 +423,7 @@ where
         // allocate checkpoint info
         let mut nsteps = 0;
         let t0 = self.state().t;
-        let mut checkpoints = vec![self.checkpoint()];
+        let mut checkpoints = vec![self.state_clone()];
         let mut ts = vec![t0];
         let mut ys = vec![self.state().y.clone()];
         let mut ydots = vec![self.state().dy.clone()];
@@ -450,7 +453,7 @@ where
         assert_eq!(step_reason, OdeSolverStopReason::TstopReached);
 
         // add final checkpoint
-        checkpoints.push(self.checkpoint());
+        checkpoints.push(self.state_clone());
 
         // construct the adjoint equations
         let last_segment = HermiteInterpolator::new(ys, ydots, ts);

diffsol/src/ode_solver/sdirk.rs

@@ -337,6 +337,10 @@ where
         self.rk.state().clone()
     }
 
+    fn state_clone(&self) -> Self::State {
+        self.rk.state().clone()
+    }
+
     fn step(&mut self) -> Result<OdeSolverStopReason<Eqn::T>, DiffsolError> {
         debug!(
             "Taking SDIRK step at time {}, step size {}",
@@ -695,8 +699,8 @@ mod test {
         test_adjoint(adjoint_solver, dgdu);
         insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
         number_of_calls: 419
-        number_of_jac_muls: 12
-        number_of_matrix_evals: 6
+        number_of_jac_muls: 10
+        number_of_matrix_evals: 5
         number_of_jac_adj_muls: 271
         "###);
     }
@@ -717,8 +721,8 @@ mod test {
         test_adjoint_sum_squares(adjoint_solver, dgdp, soln, data, times.as_slice());
         insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
         number_of_calls: 362
-        number_of_jac_muls: 12
-        number_of_matrix_evals: 4
+        number_of_jac_muls: 9
+        number_of_matrix_evals: 3
         number_of_jac_adj_muls: 686
         "###);
     }
@@ -736,8 +740,8 @@ mod test {
         test_adjoint(adjoint_solver, dgdu);
         insta::assert_yaml_snapshot!(problem.eqn.rhs().statistics(), @r###"
         number_of_calls: 433
-        number_of_jac_muls: 30
-        number_of_matrix_evals: 10
+        number_of_jac_muls: 27
+        number_of_matrix_evals: 9
         number_of_jac_adj_muls: 169
         "###);
     }