Set closest_event_idx in IDAKLUSolver after IDA_ROOT_RETURN#5502
Merged
rtimms merged 6 commits intopybamm-team:mainfrom May 8, 2026
Merged
Set closest_event_idx in IDAKLUSolver after IDA_ROOT_RETURN#5502rtimms merged 6 commits intopybamm-team:mainfrom
rtimms merged 6 commits intopybamm-team:mainfrom
Conversation
After a root return, BaseSolver.get_termination_reason needs to know which TERMINATION event fired so it can label solution.termination. The casadi solver records this in solution.closest_event_idx and get_termination_reason short-circuits. IDAKLU never set it, so the slow fallback path re-walked every event's symbolic expression on the Python side every step — generating tens of thousands of small numpy allocations from StateVector._base_evaluate and scipy._matmul_vector on long event-terminated cycling runs. The combined-events root function (rootfn) is already built and serialised for the C++ wrapper. Keep a Python-callable handle on it in self._setup so _post_process_solution can evaluate it once at (t_event, y_event, p_stacked) and set closest_event_idx via np.nanargmin(|values|). Round-trip through __getstate__/__setstate__ is supported by deserialising from the same rootfn_pkl the C++ wrapper consumes. Also extracts the duplicated inputs-dict-flattening pattern in _integrate into a small _flatten_inputs helper, reused at both call sites. Benchmark: 1000-cycle SPM with output_variables set, IDAKLUSolver, save_at_cycles=N — baseline 1756 MB / 4.43M allocs / 13.10 s wall; with fix 1311 MB / 3.98M allocs / 10.95 s wall (25% bytes, 10% allocs, 16% wall). StateVector._base_evaluate count drops 326k → 0; _matmul_vector 305k → 0. Adds tests/unit/test_solvers/test_idaklu_solver.py:: test_closest_event_idx_set_after_root_return — fails on main (closest_event_idx is None for IDAKLU), passes after this change, and verifies the index resolves to the same event the slow path would have picked.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #5502 +/- ##
==========================================
- Coverage 98.17% 98.17% -0.01%
==========================================
Files 338 338
Lines 31160 31222 +62
==========================================
+ Hits 30591 30651 +60
- Misses 569 571 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
The behavioural test in test_idaklu_solver.py asserts the contract; this adds a count-based memory test alongside the existing test_memory_leaks.py to catch the regression at its measurable symptom. Counts direct calls to StateVector._base_evaluate during a 5-cycle solve — 1266 without the fix, 161 with — separated by a threshold of 300. Comments trimmed in idaklu_solver.py: the long WHY paragraphs were too prescriptive next to short code; the changelog and PR description already cover the rationale.
6 tasks
Adds tests/unit/test_solvers/test_idaklu_solver.py:: test_pickle_roundtrip_preserves_closest_event_idx which exercises the __getstate__/__setstate__ path: serialise an IDAKLUSolver, deserialise, run a fresh solve, and confirm closest_event_idx is still populated for event-terminated steps. This covers the new self._setup["rootfn_casadi"] entry and its rehydration via casadi.Function.deserialize from the existing rootfn_pkl.
BradyPlanden
approved these changes
May 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
After a SUNDIALS root return,
BaseSolver.get_termination_reasonneeds to labelsolution.terminationwith the specific event that fired. The casadi solver records this inSolution.closest_event_idxandget_termination_reasonshort-circuits. IDAKLU never set it, so the slow fallback path re-walked every TERMINATION event's symbolic expression on the Python side every step — dominating per-step allocations on long event-terminated cycling runs.The combined-events root function (
rootfn) is already built and serialised for the C++ wrapper. This PR keeps a Python-callable handle on it inself._setup["rootfn_casadi"]so_post_process_solutioncan evaluate it once at(t_event, y_event, p_stacked)and setclosest_event_idx = nanargmin(|values|). Round-trip through__getstate__/__setstate__is supported by deserialising from the samerootfn_pklthe C++ wrapper consumes.Also extracts the duplicated inputs-dict-flattening pattern in
_integrateinto a small_flatten_inputshelper, reused at both call sites.Where the allocations were
memray on a 1000-cycle SPM (
output_variablesset, IDAKLUSolver,save_at_cycles=Nso most cycles are discarded):StateVector._base_evaluate(state_vector.py:299)_matmul_vector(scipy.sparse._compressed.py:395)Both eliminated. Top-line:
Test plan
test_closest_event_idx_set_after_root_returnintests/unit/test_solvers/test_idaklu_solver.py— asserts every event-terminated step hasclosest_event_idxpopulated AND that the index resolves to the same event namestep.terminationcarries. Fails onmain(alwaysNonefor IDAKLU), passes after this change.tests/memory/test_memory_leaks.pytests pass.