Chain PEtab v2 experiment periods natively in the JAX simulator - #3198
Draft
FFroehlich wants to merge 24 commits into
Draft
Chain PEtab v2 experiment periods natively in the JAX simulator#3198FFroehlich wants to merge 24 commits into
FFroehlich wants to merge 24 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.
Contributor
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 78.46% 77.88% -0.59%
==========================================
Files 317 318 +1
Lines 20974 21293 +319
Branches 1483 1482 -1
==========================================
+ Hits 16458 16584 +126
- Misses 4508 4701 +193
Partials 8 8
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
FFroehlich
marked this pull request as draft
July 4, 2026 21:48
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.
…t-refactor-j48b04 # Conflicts: # python/sdist/amici/sim/jax/petab.py
…t-refactor-j48b04 Resolves conflicts between the native per-period JAX chaining introduced in this branch and main's independent SciML/PEtab-v2 refactors landed since the last merge (libpetab linting follow-ups, CSE/ImplicitAdjoint JAX perf work, and several JAX PEtab v2 condition-handling bugfixes). Where both sides had independently rewritten the same functionality, this keeps the per-period-chaining design (not present on main) while adopting main's genuine fixes/optimizations by threading them into that design rather than reverting to main's simpler single/two-phase structure: - _handle_t0_event: adopted main's equivalent live-reevaluation fix (main independently arrived at the same fix as this branch's earlier commit) with this branch's fuller comment covering the per-period reinit case. - Observable/noise parameter override resolution: kept this branch's safe per-entry-type-checked parsing, but threads through main's precomputed fixed_parameter_values dict (avoids rebuilding petab_problem.parameter_df per period, and excludes array-valued SciML parameters from substitution). - _prepare_experiments: kept this branch's N-period-aware parameter/ reinit array construction, adopted main's parameter-scale lookup for the (legacy, non-petabv2.Problem) unscaled-parameters branch. - Removed now-dead additive helpers from main's non-N-period design that this branch's equivalents already supersede (_get_period_condition_ids, _experiment_indices, _resolve_condition_target_value). Also fixes a latent jnp.stack([]) crash in load_model_parameters for models with no free SBML parameters (only literal rate constants), surfaced by a new test added upstream. Verified against the full JAX PEtab v2 (multiperiod + general) test suite, the SBML JAX semantic suite's event-tolerance-flagged cases, and the JAX performance regression suite -- all green, plus two new upstream regression tests for noise-parameter placeholder handling.
…le test call site _prepare_experiments's is_preeq branch resolved reinitialisation condition ids from a globally deduplicated set (_get_preequilibration_condition_ids), rather than per-experiment, so mask_reinit_array/x_reinit_array could end up shorter than p_array whenever experiments shared a preequilibration condition id, causing a vmap shape mismatch in run_preequilibration. Resolve each experiment's own preequilibration period condition ids directly, mirroring how load_model_parameters already does it for parameters. Also fix test_steady_state_event_no_recompile_across_conditions (added independently on main before the period-chaining merge), whose simulate_condition call was still missing the period axis on p/ts_dyn/ts_posteq/my/iys/iy_trafos/ops/nps. Also apply two previously-identified fixes surfaced by CI: empty-string override tokens in _split_override_column, and NaN-experiment-id row selection in _build_simulation_df_v2.
… stale simulation_df experimentId column
_get_iy_trafos built its return array by iterating over
petab_problem.observables (one entry per observable in the model)
instead of gathering per measurement row via iys, silently producing an
array of the wrong length whenever the number of measurement rows in a
period differs from the number of observables. This caused a spurious
"index can't contain negative values" crash in _pad_and_stack for
benchmark models with more than one observable (e.g. SalazarCavazos_
MBoC2020, Brannmark_JBC2010, Laske_PLOSComputBiol2019), and would have
silently mismatched sigma/observable-transform lookups otherwise.
Resolve the transformation by observable id first, then gather onto
iys's own length.
Also fix _build_simulation_df_v2's experiment-id row matching: the
"__default__" experiment sentinel is coerced to jnp.nan for the JAX
side, but the underlying measurement_df always stores the literal
string "__default__" (never a real NaN), so neither a string .query()
nor an .isna() mask ever matched it, leaving observableParameters/
noiseParameters all-NaN in the simulation output. Match against the
literal sentinel string instead. This fixes PEtab Testsuite cases with
implicit ("__default__") experiments (e.g. cases 0003, 0006, 0014, 0015
upgraded from PEtab v1).
tests/petab_test_suite/test_petab_suite.py's JAX path added back a
v1-style simulationConditionId column for comparison against v1 ground
truth, without dropping the v2-style experimentId column also present
in AMICI's output; petabtests.evaluate_simulations determines the PEtab
version from column presence and errors out when both are present.
Mirrors the existing tests/sbml/SBMLTestModels/ entry for the non-JAX (C++) counterpart; tests/sbml/SBMLTestModelsJax/ is regenerated on every test run and was previously untracked but not ignored.
iys_real (per-measurement-row observable indices) defaulted to float64 when a period has zero real dynamic measurements (an empty list comprehension), since np.array([]) infers float64 without an explicit dtype. _get_iy_trafos now gathers via trafo_by_index[iys], which requires integer indices, so this surfaced as "IndexError: arrays used as indices must be of integer (or boolean) type" for models with such periods (e.g. Blasi_CellSystems2016 in the benchmark collection).
_get_measurements computed ts_posteq (time points) for post-equilibrium measurements correctly, but never computed the corresponding my (measured value), iys (observable index), iy_trafos, or observable/noise parameter overrides for them -- those fields only ever covered the dynamic-phase measurements. Post-equilibrium rows were still marked valid and included in the log-likelihood, but with their measurement silently zeroed, their observable identity defaulted to index 0, and their noise override defaulted to the numeric literal 0 instead of the correct free-parameter reference. This produced a near-infinite log-likelihood (division by a near-zero noise value) for any model whose observable-parameter/noise overrides differ between dynamic and post-equilibrium measurements (e.g. Blasi_CellSystems2016 in the benchmark collection, where nearly all measurements are post-equilibrium comparisons sharing a single free "sigma" parameter). Compute the post-equilibrium counterparts and concatenate them onto the dynamic-phase arrays, matching the `len(ts_dyn) + len(ts_posteq)` layout _pad_and_stack already expects.
… fix example notebook, drop accidentally-committed test model artifacts PEtab v2 nomenclature calls a chained sequence of periods an "experiment" rather than a "condition", so rename the JAX simulation entry points to match. Also: - Fix a malformed notebook cell (source stored as a single string instead of the list-of-lines format used by every other cell, and a missing trailing newline) introduced by a previous edit. - Remove two PetabImporter-generated test model directories under 1.0.1/ that were accidentally committed, and gitignore that pattern: python/tests/conftest.py points AMICI_MODELS_ROOT at the repo root for the test session, so these are regenerated fresh on every local test run and should never be tracked.
simulate_experiment[_unjitted] requires a leading period axis on p/ts_dyn/ts_posteq/my/iys/iy_trafos/ops/nps; this call site predates that requirement and was missed by the earlier period-axis fixes in test_jax.py and the example notebook, causing "TypeError: iteration over a 0-d array" in _x0's p[0] indexing for any SBML test suite case in sensitivity_check_cases.
… expressions to JAX _get_measurements previously combined a period's dynamic-phase and post-equilibrium-phase measurement data into single arrays, then re-split them by len(ts_dyn) in three separate places (_pad_and_stack, the ts_masks padding, and the petab_indices padding). That split point was only recoverable by convention across every field, which is what let post-equilibrium overrides silently fall back to zero-filled placeholders in an earlier bug. _PeriodMeasurements now tracks the dynamic and post-equilibrium portions as separate fields throughout (mirroring ts_dyn/ts_posteq), removing the concatenate-then-reslice step entirely. Condition table changes with a compound symbolic target_value (e.g. k1 + k2) previously raised NotImplementedError; _resolve_petab_change_value now compiles any target_value expression to JAX via the same sympy-to-JAX code printer (AmiciJaxCodePrinter) used to generate the model's own equations, with each free symbol resolved by the calling site's existing rules (model parameter, estimated PEtab parameter, or fixed nominal value). A numeric literal or single parameter reference is just the zero/one-free-symbol case of the same mechanism, so the previous separate number/symbol special-casing is gone too.
…crashing Compiling a condition's target_value can now succeed for expressions that were previously rejected outright, including ones referencing another state (e.g. A = "A + 5.0", found by the PEtab v2 test suite's case 0028/0031). resolve_symbol had no case for a state id, so it fell through to a parameter_df lookup and raised a confusing pandas KeyError instead of a clean, catchable NotImplementedError. Resolving a state's value would require the actual simulated trajectory at that period boundary, which load_reinitialisation cannot provide: x_reinit is precomputed once per experiment in _prepare_experiments, before any period is integrated. Both resolve_symbol closures now raise NotImplementedError for a state-referencing symbol, restoring the same graceful skip the PEtab test suite's wrapper already applies for genuinely unsupported cases.
It had exactly one call site and was a single dict.get(value, value) lookup; the wrapper added a function and a docstring for something that reads just as clearly inline.
…ndant with the v2 testsuite simulate_condition[_unjitted] were renamed to simulate_experiment[_unjitted] earlier in this branch; restore them as thin deprecated wrappers for backward compatibility, with a regression test confirming they still work and match the new names' output. Removed test_two_period_preequilibration_matches_analytical_solution and test_single_period_matches_analytical_solution: cross-checked against all 32 official PEtab v2 test-suite cases and confirmed cases 0009/0010/0017/0018 already exercise preeq+one-period chaining with plain numeric reinits under jax=True, and single-period (no chaining) is exercised throughout the wider suite already. The remaining tests in this file (three-period chaining, gradient-through-chain, event/heaviside-at-reinit, no-event-conversion) each cover ground no official test-suite case reaches.
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