Fix JAX PEtab backend: discontinuities, PEtab v1 import, event assignments#3196
Open
FFroehlich wants to merge 2 commits into
Open
Fix JAX PEtab backend: discontinuities, PEtab v1 import, event assignments#3196FFroehlich wants to merge 2 commits into
FFroehlich wants to merge 2 commits into
Conversation
Several related fixes to the JAX simulation backend surfaced while
investigating discontinuity handling:
* Clip the step-size controller onto known (explicit) discontinuity
times so time-triggered events are integrated precisely instead of
relying solely on root-finding.
* Represent explicit (time-triggered) Heaviside placeholders as
tracked variables instead of inlining them as `jnp.select`, matching
how implicit/state-dependent triggers are already handled. This
keeps the vector field smooth within each integration segment.
* `PetabImporter` now upgrades in-memory PEtab v1 problems to v2
instead of raising `NotImplementedError`, so the JAX backend (which
requires v2) can consume v1 PEtab problems.
* Fix wrong sigmas/log-likelihood for observables using PEtab's
standard `noiseParameter{n}_{observableId}` /
`observableParameter{n}_{observableId}` placeholder naming: these
were misclassified as ordinary model parameters instead of
per-measurement overrides, so every measurement got the first
observable's noise value. Ported the substitution already used by
the legacy v1 JAX importer to the newer PEtab v2 importer.
* Fix event assignments being silently dropped whenever a model's
total root count is odd (e.g. any odd number of plain SBML events
with assignments): `_apply_event_assignments` processed roots in
fixed windows of two, assuming every event comes in a Heaviside
+/- pair. Now processes every root individually.
Adds regression tests for each of the above.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Member
Author
|
known_discs fix is here #3195 |
There was a problem hiding this comment.
Pull request overview
This PR delivers a set of targeted fixes to AMICI’s JAX simulation/PEtab integration to improve correctness around discontinuities/events and broaden PEtab compatibility (v1→v2 upgrade), while aligning JAX results with SUNDIALS for PEtab-standard noise/observable placeholder naming.
Changes:
- Clip Diffrax step-size control to known explicit discontinuity times and keep explicit time-triggered Heaviside terms as tracked variables (avoid inlined symbolic time functions).
- Fix event assignment application for models with an odd number of roots by applying assignments per-root instead of in fixed pairs.
- Upgrade in-memory PEtab v1 problems to v2 during import and correct JAX handling of per-observable noise/observable placeholder parameters.
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 |
|---|---|
| python/tests/test_jax.py | Adds regression tests for explicit discontinuity handling and odd-root-count event assignments. |
| python/tests/petab_/test_petab_v2.py | Adds tests for v1→v2 in-memory upgrade and JAX vs SUNDIALS agreement for per-observable noise placeholders. |
| python/sdist/amici/sim/jax/_simulation.py | Clips the step-size controller to known discontinuities; fixes per-root event assignment application logic. |
| python/sdist/amici/importers/petab/_petab_importer.py | Implements PEtab v1 in-memory upgrade; adds JAX-only placeholder substitution/filtering; validates upgraded problem. |
| python/sdist/amici/exporters/jax/ode_export.py | Stops inlining Heaviside variables as symbolic functions of time so discontinuities can be handled via tracked variables/events. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3196 +/- ##
==========================================
- Coverage 77.05% 76.40% -0.65%
==========================================
Files 317 317
Lines 20903 20941 +38
Branches 1484 1483 -1
==========================================
- Hits 16106 16001 -105
- Misses 4789 4932 +143
Partials 8 8
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
dweindl
approved these changes
Jul 6, 2026
| assert stats["num_rejected_steps"] == 0 | ||
|
|
||
| except NotImplementedError as err: | ||
| if "The JAX backend does not support" in str(err): |
Member
There was a problem hiding this comment.
It's supposed to be supported - do we need these checks?
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.
Summary
A cluster of related fixes to the JAX simulation/PEtab backend, found while investigating how explicit (time-triggered) discontinuities are handled:
ClipStepSizeController(controller, jump_ts=...)) instead of relying solely on root-finding. Verified 10 orders of magnitude accuracy improvement and zero rejected steps on a synthetic explicit-discontinuity model.jnp.selectfunctions of time. Keeps the vector field smooth within each integration segment.PetabImporternow upgrades an in-memorypetab.v1.Problemto v2 (viato_files_generic+petab.v2.Problem.from_yaml, which auto-upgrades) instead of raisingNotImplementedError. This lets the JAX backend (which requires v2) consume v1 PEtab problems, e.g. from the PEtab benchmark collection.noiseParameter{n}_{observableId}/observableParameter{n}_{observableId}naming convention. The anchored classification regex only matched the barenoiseParameter{n}form, so these placeholders leaked into ordinary model parameters instead of being resolved via the per-measurement override arrays — every measurement ended up using the first observable's noise value. Fixed by porting the substitution already used by the legacy PEtab v1 JAX importer (v1/_sbml_import.py) into the newerPetabImporter, JAX-gated, operating only on an ephemeral importer-internal copy (never touching the PEtab problem itself or the SUNDIALS path). Verified against the PEtab benchmark model Boehm_JProteomeRes2014: chi2/llh now match SUNDIALS exactly (was off by ~7.4 in llh)._apply_event_assignmentsprocessed roots in fixed windows of two, assuming every event contributes a Heaviside +/− root pair. A plain SBML<event>with an assignment (no associated Heaviside) contributes a single root, so whenever a model's total root count was odd, the last root (in declaration order) was silently dropped — no error, no warning, no state update. Fixed by iterating over every root individually. Investigated and confirmed this doesn't regress the paired-Heaviside mechanism, including when Heaviside pairs and plain events are mixed in the same model (explicit events always sort before Heaviside pairs internally, so a dropped tail root there is always a Heaviside root with a zero delta).Test plan
python/tests/test_jax.py: existing discontinuity tests pass; addedtest_explicit_discontinuityandtest_event_assignments_odd_root_count(verified the latter fails without the fix and passes with it)python/tests/petab_/test_petab_v2.py: addedtest_petab_importer_upgrades_v1_problemandtest_jax_matches_sundials_with_per_observable_noise_parameters(JAX vs SUNDIALS llh agreement)Notes for reviewers
Two related pre-existing bugs were found during this investigation but are out of scope for this PR (tracked separately):
_known_discs(self, p)crashes withNameErrorwhen an explicit trigger time depends on a model expression (w) rather than pure parameters (e.g.Fujita_SciSignal2010's EGF pulse timing) — predates this change (PR Update PEtab SciML support to PEtab v2 #3165).🤖 Generated with Claude Code