IDA adaptive time stepping (pybamm-team#4351)

* IDA adaptive time stepping * codecov and ubuntu runtime * update `IDAKLU` comments * better memory useage, codecov * codecov * fix ubuntu test * address comments * fix codecov, add tests * fix windows test * fix codecov * move integration test to unit fixes codecov * updates changes --------- Co-authored-by: Eric G. Kratz <[email protected]>
DrSOKane · Aug 23, 2024 · 4d391fa · 4d391fa
1 parent fcab586
commit 4d391fa
Show file tree

Hide file tree

Showing 32 changed files with 1,251 additions and 453 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@
 
 ## Optimizations
 
+- 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))
 

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -91,6 +91,7 @@ pybind11_add_module(idaklu
   src/pybamm/solvers/c_solvers/idaklu/IdakluJax.cpp
   src/pybamm/solvers/c_solvers/idaklu/IdakluJax.hpp
   src/pybamm/solvers/c_solvers/idaklu/common.hpp
+  src/pybamm/solvers/c_solvers/idaklu/common.cpp
   src/pybamm/solvers/c_solvers/idaklu/python.hpp
   src/pybamm/solvers/c_solvers/idaklu/python.cpp
   src/pybamm/solvers/c_solvers/idaklu/Solution.cpp

diff --git a/setup.py b/setup.py
@@ -322,6 +322,7 @@ def compile_KLU():
         "src/pybamm/solvers/c_solvers/idaklu/IdakluJax.cpp",
         "src/pybamm/solvers/c_solvers/idaklu/IdakluJax.hpp",
         "src/pybamm/solvers/c_solvers/idaklu/common.hpp",
+        "src/pybamm/solvers/c_solvers/idaklu/common.cpp",
         "src/pybamm/solvers/c_solvers/idaklu/python.hpp",
         "src/pybamm/solvers/c_solvers/idaklu/python.cpp",
         "src/pybamm/solvers/c_solvers/idaklu/Solution.cpp",

diff --git a/src/pybamm/batch_study.py b/src/pybamm/batch_study.py
@@ -106,6 +106,7 @@ def solve(
         calc_esoh=True,
         starting_solution=None,
         initial_soc=None,
+        t_interp=None,
         **kwargs,
     ):
         """

diff --git a/src/pybamm/experiment/experiment.py b/src/pybamm/experiment/experiment.py
@@ -40,7 +40,7 @@ class Experiment:
     def __init__(
         self,
         operating_conditions: list[str | tuple[str]],
-        period: str = "1 minute",
+        period: str | None = None,
         temperature: float | None = None,
         termination: list[str] | None = None,
     ):

diff --git a/src/pybamm/experiment/step/base_step.py b/src/pybamm/experiment/step/base_step.py
@@ -266,6 +266,84 @@ def default_duration(self, value):
         else:
             return 24 * 3600  # one day in seconds
 
+    @staticmethod
+    def default_period():
+        return 60.0  # seconds
+
+    def default_time_vector(self, tf, t0=0):
+        if self.period is None:
+            period = self.default_period()
+        else:
+            period = self.period
+        npts = max(int(round(np.abs(tf - t0) / period)) + 1, 2)
+
+        return np.linspace(t0, tf, npts)
+
+    def setup_timestepping(self, solver, tf, t_interp=None):
+        """
+        Setup timestepping for the model.
+
+        Parameters
+        ----------
+        solver: :class`pybamm.BaseSolver`
+            The solver
+        tf: float
+            The final time
+        t_interp: np.array | None
+            The time points at which to interpolate the solution
+        """
+        if solver.supports_interp:
+            return self._setup_timestepping(solver, tf, t_interp)
+        else:
+            return self._setup_timestepping_dense_t_eval(solver, tf, t_interp)
+
+    def _setup_timestepping(self, solver, tf, t_interp):
+        """
+        Setup timestepping for the model. This returns a t_eval vector that stops
+        only at the first and last time points. If t_interp and the period are
+        unspecified, then the solver will use adaptive time-stepping. For a given
+        period, t_interp will be set to return the solution at the end of each period
+        and at the final time.
+
+        Parameters
+        ----------
+        solver: :class`pybamm.BaseSolver`
+            The solver
+        tf: float
+            The final time
+        t_interp: np.array | None
+            The time points at which to interpolate the solution
+        """
+        t_eval = np.array([0, tf])
+        if t_interp is None:
+            if self.period is not None:
+                t_interp = self.default_time_vector(tf)
+            else:
+                t_interp = solver.process_t_interp(t_interp)
+
+        return t_eval, t_interp
+
+    def _setup_timestepping_dense_t_eval(self, solver, tf, t_interp):
+        """
+        Setup timestepping for the model. By default, this returns a dense t_eval which
+        stops the solver at each point in the t_eval vector. This method is for solvers
+        that do not support intra-solve interpolation for the solution.
+
+        Parameters
+        ----------
+        solver: :class`pybamm.BaseSolver`
+            The solver
+        tf: float
+            The final time
+        t_interp: np.array | None
+            The time points at which to interpolate the solution
+        """
+        t_eval = self.default_time_vector(tf)
+
+        t_interp = solver.process_t_interp(t_interp)
+
+        return t_eval, t_interp
+
     def process_model(self, model, parameter_values):
         new_model = model.new_copy()
         new_parameter_values = parameter_values.copy()

diff --git a/src/pybamm/simulation.py b/src/pybamm/simulation.py
@@ -352,6 +352,7 @@ def solve(
         callbacks=None,
         showprogress=False,
         inputs=None,
+        t_interp=None,
         **kwargs,
     ):
         """
@@ -361,11 +362,14 @@ def solve(
         Parameters
         ----------
         t_eval : numeric type, optional
-            The times (in seconds) at which to compute the solution. Can be
-            provided as an array of times at which to return the solution, or as a
-            list `[t0, tf]` where `t0` is the initial time and `tf` is the final time.
-            If provided as a list the solution is returned at 100 points within the
-            interval `[t0, tf]`.
+            The times at which to stop the integration due to a discontinuity in time.
+            Can be provided as an array of times at which to return the solution, or as
+            a list `[t0, tf]` where `t0` is the initial time and `tf` is the final
+            time. If the solver does not support intra-solve interpolation, providing
+            `t_eval` as a list returns the solution at 100 points within the interval
+            `[t0, tf]`. Otherwise, the solution is returned at the times specified in
+            `t_interp` or as a result of the adaptive time-stepping solution. See the
+            `t_interp` argument for more details.
 
             If not using an experiment or running a drive cycle simulation (current
             provided as data) `t_eval` *must* be provided.
@@ -400,6 +404,9 @@ def solve(
             Whether to show a progress bar for cycling. If true, shows a progress bar
             for cycles. Has no effect when not used with an experiment.
             Default is False.
+        t_interp : None, list or ndarray, optional
+            The times (in seconds) at which to interpolate the solution. Defaults to None.
+            Only valid for solvers that support intra-solve interpolation (`IDAKLUSolver`).
         **kwargs
             Additional key-word arguments passed to `solver.solve`.
             See :meth:`pybamm.BaseSolver.solve`.
@@ -486,7 +493,7 @@ def solve(
                         )
 
             self._solution = solver.solve(
-                self._built_model, t_eval, inputs=inputs, **kwargs
+                self._built_model, t_eval, inputs=inputs, t_interp=t_interp, **kwargs
             )
 
         elif self.operating_mode == "with experiment":
@@ -687,13 +694,17 @@ def solve(
                         "start time": start_time,
                     }
                     # Make sure we take at least 2 timesteps
-                    npts = max(int(round(dt / step.period)) + 1, 2)
+                    t_eval, t_interp_processed = step.setup_timestepping(
+                        solver, dt, t_interp
+                    )
+
                     try:
                         step_solution = solver.step(
                             current_solution,
                             model,
                             dt,
-                            t_eval=np.linspace(0, dt, npts),
+                            t_eval,
+                            t_interp=t_interp_processed,
                             save=False,
                             inputs=inputs,
                             **kwargs,

diff --git a/src/pybamm/solvers/algebraic_solver.py b/src/pybamm/solvers/algebraic_solver.py
@@ -36,7 +36,7 @@ def __init__(self, method="lm", tol=1e-6, extra_options=None):
         self.tol = tol
         self.extra_options = extra_options or {}
         self.name = f"Algebraic solver ({method})"
-        self.algebraic_solver = True
+        self._algebraic_solver = True
         pybamm.citations.register("Virtanen2020")
 
     @property
@@ -47,7 +47,7 @@ def tol(self):
     def tol(self, value):
         self._tol = value
 
-    def _integrate(self, model, t_eval, inputs_dict=None):
+    def _integrate(self, model, t_eval, inputs_dict=None, t_interp=None):
         """
         Calculate the solution of the algebraic equations through root-finding