Fix bug with out-of-bounds discontinuities in time (#5205)

* filter constant-time discontinuities

Co-Authored-By: isaacbasil <[email protected]>

* Update CHANGELOG.md

---------

Co-authored-by: isaacbasil <[email protected]>

Author: 2 people

CHANGELOG.md

@@ -10,6 +10,7 @@
 
 ## Bug fixes
 
+- Fixed a bug where time-based Heaviside or modulo discontinuities could trigger out-of-bounds errors in time arrays. ([#5205](https://github.com/pybamm-team/PyBaMM/pull/5205))
 - Fixed a bug using a time-varying input with heaviside or modulo functions using the `IDAKLUSolver`. ([#4994](https://github.com/pybamm-team/PyBaMM/pull/4994))
 - Fix a bug in setting initial stoichiometries where the reference temperature was used instead of the initial temperature. ([#5189](https://github.com/pybamm-team/PyBaMM/pull/5189))
 - Fix a bug in the calculation of "Bulk" OCP terms in hysteresis models ([#5169](https://github.com/pybamm-team/PyBaMM/pull/5169))

src/pybamm/solvers/base_solver.py

@@ -224,6 +224,7 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
             casadi_switch_events,
             terminate_events,
             interpolant_extrapolation_events,
+            t_discon_constant,
             discontinuity_events,
         ) = self._set_up_events(model, t_eval, inputs, vars_for_processing)
 
@@ -233,8 +234,9 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
         model.rhs_algebraic_eval = rhs_algebraic
 
         model.terminate_events_eval = terminate_events
-        model.discontinuity_events_eval = discontinuity_events
         model.interpolant_extrapolation_events_eval = interpolant_extrapolation_events
+        model.discontinuity_events_eval = discontinuity_events
+        model.t_discon_constant = t_discon_constant
 
         model.jac_rhs_eval = jac_rhs
         model.jac_rhs_action_eval = jac_rhs_action
@@ -482,23 +484,17 @@ def _set_up_events(self, model, t_eval, inputs, vars_for_processing):
         # discontinuity events if these exist.
         # Note: only checks for the case of t < X, t <= X, X < t, or X <= t,
         # but also accounts for the fact that t might be dimensional
-
-        t0 = np.min(t_eval)
         tf = np.max(t_eval)
 
         def supports_t_eval_discontinuities(expr):
             # Only IDAKLUSolver supports discontinuities represented by t_eval
-            return (
-                self.supports_t_eval_discontinuities
-                and (t_eval is not None)
-                and expr.is_constant()
-            )
+            return self.supports_t_eval_discontinuities and expr.is_constant()
+
+        # Find all the constant time-based discontinuities
+        t_discon = []
 
-        def append_t_eval(t):
-            if t0 <= t <= tf and t not in t_eval:
-                # Insert t in the correct position to maintain sorted order
-                idx = np.searchsorted(t_eval, t)
-                t_eval.insert(idx, t)
+        def append_t_discon(t):
+            t_discon.append(t)
 
         def heaviside_event(symbol, expr):
             model.events.append(
@@ -509,28 +505,28 @@ def heaviside_event(symbol, expr):
                 )
             )
 
-        def heaviside_t_eval(symbol, expr):
+        def heaviside_t_discon(symbol, expr):
             value = expr.evaluate(0, model.y0.full(), inputs=inputs)
-            append_t_eval(value)
+            append_t_discon(value)
 
             if isinstance(symbol, pybamm.EqualHeaviside):
                 if symbol.left == pybamm.t:
                     # t <= x
                     # Stop at t = x and right after t = x
-                    append_t_eval(np.nextafter(value, np.inf))
+                    append_t_discon(np.nextafter(value, np.inf))
                 else:
                     # t >= x
                     # Stop at t = x and right before t = x
-                    append_t_eval(np.nextafter(value, -np.inf))
+                    append_t_discon(np.nextafter(value, -np.inf))
             elif isinstance(symbol, pybamm.NotEqualHeaviside):
                 if symbol.left == pybamm.t:
                     # t < x
                     # Stop at t = x and right before t = x
-                    append_t_eval(np.nextafter(value, -np.inf))
+                    append_t_discon(np.nextafter(value, -np.inf))
                 else:
                     # t > x
                     # Stop at t = x and right after t = x
-                    append_t_eval(np.nextafter(value, np.inf))
+                    append_t_discon(np.nextafter(value, np.inf))
             else:
                 raise ValueError(
                     f"Unknown heaviside function: {symbol}"
@@ -546,13 +542,13 @@ def modulo_event(symbol, expr, num_events):
                     )
                 )
 
-        def modulo_t_eval(symbol, expr, num_events):
+        def modulo_t_discon(symbol, expr, num_events):
             value = expr.evaluate(0, model.y0.full(), inputs=inputs)
             for i in np.arange(num_events):
                 t = value * (i + 1)
                 # Stop right before t and at t
-                append_t_eval(np.nextafter(t, -np.inf))
-                append_t_eval(t)
+                append_t_discon(np.nextafter(t, -np.inf))
+                append_t_discon(t)
 
         for symbol in itertools.chain(
             model.concatenated_rhs.pre_order(),
@@ -569,7 +565,7 @@ def modulo_t_eval(symbol, expr, num_events):
                     continue  # pragma: no cover
 
                 if supports_t_eval_discontinuities(expr):
-                    heaviside_t_eval(symbol, expr)
+                    heaviside_t_discon(symbol, expr)
                 else:
                     heaviside_event(symbol, expr)
 
@@ -578,7 +574,7 @@ def modulo_t_eval(symbol, expr, num_events):
                 num_events = 200 if (t_eval is None) else (tf // expr.value)
 
                 if supports_t_eval_discontinuities(expr):
-                    modulo_t_eval(symbol, expr, num_events)
+                    modulo_t_discon(symbol, expr, num_events)
                 else:
                     modulo_event(symbol, expr, num_events)
             else:
@@ -641,6 +637,7 @@ def modulo_t_eval(symbol, expr, num_events):
             casadi_switch_events,
             terminate_events,
             interpolant_extrapolation_events,
+            t_discon,
             discontinuity_events,
         )
 
@@ -1053,39 +1050,53 @@ def solve(
             return solutions
 
     @staticmethod
-    def _get_discontinuity_start_end_indices(model, inputs, t_eval):
+    def filter_discontinuities(t_discon: list, t_eval: list) -> np.ndarray:
+        """
+        Filter the discontinuities to only include the unique and sorted
+        values within the t_eval range (non-exclusive of end points).
+
+        Parameters
+        ----------
+        t_discon : list
+            The list of all possible discontinuity times.
+        t_eval : list
+            The integration time points.
+
+        Returns
+        -------
+        np.ndarray
+            The filtered list of discontinuities within the range of t_eval.
+        """
+        t_discon_unique = np.unique(t_discon)
+
+        # Find the indices within t_eval (non-exclusive of end points)
+        idx_start = np.searchsorted(t_discon_unique, t_eval[0], side="right")
+        idx_end = np.searchsorted(t_discon_unique, t_eval[-1], side="left")
+        return t_discon_unique[idx_start:idx_end]
+
+    def _get_discontinuity_start_end_indices(self, model, inputs, t_eval):
+        if self.supports_t_eval_discontinuities:
+            t_discon_constant = self.filter_discontinuities(
+                model.t_discon_constant, t_eval
+            )
+            t_eval = np.union1d(t_eval, t_discon_constant)
+
         if not model.discontinuity_events_eval:
             pybamm.logger.verbose("No discontinuity events found")
             return [0], [len(t_eval)], t_eval
 
-        # Calculate discontinuities
-        discontinuities = [
+        # Calculate all possible discontinuities
+        _t_discon_full = [
             # Assuming that discontinuities do not depend on
             # input parameters when len(input_list) > 1, only
             # `inputs` is passed to `evaluate`.
             # See https://github.com/pybamm-team/PyBaMM/pull/1261
             event.expression.evaluate(inputs=inputs)
             for event in model.discontinuity_events_eval
         ]
+        t_discon = self.filter_discontinuities(_t_discon_full, t_eval)
 
-        # make sure they are increasing in time
-        discontinuities = sorted(discontinuities)
-
-        # remove any identical discontinuities
-        discontinuities = [
-            v
-            for i, v in enumerate(discontinuities)
-            if (
-                i == len(discontinuities) - 1
-                or discontinuities[i] < discontinuities[i + 1]
-            )
-            and v > 0
-        ]
-
-        # remove any discontinuities after end of t_eval
-        discontinuities = [v for v in discontinuities if v < t_eval[-1]]
-
-        pybamm.logger.verbose(f"Discontinuity events found at t = {discontinuities}")
+        pybamm.logger.verbose(f"Discontinuity events found at t = {t_discon}")
         if isinstance(inputs, list):
             raise pybamm.SolverError(
                 "Cannot solve for a list of input parameters sets with discontinuities"
@@ -1096,7 +1107,7 @@ def _get_discontinuity_start_end_indices(model, inputs, t_eval):
         start_indices = [0]
         end_indices = []
         eps = sys.float_info.epsilon
-        for dtime in discontinuities:
+        for dtime in t_discon:
             dindex = np.searchsorted(t_eval, dtime, side="left")
             end_indices.append(dindex + 1)
             start_indices.append(dindex + 1)

tests/unit/test_simulation.py

@@ -6,6 +6,7 @@
 from scipy.integrate import trapezoid
 
 import pybamm
+from pybamm.solvers.base_solver import BaseSolver
 from tests import no_internet_connection
 
 
@@ -895,6 +896,19 @@ def prevfloat(t):
         for t_node in t_nodes:
             assert current(t_node) == pytest.approx(sawtooth_current(t_node))
 
+    def test_filter_discontinuities_simple(self):
+        t_eval = [0.0, 3.0, 10.0]
+        t_discon = [-5.0, 0.0, 1.0, 3.0, 3.0, 5.0, 10.0, 12.0]
+
+        result = BaseSolver.filter_discontinuities(t_discon, t_eval)
+        expected = np.array([1.0, 3.0, 5.0])
+
+        # Exclusive of endpoints
+        t_eval_endpoints = [t_eval[0], t_eval[-1]]
+        assert all(t not in result for t in t_eval_endpoints)
+
+        np.testing.assert_array_equal(result, expected)
+
     def test_t_eval(self):
         model = pybamm.lithium_ion.SPM()
         sim = pybamm.Simulation(model)