Fix JAX PEtab condition-table resolution: state reinit, noise distribution, and differentiability - #3207
Merged
Merged
Conversation
Fixes a cluster of bugs surfaced by CI once the v1->v2 upgrade path (from the two separate PRs merged into this branch) started actually reaching previously-unreachable code: * get_simulation_conditions_v2: a dynamic period may reference multiple condition ids (e.g. a synthetic preequilibration-indicator condition alongside the real experiment condition). Measurements are only ever queried by experiment id, so multiple condition-id rows per experiment produced duplicate/misaligned measurement arrays, causing vmap batch-size mismatches (`vmap got inconsistent sizes`). Collapse to one row per experiment. * add_default_experiment_names_to_v2_problem: read condition ids from condition table elements instead of the long-format `condition_df`, which contributes zero rows for a condition with no changes (e.g. the default condition, or any no-op condition) -- exactly the "Experiment has no dynamic period with a condition id" case. * _build_simulation_df_v2: the synthetic default experiment id was overwritten with NaN before being reused to query observableParameters/noiseParameters from the measurement table, silently matching nothing. * _get_parameter_mappings: a condition table's target value can be a reference to another PEtab parameter id (not just a numeric literal), which crashed trying to cast the symbol straight to a jax array. Resolve it the same way `_map_experiment_model_parameter_value` already resolves other parameter references. * import_petab_problem (legacy v1 path): snapshot the pristine v1 problem before SBML/PySB compilation mutates it in place, so the later v1->v2 upgrade doesn't serialize an already-mutated (and potentially v1-lint-failing) problem. * pytest.ini: ignore the benign, documented petab1to2 warning when falling back from a v1-only noise distribution (log10-normal) to log-normal for v2 -- was being promoted to a hard error by this repo's `filterwarnings = error` policy, exactly when the v1->v2 upgrade path first became reachable. * ExampleJaxPEtab.ipynb / test_petab_suite.py: two consumers still expected `dynamic_conditions` to hold bare condition-id strings; it now holds tuples (to support multi-condition periods). Updated to match. Brings the official PEtab v1/v2 test suite (jax=True) from 28 failing down to 18 (case 0007's chi2 mismatch is a likely-inherent consequence of the log10-normal->log-normal fallback; the remaining ~8 cases cluster around condition-table-driven state reinitialisation and are tracked separately). jax=False path re-verified unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Fixes a regression: petab1to2's "Parameter scales are not supported in PEtab v2" warning (emitted whenever a v1 problem uses non-linear parameterScale, e.g. the lotka_volterra test fixture) was being promoted to a hard error by this repo's `filterwarnings = error` policy, breaking test_preequilibration_failure/test_serialisation. Verified benign via direct SUNDIALS simulation of the same converted v2 problem (bypassing JAX) for petab test suite cases 0019/0020, which also trip this warning: llh matches the ground truth solution exactly, confirming the dropped parameterScale is purely estimation-scale metadata that doesn't affect simulation values (petab v1's nominalValue is always stored in linear units). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…g in JAX PEtab backend Resolves the remaining 18 official PEtab test suite (jax=True) failures (cases 0007, 0010, 0011, 0013, 0016-0020, both formats): * JAXProblem._state_needs_reinitialisation/_state_reinitialisation_value/ load_reinitialisation looked up species-level condition-table overrides via the old wide-format `condition_df.loc[condition, state_id]`, but the current petab.v2 API returns `condition_df` in long format (conditionId/targetId/targetValue columns), so the lookup always missed and every such override silently fell back to the SBML default. Rewired to reuse `_parameter_mappings["targets_map"]` (built from `c.changes`), which parameter-target lookups already relied on correctly. Fixes cases 0010, 0011, 0013, 0017, 0018, 0019, 0020. * Traced case 0007/0016's chi2-only mismatches (LLH and simulated values already matched) to an upstream libpetab-python bug: petab1to2's `update_noise_dist` computes the merged v1->v2 `noiseDistribution` (e.g. folding `observableTransformation=log` into `log-normal`) but never returns it, so every upgraded observable silently reverts to `normal`, discarding any log/log10 transform. This corrupted `iy_trafos` (and thus chi2) even though AMICI's own log-likelihood code generation is unaffected, since it's derived from the pristine v1 problem directly. Added a workaround in `import_petab_problem` that recomputes the correct value from the pristine v1 problem and patches it onto the upgraded v2 observables. Fixes case 0016 outright. * With the above fixed, case 0007 has one genuinely inherent residual: PEtab v2 has no `log10-normal` distribution, so `log-normal` is substituted (with a warning); recomputing chi2 with log10 in place of log reproduces the expected ground-truth value exactly, confirming this is a real v1->v2 upgrade limitation, not an AMICI bug. Marked as an explicit, documented `pytest.xfail` rather than silently skipped. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The previous commit resolved condition-table state initial values through `_parameter_mappings["targets_map"]`, which `_get_parameter_mappings` builds once at construction time. For a value referencing an estimated parameter, that captured `self.parameters[i]` as a constant in a separate (frozen) pytree leaf. Consequences: * `update_parameters(...)` only replaces `.parameters`, never the cached `targets_map`, so the initial value stayed pinned at the nominal parameter value -- re-simulating after an update silently ignored it. * Differentiating via the documented `eqx.filter_grad(run_simulations)` idiom put the sensitivity into `grad._parameter_mappings[...]` rather than `grad.parameters`, so `grad.parameters` (what callers read) was 0 for any parameter used as an initial value. Verified on petab test suite case 0020 (initial_A estimated, entering the likelihood only through A(0)): pre-fix, updating initial_A left llh unchanged and grad.parameters[initial_A] was 0 while -8.70 leaked into the cache. Fix: resolve the reinitialisation value live from the raw condition-table change (`_condition_reinit_target_value` + `_resolve_condition_target_value`) inside `_state_reinitialisation_value`, which runs within the traced region via `_prepare_experiments`. Reading `self.parameters` there keeps the value a function of the current parameters. After the fix, autodiff matches central finite differences to ~1e-10 with zero gradient leaking into the cache, for estimated initial values (cases 0020, 0019), parameter-referenced initial values (case 0013) and ordinary rate parameters alike. Adds `test_condition_table_initial_value_is_differentiable` (the petab test suite skips derivative checks for jax, so this bug was uncaught). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Same construction-time freezing bug as the preceding commit, on the parameter-mapping path: `_map_experiment_model_parameter_value` read the override value from `_parameter_mappings["targets_map"]`, which caches `self.parameters[i]` at construction time. So a model parameter mapped by the condition table to an estimated parameter -- the standard PEtab pattern for condition-specific estimated parameters -- was frozen at the nominal value: `update_parameters` had no effect and its gradient leaked into the cache instead of `grad.parameters`. Verified: a condition setting model parameter `k1` to estimated `k1_c0` previously left llh unchanged under `update_parameters(k1_c0)` with `grad.parameters[k1_c0] == 0` (and -0.40 leaking into the cache); after the fix, forward responds and autodiff matches central finite differences. Fix: build the override lookup from the raw condition-table changes and resolve it live via `_resolve_condition_target_value` (which reads the live `self.parameters` inside the traced region), mirroring the reinitialisation fix. Adds `test_condition_table_parameter_override_is_differentiable`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes remaining correctness gaps in the JAX PEtab (v1→v2-upgraded) workflow around condition-table handling and observable noise distributions, and adds regression tests to ensure condition-table–driven initial values/parameter overrides remain differentiable w.r.t. JAXProblem.parameters.
Changes:
- Fix condition-table resolution in the JAX PEtab backend (state reinitialisation + condition-mapped parameter overrides), ensuring values are resolved “live” so gradients/
update_parameters(...)behave correctly. - Work around a
petab1to2upgrade bug where v1 observable transformations fail to propagate into v2noiseDistribution, restoring correct chi2/transform behavior. - Add JAX-specific test-suite adjustments (xfail for PEtab case 0007) and new gradient regression tests for differentiability.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/petab_test_suite/test_petab_suite.py | Adds a JAX-only xfail for PEtab case 0007 and aligns simulation-condition IDs for v1 expectations. |
| python/tests/test_jax.py | Adds regression tests ensuring condition-table initial values and parameter overrides remain differentiable. |
| python/sdist/amici/sim/jax/petab.py | Fixes condition-table lookup for v2 long format, resolves targets live for differentiability, and adjusts v2 experiment/condition handling. |
| python/sdist/amici/importers/petab/v1/_petab_import.py | Adds a workaround to patch v2 noiseDistribution after v1→v2 upgrade. |
| pytest.ini | Filters additional known/benign petab1to2 warnings. |
| doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb | Updates example to use tuple-form experiment conditions consistent with runtime output. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+2092
to
+2096
| @@ -2056,7 +2093,7 @@ def _build_simulation_df_v2(problem, y, dyn_conditions): | |||
| { | |||
| petabv2.C.MODEL_ID: [float("nan")] * len(t), | |||
| petabv2.C.OBSERVABLE_ID: obs, | |||
| petabv2.C.EXPERIMENT_ID: [experiment_id] * len(t), | |||
| petabv2.C.EXPERIMENT_ID: [reported_experiment_id] * len(t), | |||
Member
Author
There was a problem hiding this comment.
Fixed in ab21bf9. The DataFrame now applies the per-experiment padding mask (_ts_masks[ic]) consistently to the index and every column, so experiments with differing numbers of timepoints no longer cause length mismatches or leak padded/duplicated indices. Added test_petab_simulate_ragged_experiments (two experiments with 1 vs 2 timepoints) as a regression test, since no PEtab test-suite case exercises this. Also dropped the fragile sc[0] experiment-id reverse-lookup in favor of the experiment ids already produced per-experiment by run_simulations.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3207 +/- ##
===========================================
+ Coverage 41.45% 75.75% +34.30%
===========================================
Files 311 311
Lines 20638 20642 +4
Branches 1484 1483 -1
===========================================
+ Hits 8556 15638 +7082
+ Misses 12057 4992 -7065
+ Partials 25 12 -13
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
The `update_noise_dist` missing-return bug in petab.v2.petab1to2 (which silently reverted every upgraded observable's noiseDistribution to `normal`, corrupting iy_trafos/chi2 for non-linear observable transformations) has been fixed upstream in PEtab-dev/libpetab-python#502. Remove the `_fix_petab1to2_noise_distribution_bug` shim and bump the v1 petab-suite CI job's libpetab pin from 44c8062 to 1b8599dd (the #502 merge commit) so the fix is present. Case 0016 now passes via the upstream conversion; case 0007's xfail (inherent log10-normal -> log-normal substitution) and its warning filter remain. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The pristine-v1-problem snapshot (deepcopy before compilation, used for the v1->v2 upgrade instead of the possibly-mutated problem) is not needed: the full v1.0.0 jax petab test suite -- including the placeholder cases (0003/0004/0005/0009) that trigger `_workaround_observable_parameters`'s in-place mutation -- passes upgrading `petab_problem` directly, as on main. Restore the file to match main so this PR leaves it untouched. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
Claude Code Review is paused for this repository. To reconnect it, an admin of this repository's GitHub organization (or the account owner, for personal repositories) who can also manage your Claude organization's Code Review settings needs to re-link GitHub in Code Review settings. This is a one-time step.
Tip: disable this comment in your organization's Code Review settings.
…ookup Two issues in the PEtab v2 simulation-DataFrame builder: * Length inconsistency (reviewer comment): the DataFrame mixed masked/valid-only arrays (obs, t[mask], y[mask]) with unmasked lengths (`len(t)`, `index=_petab_measurement_indices[ic, :]`). Since `_get_measurements` pads every experiment to a common length, any problem whose experiments have differing numbers of timepoints would raise `ValueError: arrays must all be same length` or leak padded/duplicated indices. Apply the per-experiment `_ts_masks[ic]` mask consistently to the index and every column. * Fragile zero-indexing: the experiment id was recovered by reverse- mapping `dyn_conditions`' first condition id (`sc[0]`) through `_conditions_to_experiment_map`. `run_simulations` already builds these per experiment, so thread the experiment ids through directly and drop both the `sc[0]` lookup and the now-unused `_conditions_to_experiment_map`. No PEtab test-suite case exercises ragged (differing-timepoint) experiments, so add `test_petab_simulate_ragged_experiments` as a regression test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Member
Author
BSnelling
approved these changes
Jul 10, 2026
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.
Now independent of #3196 (rebased onto
main).To stand alone it carries two prerequisite commits that also live in #3196 (they'll be reconciled when #3196 is rebased):
Fix CI failures in JAX PEtab v1/v2 backend(adds_resolve_condition_target_value, the pristine-v1 snapshot, and the v2 simulation-df/experiment fixes) andIgnore another benign petab1to2 warning-turned-error(parameterScale-warning filter, needed for the log10 cases). Reviewers can focus on the last four commits — the rest is the actual new work.1. Correctness: remaining PEtab test-suite failures (cases 0010–0013, 0016–0020; 0007 xfail)
sim/jax/petab.py): the reinit methods read condition-table state overrides via the old wide-formatcondition_df.loc[condition, state_id], butpetab.v2'scondition_dfis long format, so every override silently fell back to the SBML default. Rewired to the structured condition changes.importers/petab/v1/_petab_import.py):petab1to2'supdate_noise_distcomputes the merged v1→v2noiseDistribution(e.g.log→log-normal) but never returns it, so upgraded observables revert tonormal, corruptingiy_trafos/chi2 (LLH is unaffected). Added_fix_petab1to2_noise_distribution_bug(with aTODOto drop once fixed upstream). Fixes case 0016.log10-normal;log-normalis substituted and reproduces the ground-truth chi2 exactly once recomputed withlog10. Inherent v1→v2 limitation → documentedpytest.xfail.2. Differentiability: condition-table target values were frozen at construction
Condition-table target values — both species initial values (
_state_reinitialisation_value) and parameter overrides (_map_experiment_model_parameter_value) — were resolved through_parameter_mappings["targets_map"], which is built once at__init__, bakingself.parameters[i]into a frozen pytree leaf. For any parameter used as an initial value, or mapped to a model parameter by the condition table (the standard pattern for condition-specific estimated parameters):update_parameters(...)had no effect on it, andeqx.filter_grad(run_simulations)routed the sensitivity intograd._parameter_mappings[...], sograd.parameters(what callers read) was 0.Fix: resolve these values live from the raw condition-table changes via
_resolve_condition_target_value, called inside the traced region (_prepare_experiments), soself.parametersis the live array. Autodiff then matches central finite differences to ~1e-10 with zero cache leak.The petab suite skips derivative checks for jax (
if jax: pass), so these were uncaught → addedtest_condition_table_initial_value_is_differentiableandtest_condition_table_parameter_override_is_differentiable.Verification
mainbase.python/tests/test_jax.pypasses locally (incl. the two new tests). Not chasing unrelated suite cases here to keep the diff reviewable.🤖 Generated with Claude Code