Update IDAKLU time stepping in tests and benchmarks (pybamm-team#4390)

* time stepping updates

* Update test_idaklu_solver.py

* Update CHANGELOG.md

CHANGELOG.md

@@ -7,13 +7,14 @@
 
 ## Optimizations
 
+- Update `IDAKLU` tests and benchmarks to use adaptive time stepping. ([#4390](https://github.com/pybamm-team/PyBaMM/pull/4390))
 - Improved adaptive time-stepping performance of the (`IDAKLUSolver`). ([#4351](https://github.com/pybamm-team/PyBaMM/pull/4351))
 - Improved performance and reliability of DAE consistent initialization. ([#4301](https://github.com/pybamm-team/PyBaMM/pull/4301))
 - Replaced rounded Faraday constant with its exact value in `bpx.py` for better comparison between different tools. ([#4290](https://github.com/pybamm-team/PyBaMM/pull/4290))
 
 ## Bug Fixes
 
-- Fixed memory issue that caused failure when `output variables` were specified with (`IDAKLUSolver`). ([#4379](https://github.com/pybamm-team/PyBaMM/issues/4379))
+- Fixed memory issue that caused failure when `output variables` were specified with (`IDAKLUSolver`). ([#4379](https://github.com/pybamm-team/PyBaMM/pull/4379))
 - Fixed bug where IDAKLU solver failed when `output variables` were specified and an event triggered. ([#4300](https://github.com/pybamm-team/PyBaMM/pull/4300))
 
 ## Breaking changes

benchmarks/different_model_options.py

@@ -34,6 +34,7 @@ class SolveModel:
     solver: pybamm.BaseSolver
     model: pybamm.BaseModel
     t_eval: np.ndarray
+    t_interp: np.ndarray | None
 
     def solve_setup(self, parameter, model_, option, value, solver_class):
         import importlib
@@ -51,8 +52,13 @@ def solve_setup(self, parameter, model_, option, value, solver_class):
         self.model = model_({option: value})
         c_rate = 1
         tmax = 4000 / c_rate
-        nb_points = 500
-        self.t_eval = np.linspace(0, tmax, nb_points)
+        if self.solver.supports_interp:
+            self.t_eval = np.array([0, tmax])
+            self.t_interp = None
+        else:
+            nb_points = 500
+            self.t_eval = np.linspace(0, tmax, nb_points)
+            self.t_interp = None
         geometry = self.model.default_geometry
 
         # load parameter values and process model and geometry
@@ -77,7 +83,7 @@ def solve_setup(self, parameter, model_, option, value, solver_class):
         disc.process_model(self.model)
 
     def solve_model(self, _model, _params):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeBuildModelLossActiveMaterial:
@@ -109,7 +115,7 @@ def setup(self, model, params, solver_class):
         )
 
     def time_solve_model(self, _model, _params, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeBuildModelLithiumPlating:
@@ -141,7 +147,7 @@ def setup(self, model, params, solver_class):
         )
 
     def time_solve_model(self, _model, _params, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeBuildModelSEI:
@@ -187,7 +193,7 @@ def setup(self, model, params, solver_class):
         SolveModel.solve_setup(self, "Marquis2019", model, "SEI", params, solver_class)
 
     def time_solve_model(self, _model, _params, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeBuildModelParticle:
@@ -229,7 +235,7 @@ def setup(self, model, params, solver_class):
         )
 
     def time_solve_model(self, _model, _params, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeBuildModelThermal:
@@ -261,7 +267,7 @@ def setup(self, model, params, solver_class):
         )
 
     def time_solve_model(self, _model, _params, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeBuildModelSurfaceForm:
@@ -299,4 +305,4 @@ def setup(self, model, params, solver_class):
         )
 
     def time_solve_model(self, _model, _params, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)

benchmarks/time_solve_models.py

@@ -6,8 +6,8 @@
 import numpy as np
 
 
-def solve_model_once(model, solver, t_eval):
-    solver.solve(model, t_eval=t_eval)
+def solve_model_once(model, solver, t_eval, t_interp):
+    solver.solve(model, t_eval=t_eval, t_interp=t_interp)
 
 
 class TimeSolveSPM:
@@ -31,15 +31,22 @@ class TimeSolveSPM:
     model: pybamm.BaseModel
     solver: pybamm.BaseSolver
     t_eval: np.ndarray
+    t_interp: np.ndarray | None
 
     def setup(self, solve_first, parameters, solver_class):
         set_random_seed()
         self.solver = solver_class()
         self.model = pybamm.lithium_ion.SPM()
         c_rate = 1
         tmax = 4000 / c_rate
-        nb_points = 500
-        self.t_eval = np.linspace(0, tmax, nb_points)
+        if self.solver.supports_interp:
+            self.t_eval = np.array([0, tmax])
+            self.t_interp = None
+        else:
+            nb_points = 500
+            self.t_eval = np.linspace(0, tmax, nb_points)
+            self.t_interp = None
+
         geometry = self.model.default_geometry
 
         # load parameter values and process model and geometry
@@ -63,10 +70,10 @@ def setup(self, solve_first, parameters, solver_class):
         disc = pybamm.Discretisation(mesh, self.model.default_spatial_methods)
         disc.process_model(self.model)
         if solve_first:
-            solve_model_once(self.model, self.solver, self.t_eval)
+            solve_model_once(self.model, self.solver, self.t_eval, self.t_interp)
 
     def time_solve_model(self, _solve_first, _parameters, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeSolveSPMe:
@@ -97,8 +104,13 @@ def setup(self, solve_first, parameters, solver_class):
         self.model = pybamm.lithium_ion.SPMe()
         c_rate = 1
         tmax = 4000 / c_rate
-        nb_points = 500
-        self.t_eval = np.linspace(0, tmax, nb_points)
+        if self.solver.supports_interp:
+            self.t_eval = np.array([0, tmax])
+            self.t_interp = None
+        else:
+            nb_points = 500
+            self.t_eval = np.linspace(0, tmax, nb_points)
+            self.t_interp = None
         geometry = self.model.default_geometry
 
         # load parameter values and process model and geometry
@@ -122,10 +134,10 @@ def setup(self, solve_first, parameters, solver_class):
         disc = pybamm.Discretisation(mesh, self.model.default_spatial_methods)
         disc.process_model(self.model)
         if solve_first:
-            solve_model_once(self.model, self.solver, self.t_eval)
+            solve_model_once(self.model, self.solver, self.t_eval, self.t_interp)
 
     def time_solve_model(self, _solve_first, _parameters, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)
 
 
 class TimeSolveDFN:
@@ -161,8 +173,13 @@ def setup(self, solve_first, parameters, solver_class):
         self.model = pybamm.lithium_ion.DFN()
         c_rate = 1
         tmax = 4000 / c_rate
-        nb_points = 500
-        self.t_eval = np.linspace(0, tmax, nb_points)
+        if self.solver.supports_interp:
+            self.t_eval = np.array([0, tmax])
+            self.t_interp = None
+        else:
+            nb_points = 500
+            self.t_eval = np.linspace(0, tmax, nb_points)
+            self.t_interp = None
         geometry = self.model.default_geometry
 
         # load parameter values and process model and geometry
@@ -186,7 +203,7 @@ def setup(self, solve_first, parameters, solver_class):
         disc = pybamm.Discretisation(mesh, self.model.default_spatial_methods)
         disc.process_model(self.model)
         if solve_first:
-            solve_model_once(self.model, self.solver, self.t_eval)
+            solve_model_once(self.model, self.solver, self.t_eval, self.t_interp)
 
     def time_solve_model(self, _solve_first, _parameters, _solver_class):
-        self.solver.solve(self.model, t_eval=self.t_eval)
+        self.solver.solve(self.model, t_eval=self.t_eval, t_interp=self.t_interp)

benchmarks/work_precision_sets/time_vs_abstols.py

@@ -42,8 +42,13 @@
         model = i[1].new_copy()
         c_rate = 1
         tmax = 3500 / c_rate
-        nb_points = 500
-        t_eval = np.linspace(0, tmax, nb_points)
+        if solver.supports_interp:
+            t_eval = np.array([0, tmax])
+            t_interp = None
+        else:
+            nb_points = 500
+            t_eval = np.linspace(0, tmax, nb_points)
+            t_interp = None
         geometry = model.default_geometry
 
         # load parameter values and process model and geometry
@@ -69,11 +74,11 @@
 
         for tol in abstols:
             solver.atol = tol
-            solver.solve(model, t_eval=t_eval)
+            solver.solve(model, t_eval=t_eval, t_interp=t_interp)
             time = 0
             runs = 20
             for _ in range(0, runs):
-                solution = solver.solve(model, t_eval=t_eval)
+                solution = solver.solve(model, t_eval=t_eval, t_interp=t_interp)
                 time += solution.solve_time.value
             time = time / runs
 

benchmarks/work_precision_sets/time_vs_mesh_size.py

@@ -17,6 +17,7 @@
 npts = [4, 8, 16, 32, 64]
 
 solvers = {
+    "IDAKLUSolver": pybamm.IDAKLUSolver(),
     "Casadi - safe": pybamm.CasadiSolver(),
     "Casadi - fast": pybamm.CasadiSolver(mode="fast"),
 }

benchmarks/work_precision_sets/time_vs_no_of_states.py

@@ -16,6 +16,7 @@
 npts = [4, 8, 16, 32, 64]
 
 solvers = {
+    "IDAKLUSolver": pybamm.IDAKLUSolver(),
     "Casadi - safe": pybamm.CasadiSolver(),
     "Casadi - fast": pybamm.CasadiSolver(mode="fast"),
 }

benchmarks/work_precision_sets/time_vs_reltols.py

@@ -48,8 +48,13 @@
         model = i[1].new_copy()
         c_rate = 1
         tmax = 3500 / c_rate
-        nb_points = 500
-        t_eval = np.linspace(0, tmax, nb_points)
+        if solver.supports_interp:
+            t_eval = np.array([0, tmax])
+            t_interp = None
+        else:
+            nb_points = 500
+            t_eval = np.linspace(0, tmax, nb_points)
+            t_interp = None
         geometry = model.default_geometry
 
         # load parameter values and process model and geometry
@@ -75,11 +80,11 @@
 
         for tol in reltols:
             solver.rtol = tol
-            solver.solve(model, t_eval=t_eval)
+            solver.solve(model, t_eval=t_eval, t_interp=t_interp)
             time = 0
             runs = 20
             for _ in range(0, runs):
-                solution = solver.solve(model, t_eval=t_eval)
+                solution = solver.solve(model, t_eval=t_eval, t_interp=t_interp)
                 time += solution.solve_time.value
             time = time / runs
 

tests/integration/test_solvers/test_idaklu.py

@@ -32,7 +32,7 @@ def test_on_spme_sensitivities(self):
         disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
         disc.process_model(model)
         t_interp = np.linspace(0, 3500, 100)
-        t_eval = np.array([t_interp[0], t_interp[-1]])
+        t_eval = [t_interp[0], t_interp[-1]]
         solver = pybamm.IDAKLUSolver(rtol=1e-10, atol=1e-10)
         solution = solver.solve(
             model,