Fix wrong branch for piecewise nested in another piecewise's condition - #3234
Open
FFroehlich wants to merge 3 commits into
Open
Fix wrong branch for piecewise nested in another piecewise's condition#3234FFroehlich wants to merge 3 commits into
FFroehlich wants to merge 3 commits into
Conversation
`Model::reinit_events` computed the Heaviside variables `h` in a single pass: it evaluated `froot` once and then derived `h` from the root function values. Root functions are evaluated from `w`, though, and `w` may itself depend on `h` -- which is exactly what happens when a piecewise expression occurs inside the condition of another piecewise expression. In that situation the outer root function was evaluated using the Heaviside variables of the inner expression as they were *before* the same pass updated them (i.e., all zeros at `t0`), so the outer Heaviside variable was initialized from a stale root value and the wrong branch was taken for the rest of the simulation. For the model from #3233 (`w = [p1, p2 = h0*p1, p3 = p2 - 1e-4, p4 = h2*p3]`, `root = [p1, -p1, p3, -p3]`), `root[2]` was evaluated as `-1e-4` with `h0` still `0`, so `h2` became `0` and `p4` stayed `0` although `p3` evaluated to `9e-4` afterwards. `h` is now iterated to a fixed point (bounded by `ne + 1` iterations, so cyclic dependencies cannot loop forever), and `roots_found` is derived from the converged values. As a side effect, `reinit_explicit_roots` now also sees a consistent `w`. Models whose root functions don't depend on `h` -- the common case -- just do one additional `froot` evaluation confirming that nothing changed. Closes #3233 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016ZKvTDruCaD3vHVXXXjz8o
`_handle_t0_event` evaluated the root functions exactly once, at the `h` it
seeded from `event_initial_values`. Root functions are evaluated from the
model expressions, which may themselves depend on `h` -- so for a piecewise
expression inside the condition of another piecewise expression, the outer
root function saw a stale value for the inner one, exactly as in
`Model::reinit_events`.
The only difference between the backends was the seed: JAX seeds `h = 1`
from the trigger initial values, sundials started from a zero-initialized
buffer. Both are arbitrary, so each was wrong on a different class of model.
For
p1 := -0.001
p2 := piecewise(0, p1 < 0, p1) # -> p1*H(p1); consistent H(p1) = 0
p3 := p2 + 0.0001 # correct: +1e-4
p4 := piecewise(0, p3 < 0, p3) # correct: +1e-4
a stale `h(p1) = 1` gives `p3 = -9e-4`, so `h(p3) = 0` and `p4` collapsed to
zero under JAX, while the model from #3233 was wrong under sundials instead.
`h` is now refined to a fixed point before the root functions are evaluated.
The trip count is static (`jax.lax.scan` over `n_events`, as already used for
cascading event assignments) because this sits in the AD path and
`jax.lax.while_loop` is not reverse-mode differentiable.
`h` itself keeps its seeded, pre-event value -- it is the reference that
`roots_found` flips against -- so only the point at which the root functions
are evaluated changes. Models whose root functions do not depend on `h`
therefore get a bit-identical result.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016ZKvTDruCaD3vHVXXXjz8o
Covers the mirror-image case of #3233, where the Heaviside variable of the inner piecewise expression is consistently 0 rather than 1 -- the case the JAX seed (`h = 1` from the trigger initial values) got wrong. Fails without the preceding commit. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016ZKvTDruCaD3vHVXXXjz8o
Contributor
There was a problem hiding this comment.
Pull request overview
Fixes incorrect branch selection for nested piecewise conditions by ensuring Heaviside/event states are initialized from root functions consistently when root functions depend on h (e.g., inner piecewise influencing an outer piecewise’s condition). This aligns both the Sundials (C++) and JAX backends and adds regression coverage for the reported failure mode (#3233).
Changes:
- Iterate Heaviside variables to a fixed point during event initialization in the C++ backend (
Model::reinit_events). - Refine the
hused for root evaluation att0in the JAX backend before determiningroots_found. - Add regression tests for both backends and document the fix in the changelog.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
src/model.cpp |
Iterates h until consistent with root function values before computing roots_found. |
python/sdist/amici/sim/jax/model.py |
Refines the h used for root evaluation at t0 to avoid stale inner-piecewise values. |
python/tests/test_jax.py |
Adds a JAX regression test reproducing nested-piecewise-in-condition behavior. |
python/tests/test_heavisides.py |
Adds a Sundials regression test including deeper nesting (3 levels). |
CHANGELOG.md |
Documents the fixed-point iteration approach and the user-visible behavioral fix. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+435
to
+449
| for (int iter = 0; iter <= ne; ++iter) { | ||
| froot(t, x, dx, rootvals); | ||
| bool h_changed = false; | ||
| for (int ie = 0; ie < ne; ie++) { | ||
| realtype const h_new = rootvals.at(ie) < 0.0 ? 0.0 : 1.0; | ||
| if (h_new != state_.h.at(ie)) { | ||
| state_.h.at(ie) = h_new; | ||
| h_changed = true; | ||
| } | ||
| } | ||
| if (!h_changed) { | ||
| break; | ||
| } | ||
| } | ||
|
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3234 +/- ##
==========================================
+ Coverage 78.61% 78.63% +0.01%
==========================================
Files 318 318
Lines 21122 21136 +14
Branches 1487 1490 +3
==========================================
+ Hits 16606 16620 +14
Misses 4508 4508
Partials 8 8
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
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
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.
Model::reinit_eventscomputed the Heaviside variableshin a singlepass: it evaluated
frootonce and then derivedhfrom the rootfunction values. Root functions are evaluated from
w, though, andwmay itself depend on
h-- which is exactly what happens when apiecewise expression occurs inside the condition of another piecewise
expression.
In that situation the outer root function was evaluated using the
Heaviside variables of the inner expression as they were before the
same pass updated them (i.e., all zeros at
t0), so the outer Heavisidevariable was initialized from a stale root value and the wrong branch was
taken for the rest of the simulation.
For the model from #3233 (
w = [p1, p2 = h0*p1, p3 = p2 - 1e-4, p4 = h2*p3],root = [p1, -p1, p3, -p3]),root[2]was evaluated as-1e-4withh0still0, soh2became0andp4stayed0although
p3evaluated to9e-4afterwards.his now iterated to a fixed point (bounded byne + 1iterations, socyclic dependencies cannot loop forever), and
roots_foundis derivedfrom the converged values. As a side effect,
reinit_explicit_rootsnowalso sees a consistent
w. Models whose root functions don't depend onh-- the common case -- just do one additionalfrootevaluationconfirming that nothing changed.
Closes #3233
Co-Authored-By: Claude Opus 5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_016ZKvTDruCaD3vHVXXXjz8o