Chain PEtab v2 experiment periods natively in the JAX simulator#3198
Draft
FFroehlich wants to merge 10 commits into
Draft
Chain PEtab v2 experiment periods natively in the JAX simulator#3198FFroehlich wants to merge 10 commits into
FFroehlich wants to merge 10 commits into
Conversation
Previously, PEtab v2 experiments with more than two periods were collapsed into SBML events at import time via ExperimentsToSbmlConverter, for both the sundials and JAX backends. For JAX, this meant period switches were driven by root-finding on synthetic indicator parameters baked into the compiled model rather than by directly chaining simulation calls. For the JAX backend, skip that conversion entirely and instead run one ODE integration per experiment period directly in JAXModel.simulate_condition, carrying state and heaviside/event state across period boundaries the same way pre-equilibration already hands off into the main simulation. JAXProblem's measurement bucketing, parameter mapping, and reinitialisation resolution are generalised from a hardcoded two-phase (preeq + main) model to arbitrary period counts. The sundials backend is unaffected. Along the way, fixes several latent bugs that were only reachable once JAX stopped seeing SBML-converted (indicator-only) condition tables: condition tables with multiple simultaneous changes, state reinitialisation lookups against the (long-format) condition table, a "preequilibration" substring-matching heuristic that depended on the converter's naming convention, and a couple of shape bugs in JAXModel for single-state models and models without observable/noise parameter overrides. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
- _petab_importer.py: raise NotImplementedError for condition-table changes targeting anything other than a species or parameter (e.g. compartments), which the JAX backend has no runtime mechanism to apply; keep is_state_variable() (species, compartments, and rule-governed entities) as the fixed-parameter exclusion filter. - petab.py: add_default_experiment_names_to_v2_problem now reads condition ids from petab_problem.conditions instead of the long-format condition_df, which omits conditions with zero changes (e.g. default no-op conditions). - petab.py: rewrite _build_simulation_df_v2 (and add the _dynamic_condition_index_map helper) to index into the 3D (experiment, period, timepoint) measurement arrays introduced by the period-chaining refactor, instead of a stale flat condition index.
- petab.py: PEtab v2 has no parameterScale column at all (unlike v1); replace the now-broken parameter_df lookups with the LIN scale constant already used elsewhere for v2 parameters. - petab.py: give every experiment period exactly one dynamic-condition label (first non-preequilibration condition id, or a synthesized one for periods without any condition table changes) instead of one label per condition id attached to the period. Periods with several simultaneous condition ids -- e.g. PySB's converted indicator encoding, which tags every kept period with both an experiment-indicator and a preequilibration-toggle condition id -- were otherwise being counted as multiple simulation legs, producing duplicate/misaligned rows in the simulation dataframe. - _petab_importer.py: clarify (comment only) that PySB models keep going through ExperimentsToPySBConverter for both backends, since PySB condition-table targets are frequently pysb.Observable names aliasing an underlying pysb.Initial/Expression, which JAXProblem's native per-period resolution has no equivalent for.
…radient check - petab.py: split the unwieldy nested get_overrides closure in JAXProblem._get_measurements into small, module-level, independently testable helpers (_override_placeholder, _resolve_override_symbol, _split_override_column, _override_triple_from_matrix, _column_overrides). petab_problem.parameter_df is threaded through lazily (accessed only once actually needed, inside the string-override branch) rather than eagerly per call -- eager access was tried first but broke SciML models with array-valued parameters, where building parameter_df emits a pydantic serialization warning. - petab.py: apply walrus-operator (:=) assignments at a few single-use assignment-then-check sites, and remove incidental dead code found along the way (an if/else with identical branches in _get_measurements, a redundant two-pass list comprehension in add_default_experiment_names_to_v2_problem). - test_petab_v2_multiperiod.py: replace the finite-difference-based gradient cross-check with a closed-form analytical derivative of the segment-wise exponential-decay solution, avoiding the need for a numerical approximation in the test.
np.where(par_mask, 0.0, mat) fails when mat is a fixed-width numpy string array (all entries are parameter references, no numeric values for np.stack to promote against object dtype). Revert to in-place assignment, which numpy handles correctly regardless of mat's dtype. Introduced in 58d7e7e; broke 7 previously-passing petabtests v2 suite cases (0003/0014/0015/0021 sbml, 0003/0014/0015 pysb, jax=True).
…urements Replace the bare 3-tuples threaded through the observable/noise parameter override machinery (_column_overrides, _override_triple_from_matrix, get_overrides's dict-of-tuples, and a hand-rolled zip-and-concatenate loop) with an OverrideColumn NamedTuple exposing .numeric/.mask/.index fields plus .placeholder()/.concatenate() constructors. De-nest _get_measurements's five closures (get_overrides, placeholder_row, get_iy_trafos, pad_measurement, pad_and_stack), which existed purely to close over local state, into module-level functions taking that state as explicit parameters. Replace the flat, comment-numbered 12-element tuple stored per (experiment, period) with a _PeriodMeasurements NamedTuple, so downstream padding/stacking reads named fields instead of positional indices. Consolidate the two near-identical "all-masked, single-timepoint placeholder period" blocks (a real period with no measurements in its window, and a padding period that doesn't exist for a given experiment) into a single _masked_placeholder_period helper. Purely structural; verified against the full 94-case tests/petab_test_suite/test_petab_v2_suite.py (71 passed, 23 skipped, 0 failed - unchanged from baseline), the multiperiod chaining tests, the SciML tests, and the JAX performance regression suite.
There was a problem hiding this comment.
Pull request overview
This PR updates AMICI’s JAX PEtab v2 simulation path to support native chaining of arbitrary numbers of experiment periods (instead of collapsing periods into SBML events), and adjusts measurement/parameter handling and tests accordingly.
Changes:
- Implement per-period measurement bucketing/padding and per-period parameter + reinitialisation resolution in
amici.sim.jax.petab.JAXProblem. - Update
amici.sim.jax.model.JAXModel.simulate_condition(_unjitted)to accept a leading “period” axis and chain one ODE integration per period. - Add/adjust regression and performance tests to match the new period-axis API and validate multi-period correctness + gradients.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/sbml/testSBMLSuiteJax.py | Updates direct simulate_condition call sites to add a leading period axis (size 1). |
| tests/performance/test_jax_regression.py | Adapts performance harness to the new period-axis inputs and list-based per-period stats. |
| python/tests/petab_/test_petab_v2_multiperiod.py | Adds new functional + gradient tests for native multi-period chaining and importer behavior. |
| python/sdist/amici/sim/jax/petab.py | Major refactor: per-experiment/per-period measurement bucketing, override parsing, and generalized (N-period) preparation logic. |
| python/sdist/amici/sim/jax/model.py | Refactors simulation to chain per-period integrations; updates likelihood/observable evaluation to accept per-timepoint parameters/TCL. |
| python/sdist/amici/sim/jax/_simulation.py | Fixes jnp.repeat usage to repeat along axis 0 for eventless solve paths. |
| python/sdist/amici/importers/petab/_petab_importer.py | Skips experiments-to-events conversion for JAX SBML imports and adds JAX-specific condition-target validation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+125
to
+132
| def resolve_row(entry: list | float) -> list: | ||
| if isinstance(entry, list): | ||
| return [_resolve_override_symbol(v, parameter_df) for v in entry] | ||
| return [] if pd.isna(entry) else [entry] | ||
|
|
||
| rows = col_values.str.split(petabv2.C.PARAMETER_SEPARATOR).apply( | ||
| resolve_row | ||
| ) |
Comment on lines
+1621
to
+1625
| all_condition_targets = { | ||
| change.target_id | ||
| for condition in self._petab_problem.conditions | ||
| for change in condition.changes | ||
| } |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #3198 +/- ##
==========================================
+ Coverage 75.73% 76.14% +0.41%
==========================================
Files 317 318 +1
Lines 20903 21207 +304
Branches 1484 1483 -1
==========================================
+ Hits 15830 16149 +319
+ Misses 5061 5046 -15
Partials 12 12
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
JAXModel.simulate_condition[_unjitted] constructed several default argument values eagerly, at module-import time, via jnp.array(...). If jax_enable_x64 is only enabled after this module is first imported, those defaults freeze to float32 while every other (call-time- constructed) array flowing through the same call ends up float64 -- surfacing as a `body_fun must have the same input and output structure` crash inside diffrax's adaptive stepping loop whenever a caller relies on the default t_zero. Switch all such defaults to a None sentinel, constructed lazily inside the function body instead. Also fix python/tests/test_jax.py::test_conversion/test_dimerization, which weren't updated for simulate_condition[_unjitted]'s now-required leading period axis on p/ts_dyn/ts_posteq/my/iys/iy_trafos/ops/nps.
The notebook hardcoded the SBML-event-converter's synthetic condition
id ("_petab_experiment_condition___default__"), which no longer
applies now that JAX skips that conversion and uses real condition/
experiment ids directly (here, "__default__"). Also add the same
leading-period-axis fix as the previous commit to a cell that manually
reproduces JAXModel.simulate_condition's internals.
…t test - Fix _split_override_column silently dropping numeric observable/noise parameter overrides on object-dtype columns: resolve each entry's own type instead of routing the whole column through the string-only `.str.split` accessor, which turned every non-string entry into NaN. - Cache the set of condition-table override targets on JAXProblem instead of rebuilding it on every load_reinitialisation call (once per period per experiment). - Add a regression test documenting that JAXModel._handle_t0_event reuses the previous period's ending heaviside state unconditionally for i>0, so a state reinitialisation that crosses a piecewise trigger's threshold doesn't get its event state re-evaluated until/unless the ODE integrator crosses the threshold again during that period.
… simulator _handle_t0_event previously short-circuited whenever it was handed a non-empty heaviside state, unconditionally carrying it over from the preceding preequilibration or experiment period instead of checking whether the (possibly reinitialised) incoming state actually still matches it. A state reinitialisation or parameter change at a period boundary that crosses an event's trigger threshold went undetected until the ODE integrator happened to cross it again during that period. The trigger condition is now always re-evaluated against the actual incoming state, using the previous heaviside state only as the pre-transition reference for detecting a crossing, exactly as already done for a genuine t=0. Updates the regression test added for this behavior to assert the corrected (re-evaluated) result instead of pinning the previous carry-over behavior.
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.
Previously, PEtab v2 experiments with more than two periods were
collapsed into SBML events at import time via
ExperimentsToSbmlConverter, for both the sundials and JAX backends.
For JAX, this meant period switches were driven by root-finding on
synthetic indicator parameters baked into the compiled model rather
than by directly chaining simulation calls.
For the JAX backend, skip that conversion entirely and instead run
one ODE integration per experiment period directly in
JAXModel.simulate_condition, carrying state and heaviside/event state
across period boundaries the same way pre-equilibration already hands
off into the main simulation. JAXProblem's measurement bucketing,
parameter mapping, and reinitialisation resolution are generalised
from a hardcoded two-phase (preeq + main) model to arbitrary period
counts. The sundials backend is unaffected.
Along the way, fixes several latent bugs that were only reachable
once JAX stopped seeing SBML-converted (indicator-only) condition
tables: condition tables with multiple simultaneous changes, state
reinitialisation lookups against the (long-format) condition table,
a "preequilibration" substring-matching heuristic that depended on
the converter's naming convention, and a couple of shape bugs in
JAXModel for single-state models and models without observable/noise
parameter overrides.
Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv