Merge pull request #4440 from pipliggins/output_vars_extrapolation

Fix IDAKLU output_variables to work with extrapolation events

CHANGELOG.md

@@ -4,16 +4,17 @@
 
 - Added sensitivity calculation support for `pybamm.Simulation` and `pybamm.Experiment` ([#4415](https://github.com/pybamm-team/PyBaMM/pull/4415))
 - Added OpenMP parallelization to IDAKLU solver for lists of input parameters ([#4449](https://github.com/pybamm-team/PyBaMM/pull/4449))
+- Added phase-dependent particle options to LAM #4369
 - Added a lithium ion equivalent circuit model with split open circuit voltages for each electrode (`SplitOCVR`). ([#4330](https://github.com/pybamm-team/PyBaMM/pull/4330))
 
 ## Optimizations
 
 - Improved performance of initialization and reinitialization of ODEs in the (`IDAKLUSolver`). ([#4453](https://github.com/pybamm-team/PyBaMM/pull/4453))
 - Removed the `start_step_offset` setting and disabled minimum `dt` warnings for drive cycles with the (`IDAKLUSolver`). ([#4416](https://github.com/pybamm-team/PyBaMM/pull/4416))
 
-## Features
+## Bug Fixes
 
-- Added phase-dependent particle options to LAM #4369
+- Fixed bug where IDAKLU solver failed when `output variables` were specified and an extrapolation event is present. ([#4440](https://github.com/pybamm-team/PyBaMM/pull/4440))
 
 ## Breaking changes
 

src/pybamm/solvers/base_solver.py

@@ -1489,8 +1489,12 @@ def check_extrapolation(self, solution, events):
 
         # second pass: check if the extrapolation events are within the tolerance
         last_state = solution.last_state
-        t = last_state.all_ts[0][0]
-        y = last_state.all_ys[0][:, 0]
+        if solution.t_event:
+            t = solution.t_event[0]
+            y = solution.y_event[:, 0]
+        else:
+            t = last_state.all_ts[0][0]
+            y = last_state.all_ys[0][:, 0]
         inputs = last_state.all_inputs[0]
 
         if isinstance(y, casadi.DM):

tests/unit/test_solvers/test_idaklu_solver.py

@@ -1162,3 +1162,24 @@ def test_python_idaklu_deprecation_errors(self):
                     match="Unsupported evaluation engine for convert_to_format=jax",
                 ):
                     _ = solver.solve(model, t_eval)
+
+    def test_extrapolation_events_with_output_variables(self):
+        # Make sure the extrapolation checks work with output variables
+        model = pybamm.BaseModel()
+        v = pybamm.Variable("v")
+        c = pybamm.Variable("c")
+        model.variables = {"v": v, "c": c}
+        model.rhs = {v: -1, c: 0}
+        model.initial_conditions = {v: 1, c: 2}
+        model.events.append(
+            pybamm.Event(
+                "Triggered event",
+                v - 0.5,
+                pybamm.EventType.INTERPOLANT_EXTRAPOLATION,
+            )
+        )
+        solver = pybamm.IDAKLUSolver(output_variables=["c"])
+        solver.set_up(model)
+
+        with pytest.warns(pybamm.SolverWarning, match="extrapolation occurred for"):
+            solver.solve(model, t_eval=[0, 1])