Fix JAX PEtab condition-table resolution: state reinit, noise distribution, and differentiability#3206
Closed
FFroehlich wants to merge 6 commits into
Closed
Fix JAX PEtab condition-table resolution: state reinit, noise distribution, and differentiability#3206FFroehlich wants to merge 6 commits into
FFroehlich wants to merge 6 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR is a follow-up on top of #3196 that targets the remaining PEtab test suite failures for the JAX backend by (1) fixing condition/state reinitialisation lookup under PEtab v2’s long-format condition tables, (2) working around an upstream petab1to2 noise-distribution regression affecting chi2, and (3) explicitly documenting/xfailing an inherent PEtab v1→v2 limitation for test case 0007.
Changes:
- Rewired state reinitialisation lookup to use
_parameter_mappings["targets_map"]rather than wide-formatcondition_df.loc[...]access in the JAX PEtab adapter. - Added a
petab1to2noise-distribution workaround that patches upgraded v2 observables using the pristine v1 problem to preserve chi2 consistency. - Added an explicit
pytest.xfailfor JAX on PEtab test-suite case0007(log10-normal unsupported in v2).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| tests/petab_test_suite/test_petab_suite.py | Adds an explicit xfail for JAX on case 0007, documenting the v1→v2 log10-normal limitation. |
| python/sdist/amici/sim/jax/petab.py | Fixes reinitialisation lookup to work with PEtab v2 long-format condition changes and refactors experiment preparation around that. |
| python/sdist/amici/importers/petab/v1/_petab_import.py | Adds a workaround to correct noiseDistribution after v1→v2 upgrade to avoid chi2 corruption. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Member
Author
|
@dweindl looks like @claude suggest here that there is a bug in the v1v2 upgrade (see changes in python/sdist/amici/importers/petab/v1/_petab_import.py). Could you please confirm this is real? To me it looks real given that |
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>
FFroehlich
force-pushed
the
fix/jax-petab-reinit-noise-dist
branch
from
July 9, 2026 19:11
1044127 to
826d605
Compare
Member
Yes, unfortunately. PEtab-dev/libpetab-python#502 |
FFroehlich
changed the title
Member
Author
|
Superseded by #3207 (recreated to trigger CI, which didn't run here after the rebase/force-push). Same branch, identical history and diff. |
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