Chain PEtab v2 experiment periods natively in the JAX simulator - #3198
Draft
FFroehlich wants to merge 43 commits into
Draft
Chain PEtab v2 experiment periods natively in the JAX simulator#3198FFroehlich wants to merge 43 commits into
FFroehlich wants to merge 43 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.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #3198 +/- ##
==========================================
+ Coverage 78.62% 79.11% +0.49%
==========================================
Files 318 319 +1
Lines 21102 21743 +641
Branches 1487 1488 +1
==========================================
+ Hits 16591 17202 +611
- Misses 4503 4533 +30
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.
…d condition definitions cast_to_sym constructed sp.Float from a plain Python float without a precision argument, which defaults to ~15 significant *string* digits even though the underlying double is preserved exactly. This meant a parameter whose nominal value derives from an irrational SBML constant (e.g. `pi`) printed as a different decimal literal than the same constant emitted directly by a code printer (e.g. `jnp.pi`), so an exact equality comparison between the two silently failed. Building from repr() (the shortest round-trip-exact string) fixes the printed precision to match the actual bit pattern. Root-caused via SBML Semantic Test Suite case 00958 (P7 = piecewise(2, eq(P1, pi), 3) with P1 = pi), and verified against the exact failing CI shard (601-1200): 360 passed, 240 skipped, 0 failed. Separately, PetabImporter's fixed-parameter collection assumed every condition ID referenced by an experiment period has a matching entry in a condition table, and crashed with a KeyError otherwise. Per the PEtab v2 spec, a condition that makes no changes need not be defined at all (a problem can reference a condition purely as a timepoint label with zero condition_files). Added _get_condition_changes to treat a missing/empty condition ID as contributing no changes instead of erroring, fixing several PEtab SciML testsuite cases with no condition table at all.
…ndition
JAXProblem indexed condition_ids[0] to get a condition for keying a neural-
network output evaluation (_eval_nn), assuming a period always references
at least one condition. Per the PEtab v2 spec a period need not reference
any condition at all (e.g. one that only marks a timepoint, with nothing
to change there), so condition_ids can legitimately be empty, crashing
several PEtab SciML testsuite cases (002, 005, 006, 008, 011, 013, 014,
017, 021, 024, 027, 030, 032, 033, 034) with no condition table at all.
Added _first_condition_id to fall back to "" when there are no conditions.
"" can never collide with a real condition ID, and _eval_nn already has a
graceful default ("0"-keyed) array input lookup for exactly this case, so
this only replaces a crash with the fallback path that already existed.
Verified against the local PEtab SciML testsuite: all 15 IndexError cases
now pass (confirmed individually and via a full sequential run).
…less periods A period (or preequilibration phase) that references zero PEtab conditions is indistinguishable from an actual padding slot when both are represented as an empty condition-id list, so `load_reinitialisation` treated a real, condition-less period as padding and silently skipped reinitialisation -- breaking hybridization targets (e.g. a neural network setting a state's initial value) that don't need any condition to resolve. Use `None` to mark true padding/no-preequilibration and keep `[]` for a real period with no condition. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
#3224 fixed the JAX exporter's float truncation at the code printer, which covers every emitted Float rather than only those constructed via cast_to_sym. Building the Float from repr() is also slightly wrong: it reports 56 bits of precision for a value that only has a double's 53, and that phantom precision propagates through sympy, e.g. exp(x).evalf() no longer matches the exact double result. SBML semantic test suite case 00958 (the original motivation) still passes without it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
The `ys` loop carry in `solve` is seeded with `x0`, but the `hs` carry was seeded with plain zeros. When the loop body never runs -- `t0 >= ts[-1]`, which is exactly what a zero-duration period looks like -- the carry is returned untouched, so states were preserved while the heaviside state silently reset to all zeros instead of being chained on from the previous period. This is currently masked: zero-duration periods only arise as padding slots, which sort last and whose outputs are masked out of the likelihood. It becomes load-bearing as soon as anything is scheduled after such a slot, since `_handle_t0_event` would then re-evaluate triggers against a cleared heaviside state and could apply event assignments spuriously. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
…eriment A neural network that drives a state supplies that state's *initial value*, but `_state_needs_reinitialisation` returned True for it regardless of which period was being set up, so every period boundary in a chained experiment re-ran the network and overwrote the integrated trajectory with an initial value. Thread an `is_initial` flag through `load_reinitialisation` and apply the network-driven branches only for the first period of an experiment's own dynamic chain. Condition-table changes are unaffected -- they apply wherever their condition does. This is latent today: every problem in the PEtab SciML test suite has a single dynamic period, so the first period is the only one. It bites exactly where this branch's period chaining meets a state-valued hybridization target. Also drop a `hasattr(self, "nn_output_ids")` guard in the same code path. `nn_output_ids` is a class property, so the guard was always true and the condition collapsed to one that never consulted `nn_output_ids` at all, disagreeing with `_state_needs_reinitialisation` about which states are reinitialisable. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
Post-equilibration measurements are attached to each experiment's own last period (`is_own_last`, petab.py:698), but `eq()` is gated on the globally last period index (`is_final = i == n_periods - 1`, model.py:806, where `n_periods` is the problem-wide `_max_periods`). For an experiment with fewer periods than the problem maximum, equilibration therefore never runs where its `time=inf` rows live. Those rows are flagged valid (`posteq_valid`, petab.py:791) and so enter the likelihood, scored against the end-of-dynamics state -- and a full steady-state solve is meanwhile spent on a padding slot whose outputs are masked. Reproduced: with a two-period `exp_long` and a one-period `exp_short` carrying a `time=inf` row, the post-equilibrated value comes back as 0.6065 (i.e. exp(-0.5), the end of the dynamic phase) where the steady state of `xx' = -kk*xx` is 0. Marked `xfail(strict=True)` so that fixing the gating flips it to XPASS rather than silently passing. The fix has to be a per-experiment masked blend rather than a static slot gate: `eq()` reassigns `x_solver`/`h` (model.py:646-647), which become the carry into the next period, so equilibrating at a slot where another experiment's chain continues would clobber that experiment's state. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
Post-equilibration measurements are attached to each experiment's own last period (`is_own_last`), but `eq()` was gated on the globally last period index (`is_final = i == n_periods - 1`, where `n_periods` is the problem-wide maximum). For an experiment with fewer periods than that maximum, equilibration therefore never ran where its `time=inf` rows live: those rows are flagged valid and enter the likelihood, so they were scored against the end-of-dynamics state, while a full steady-state solve was spent on a padding slot whose outputs are masked. Reproduced with a two-period and a one-period experiment: the post-equilibrated value came back as exp(-0.5), the end of the dynamic phase, where the steady state of `xx' = -kk*xx` is 0. Which period is an experiment's last is per-experiment, but the period loop is shared across the vmapped experiment axis, so the decision is split in two. `has_posteq_slot` is static and decides whether the steady-state solve is traced at this slot at all, keeping it off slots no experiment post-equilibrates in. `do_posteq` is a per-experiment traced flag deciding whether the result is adopted; it has to be a blend rather than a branch because `eq()` replaces the state handed to the next period, so equilibrating unconditionally would clobber the trajectory of any experiment whose chain continues past that slot. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
`JAXProblem` precomputed `x_reinit` as data in host Python, before any period was integrated. Two consequences: a condition change whose targetValue referenced another state (`A = A + 5*B`) had no state to read and raised NotImplementedError; and the reinitialisation lived outside the generated model, so the model was not deployable on its own. Emit the condition-table changes that target a state into the model file instead, as an `_x_reinit(x, p, pc)` method bound exactly like `_delta_x`, and evaluate it in the period loop against the previous period's terminal state. The generated model now carries its own reinitialisation code and needs neither the PEtab problem nor AMICI's PEtab layer to run. Notes on the shape: - Fresh `_x_reinit_<i>` targets mean every change of a period reads the pre-change state, so PEtab's simultaneous-change semantics fall out. - A third argument `pc` carries PEtab parameters that are not model parameters (a parameter used only as a condition targetValue never reaches `p`). The expressions themselves are still fully baked in. - `tcl`/`w` are deliberately unavailable: reinitialisation runs before the period's conservation laws and expressions are computed, and those are derived from the reinitialised state. A targetValue referencing an assignment-rule target is rejected at code-generation time with a message naming the condition, target, value and offending symbol, rather than failing obscurely later in the code printer. - Because the condition table is now baked into the model, an edited targetValue against a cached model directory would silently simulate stale values. `JAXProblem` re-enumerates the rows and compares them to the model's, raising a regenerate-the-model error on mismatch. - Network-driven (hybridization) initial values stay on their data route and are applied after the expressions, preserving the previous precedence and the `is_initial` semantics. - MODEL_API_VERSION 0.0.4 -> 0.0.5, so a cached model generated without `_x_reinit` is reported as stale instead of failing deep in the solve. - The v1 import path upgrades to v2 before code generation (only when `jax=True`) so the exporter can see the condition table. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
The JAX regression suite's Tier 1 models are hand-written `JAXModel` subclasses rather than generated code, so each pins `api_version` in its own source. Emitting state reinitialisation into generated models bumped `MODEL_API_VERSION` to 0.0.5, which left these six models mismatched and failed the whole suite with "JAXModel API version mismatch, please regenerate the model class" (15 errors). Only the pinned version needs to change: `_x_reinit` and the two `reinitialisation_*` properties are concrete on the base class with no-op defaults, so a model with no reinitialisations inherits correct behaviour without restating them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
Emitting condition-table state changes into the generated model is a breaking change for two groups that CI cannot catch: users with existing generated JAX models (which now report an API version mismatch rather than failing later in the solve), and users with hand-written JAXModel subclasses (which pin api_version in their own source). Editing a condition table now also requires regenerating the model, since those expressions are baked into it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
`run_simulations(..., simulation_experiments=...)` raised "vmap got inconsistent sizes for array axes to be mapped" for any proper subset. The measurement and override arrays (`_ts_dyn`, `_ts_posteq`, `_my`, `_iys`, `_iy_trafos`, `_ts_masks`, and the observable/noise overrides) are built once for every experiment of the problem, while the parameter and reinitialisation arrays are built only for the experiments being simulated; both were handed to the same vmap, which only lined up when the "subset" happened to be all of them. Index the full-length arrays by the positions of the selected experiments, via a new `_experiment_indices` helper. It returns an array rather than a list because these index jax arrays too, and jax rejects a plain sequence as a multidimensional index. `h_mask` was built over all of the problem's experiments, zeroing the rows of the ones left out -- itself a workaround for this mismatch, and part of it. With the axes consistent it is one all-ones row per simulated experiment, which is exactly what it already evaluated to whenever no subset was requested. Verified beyond "no longer raises": per-experiment log-likelihoods sum to the all-experiment value, and selection is by id rather than position, so a subset that ran but scored the wrong experiment's measurements would fail. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
`JAXProblem` is an `eqx.Module`, so `_ts_masks` is a pytree leaf and becomes a tracer under `eqx.filter_jit`. Deriving the post-equilibration bookkeeping from it with `np.asarray(...)` therefore raised `TracerArrayConversionError` for any jitted call -- which is how the JAX PEtab example notebook builds its optimisation step, so it broke the notebook, sphinx and readthedocs builds while every non-jitted test kept passing. Whether an (experiment, period) cell post-equilibrates is a property of the PEtab tables, not of any traced array: an experiment has post-equilibration rows iff it has a measurement at non-finite time, and they attach to its own last dynamic period. Derive it from that instead. The remaining use of `_ts_masks` is its `.shape`, which is fine on a tracer. The regression test drives `run_simulations` through `filter_jit` + `filter_value_and_grad`, mirroring the notebook; without the fix it reproduces the CI error verbatim. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
The dtype fix in `_pad_measurement` landed in c8fd51d, whose message describes only the tracer fix it shipped alongside; this documents it and adds the test it was missing. An absent post-equilibration part is an empty array, which numpy types as float64 regardless of the data, so concatenating it against the integer row indices widened them to float. `_petab_measurement_indices` then became a float index on the simulation DataFrame, which no longer compares equal to the measurement table's integer index. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
Every experiment is padded out to the problem-wide period count with trailing placeholder periods. Those are zero-duration steps meant to leave the carried-over state alone, but they were anchored at `dyn_periods[-1].time` -- the *start* of the last real period -- rather than at the time the chain has actually reached by then. For an experiment whose last period carries measurements, that is a jump backwards: a single-period experiment measured out to t=4 got a padding period anchored at t=0. The period start times stopped being monotonic, and `JAXModel._handle_t0_event` re-evaluated trigger conditions against a time the solver had already left behind. This is latent today -- padding periods are always trailing and everything they touch is masked out of the likelihood -- but it is wrong as a representation and it blocks tightening the per-slot padding. Anchor them at the end of the last real period instead, i.e. the largest finite measurement time of the experiment (post-equilibration rows carry a non-finite time and do not advance dynamic time), falling back to the period's own start time when it has no measurements. The new `_experiment_end_time` is shared by the two places that have to agree on this: the placeholder time points in `_get_measurements` and the `t_zeros` entries in `period_start_times`. Also resolve each cell's `load_reinitialisation` once instead of calling it a second time just to take the other half of the returned pair. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
Post-equilibration was run inside the period loop, gated by a static per-slot `posteq_slots` flag (is this period index some experiment's last?) and a traced per-slot `posteq_mask` blend deciding whether to adopt the result. That was an awkward fit for something that is a property of the experiment rather than of any one period, and it cost real work: in a problem mixing 1..4-period experiments that all post-equilibrate, every slot traced a steady-state solve and ran it for the whole vmapped batch, so three of four solves per experiment were computed and then discarded by the blend. Each (experiment, period) cell also carried a full post-equilibration column block, of which all but one per experiment were dead. Run it once instead, after the whole chain. The measurement arrays are laid out along a single flat time axis -- the P dynamic blocks followed by one trailing post-equilibration block -- so `ts_posteq` loses its period axis, `my`/`iys`/`iy_trafos`/`ops`/`nps`/`ts_mask` lose theirs, and `posteq_slots`/`posteq_mask` collapse to one `do_posteq` flag per experiment. `_simulate_period` no longer knows about post-equilibration at all. For an experiment shorter than the problem-wide period count this now happens after its trailing padding periods rather than at its own last real period. That is equivalent only because those periods are zero-duration no-ops that re-evaluate `_handle_t0_event` at a time which is not in the past -- which is true as of the preceding commit. On the worst case above: llh 0.361s -> 0.294s, gradient 0.594s -> 0.582s. The gradient is dominated by the adjoint solve of the real work, so the point of this is the simplification rather than the speedup. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
The previous commit changed `simulate_experiment`'s measurement arrays to a single flat time axis, but only the two direct callers in `python/tests/test_jax.py` that fail fast locally were updated. Three more pass the old per-period shapes and broke in CI: * `tests/performance/test_jax_regression.py` (JAX regression suite) * `tests/sbml/testSBMLSuiteJax.py` and `tests/sbml/testSBMLSuite.py` (SBML semantic test suites) * `check_fields_jax` in `python/tests/test_jax.py`, reached only from `test_conversion`/`test_dimerization`, which skip locally for lack of BioNetGen -- it is beartype-wrapped, so it failed on the jaxtyping annotation rather than on the concatenate `ts_posteq` loses its period axis and `my`/`iys`/`iy_trafos`/`ops`/`nps` become one-dimensional; `p` and `ts_dyn` keep theirs. The example notebook needs no change: it slices `_ts_posteq[ic, :]`, `_my[ic, :]` and friends straight off `JAXProblem`, so it follows the new field shapes automatically. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv
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