From db00ced249f1205334081f202b14d1d60a68f0b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 9 Jul 2026 13:19:58 +0100 Subject: [PATCH 1/9] Fix CI failures in JAX PEtab v1/v2 backend 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 --- .../example_jax_petab/ExampleJaxPEtab.ipynb | 6 +- pytest.ini | 4 ++ .../amici/importers/petab/v1/_petab_import.py | 22 +++++-- python/sdist/amici/sim/jax/petab.py | 63 ++++++++++++++++--- tests/petab_test_suite/test_petab_suite.py | 11 ++++ 5 files changed, 88 insertions(+), 18 deletions(-) diff --git a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb index b4fe05dac4..ac6d054124 100644 --- a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb +++ b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb @@ -107,7 +107,7 @@ "outputs": [], "source": [ "# # Define the simulation condition\n", - "experiment_condition = \"_petab_experiment_condition___default__\"\n", + "experiment_condition = (\"_petab_experiment_condition___default__\",)\n", "\n", "# # Access the results for the specified condition\n", "ic = results[\"dynamic_conditions\"].index(experiment_condition)\n", @@ -163,7 +163,7 @@ "import numpy as np\n", "\n", "# Define the experiment condition\n", - "experiment_condition = \"_petab_experiment_condition___default__\"\n", + "experiment_condition = (\"_petab_experiment_condition___default__\",)\n", "\n", "\n", "def plot_simulation(results):\n", @@ -389,7 +389,7 @@ "from amici.sim.jax import ReturnValue\n", "\n", "# Define the simulation condition\n", - "experiment_condition = \"_petab_experiment_condition___default__\"\n", + "experiment_condition = (\"_petab_experiment_condition___default__\",)\n", "ic = 0\n", "\n", "# Load condition-specific data\n", diff --git a/pytest.ini b/pytest.ini index 7ed33543dd..fb70d051d3 100644 --- a/pytest.ini +++ b/pytest.ini @@ -13,6 +13,10 @@ filterwarnings = once:Signature .* for does not match any known type. falling back to type probe function:UserWarning # petab ignore:Using petab.v1.Problem with PEtab2.0 is deprecated:DeprecationWarning + # petab1to2 warns-and-falls-back for v1-only noise distributions not + # supported in v2 (e.g. log10-normal -> log-normal); this is an + # intentional, documented conversion decision, not an error + ignore:Noise distribution .* is not supported in PEtab v2\. Using .* instead\.:UserWarning # amici ignore:Conservation laws for non-constant species in models with RateRules are currently not supported and will be turned off.:UserWarning ignore:Conservation laws for non-constant species in models with Species-AssignmentRules are currently not supported and will be turned off.:UserWarning diff --git a/python/sdist/amici/importers/petab/v1/_petab_import.py b/python/sdist/amici/importers/petab/v1/_petab_import.py index 5c1f8800d1..e257641610 100644 --- a/python/sdist/amici/importers/petab/v1/_petab_import.py +++ b/python/sdist/amici/importers/petab/v1/_petab_import.py @@ -91,6 +91,16 @@ def import_petab_problem( "Unsupported model type " + petab_problem.model.type_id ) + if jax: + # snapshot the pristine v1 problem for the later v1->v2 upgrade + # (below) before it gets mutated by SBML/PySB model compilation + # (e.g. `_workaround_observable_parameters` adds global SBML + # parameters in place); upgrading the mutated problem can fail v1's + # own linting inside `petab1to2`. + import copy + + pristine_petab_problem = copy.deepcopy(petab_problem) + model_name = model_name or petab_problem.model.model_id if petab_problem.model.type_id == MODEL_TYPE_PYSB and model_name is None: @@ -267,12 +277,14 @@ def import_petab_problem( f"Successfully loaded jax model {model_name} from {output_dir}." ) - # JAXProblem requires a PEtab v2 problem; upgrade the v1 problem by - # serializing it to a temporary PEtab v1 problem on disk and letting - # petab auto-upgrade it (``petab.v2.Problem.from_yaml`` upgrades v1 - # YAML files via ``petab1to2``). + # JAXProblem requires a PEtab v2 problem; upgrade the pristine v1 + # problem by serializing it to a temporary PEtab v1 problem on disk + # and letting petab auto-upgrade it (``petab.v2.Problem.from_yaml`` + # upgrades v1 YAML files via ``petab1to2``). with tempfile.TemporaryDirectory() as tmp_dir: - yaml_path = petab_problem.to_files_generic(prefix_path=tmp_dir) + yaml_path = pristine_petab_problem.to_files_generic( + prefix_path=tmp_dir + ) petab_problem_v2 = petabv2.Problem.from_yaml(yaml_path) # Create and return JAXProblem diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index bf2ce40e38..ce026a6e68 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -548,11 +548,31 @@ def pad_and_stack(output_index: int): np_indices, ) + def _resolve_condition_target_value(self, target_value): + """Resolve a condition change's target value to a number. + + A condition table target value may be a numeric literal, or a + reference to another PEtab parameter id (to be substituted with + that parameter's current value, e.g. to share an estimated + parameter's value across multiple conditions). + """ + if not target_value.is_number: + pname = str(target_value) + if pname in self.parameter_ids: + return self.parameters[self.parameter_ids.index(pname)] + _petab_param_map = { + param.id: param.nominal_value + for param in self._petab_problem.parameters + } + if pname in _petab_param_map: + return _petab_param_map[pname] + return jnp.asarray(target_value, dtype=self.model.parameters.dtype) + def _get_parameter_mappings(self) -> dict[str, ...]: targets_map = { c.id: { - ch.target_id: jnp.asarray( - ch.target_value, dtype=self.model.parameters.dtype + ch.target_id: self._resolve_condition_target_value( + ch.target_value ) for ch in c.changes } @@ -1989,9 +2009,16 @@ def add_default_experiment_names_to_v2_problem(petab_problem: petabv2.Problem): petab_problem.experiment_df is None or petab_problem.experiment_df.empty ): - condition_ids = petab_problem.condition_df[ - petabv2.C.CONDITION_ID - ].values + # read condition ids from the condition table elements, not + # `condition_df`: a condition with no changes (e.g. the just-added + # default condition, or any other no-op condition) contributes zero + # rows to the long-format `condition_df`, so its id could not be + # recovered from there. + condition_ids = [ + c.id + for table in petab_problem.condition_tables + for c in table.elements + ] condition_ids = [ c for c in condition_ids if "preequilibration" not in c ] @@ -2019,7 +2046,8 @@ def get_simulation_conditions_v2(petab_problem) -> pd.DataFrame: """Get simulation conditions from PEtab v2 measurement DataFrame. Returns: - A pandas DataFrame mapping experiment_ids to condition ids. + A pandas DataFrame mapping experiment_ids to condition ids, one row + per experiment. """ experiment_df = petab_problem.experiment_df @@ -2028,6 +2056,16 @@ def get_simulation_conditions_v2(petab_problem) -> pd.DataFrame: experiment_df[petabv2.C.TIME] != petabv2.C.TIME_PREEQUILIBRATION ] experiment_df = experiment_df.drop(columns=[petabv2.C.TIME]) + # a dynamic period may reference multiple condition ids (e.g. the + # synthetic preequilibration-indicator condition alongside the actual + # experiment condition); measurements are only ever queried by + # experiment id (see `JAXProblem._get_measurements`), so collapse to + # one row per experiment -- otherwise arrays built per condition row + # here and arrays built per experiment elsewhere (e.g. `p_array` in + # `_prepare_experiments`) end up with mismatched batch sizes. + experiment_df = experiment_df.drop_duplicates( + subset=[petabv2.C.EXPERIMENT_ID] + ) return experiment_df @@ -2035,12 +2073,17 @@ def _build_simulation_df_v2(problem, y, dyn_conditions): """Build petab simulation DataFrame of similation results from a PEtab v2 problem.""" dfs = [] for ic, sc in enumerate(dyn_conditions): + # all condition ids of a period share the same experiment id, so any + # one of them (here, the first) resolves the lookup experiment_id = _conditions_to_experiment_map( problem._petab_problem.experiment_df - )[sc] + )[sc[0]] - if experiment_id == "__default__": - experiment_id = jnp.nan + # the synthetic default experiment id is reported as NaN, but the + # original id is still needed below to query the measurement table + reported_experiment_id = ( + jnp.nan if experiment_id == "__default__" else experiment_id + ) obs = [ problem.model.observable_ids[io] @@ -2056,7 +2099,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), petabv2.C.TIME: t[problem._ts_masks[ic, :]], petabv2.C.SIMULATION: y[ic, problem._ts_masks[ic, :]], }, diff --git a/tests/petab_test_suite/test_petab_suite.py b/tests/petab_test_suite/test_petab_suite.py index 8c28a55f8e..d971109cc8 100755 --- a/tests/petab_test_suite/test_petab_suite.py +++ b/tests/petab_test_suite/test_petab_suite.py @@ -102,6 +102,17 @@ def _test_case(case, model_type, version, jax): simulation_df.rename( columns={petab.SIMULATION: petab.MEASUREMENT}, inplace=True ) + # the JAX backend simulates via the (upgraded) PEtab v2 problem and + # thus reports the v2-style `experimentId` (e.g. the synthetic + # `"__default__"`, mapped to NaN) instead of the v1-style + # `simulationConditionId` expected below. Rows correspond 1:1 to + # `problem.measurement_df` (v1->v2 upgrade preserves row order), so + # recover the original column by index alignment. + simulation_df[petab.SIMULATION_CONDITION_ID] = ( + problem.measurement_df.loc[ + simulation_df.index, petab.SIMULATION_CONDITION_ID + ].values + ) else: model = imported # import_petab_problem returns Model when jax=False solver = model.create_solver() From 4646543246a24c65f59199e1a76c3de697b6c6c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 9 Jul 2026 13:33:19 +0100 Subject: [PATCH 2/9] Ignore another benign petab1to2 warning-turned-error 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 --- pytest.ini | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pytest.ini b/pytest.ini index fb70d051d3..83959468af 100644 --- a/pytest.ini +++ b/pytest.ini @@ -17,6 +17,12 @@ filterwarnings = # supported in v2 (e.g. log10-normal -> log-normal); this is an # intentional, documented conversion decision, not an error ignore:Noise distribution .* is not supported in PEtab v2\. Using .* instead\.:UserWarning + # petab1to2 drops PEtab v1's parameterScale column when upgrading to v2 + # (v2 has no equivalent); nominal/simulation values are unaffected since + # petab v1's nominalValue is always stored in linear units regardless of + # parameterScale, so this only affects estimation-scale metadata, not + # simulation correctness + ignore:Parameter scales are not supported in PEtab v2\.:UserWarning # amici ignore:Conservation laws for non-constant species in models with RateRules are currently not supported and will be turned off.:UserWarning ignore:Conservation laws for non-constant species in models with Species-AssignmentRules are currently not supported and will be turned off.:UserWarning From 89e7b09854478959498030a73c40f67d6dc8d9e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 9 Jul 2026 15:10:52 +0100 Subject: [PATCH 3/9] Fix state reinitialisation lookup and petab1to2 noise-distribution bug 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 --- .../amici/importers/petab/v1/_petab_import.py | 57 ++++++++++++ python/sdist/amici/sim/jax/petab.py | 92 ++++--------------- tests/petab_test_suite/test_petab_suite.py | 20 ++++ 3 files changed, 96 insertions(+), 73 deletions(-) diff --git a/python/sdist/amici/importers/petab/v1/_petab_import.py b/python/sdist/amici/importers/petab/v1/_petab_import.py index e257641610..cea731cce7 100644 --- a/python/sdist/amici/importers/petab/v1/_petab_import.py +++ b/python/sdist/amici/importers/petab/v1/_petab_import.py @@ -38,6 +38,60 @@ logger = get_logger(__name__, logging.WARNING) +def _fix_petab1to2_noise_distribution_bug( + pristine_v1_problem: petab.Problem, petab_problem_v2 +) -> None: + """Work around a bug in ``petab.v2.petab1to2.v1v2_observable_df``. + + That function is supposed to merge the v1 ``observableTransformation`` + (lin/log/log10) into the v2 ``noiseDistribution`` (e.g. into + ``log-normal``), but its inner ``update_noise_dist`` helper never + returns the value it computes, so every upgraded observable silently + reverts to ``normal`` regardless of the original transformation. This + corrupts chi2 (and anything else derived from ``noiseDistribution``, + e.g. :func:`amici.sim.jax.petab.JAXProblem._get_measurements`'s + ``iy_trafos``) for any v1 problem that used a non-linear observable + transformation -- even though AMICI's own log-likelihood code + generation is unaffected, since that is derived directly from the + pristine v1 problem, not from this v2 upgrade. + + Recompute the correct v2 ``noiseDistribution`` from the pristine v1 + problem and patch it into the upgraded v2 problem's observables, in + place. + + TODO: remove once fixed upstream (``update_noise_dist`` in + https://github.com/PEtab-dev/libpetab-python/blob/main/petab/v2/petab1to2.py + is missing a ``return new_dist``). + """ + v1_obs_df = pristine_v1_problem.observable_df + if petab.C.OBSERVABLE_TRANSFORMATION not in v1_obs_df: + return + + import petab.v2 as petabv2 + + transformations = v1_obs_df[petab.C.OBSERVABLE_TRANSFORMATION].fillna( + petab.C.LIN + ) + if petab.C.NOISE_DISTRIBUTION in v1_obs_df: + distributions = v1_obs_df[petab.C.NOISE_DISTRIBUTION].fillna( + petab.C.NORMAL + ) + else: + distributions = pd.Series(petab.C.NORMAL, index=v1_obs_df.index) + + observables_by_id = {obs.id: obs for obs in petab_problem_v2.observables} + for obs_id, trans in transformations.items(): + if obs_id not in observables_by_id: + continue + dist = distributions[obs_id] + new_dist = dist if trans == petab.C.LIN else f"{trans}-{dist}" + # mirrors the (already applied, and already warned-about) v1-only + # log10-normal -> log-normal substitution in the real petab1to2 + if new_dist == "log10-normal": + new_dist = petabv2.C.LOG_NORMAL + observables_by_id[obs_id].noise_distribution = new_dist + + def import_petab_problem( petab_problem: petab.Problem, output_dir: str | Path | None = None, @@ -286,6 +340,9 @@ def import_petab_problem( prefix_path=tmp_dir ) petab_problem_v2 = petabv2.Problem.from_yaml(yaml_path) + _fix_petab1to2_noise_distribution_bug( + pristine_petab_problem, petab_problem_v2 + ) # Create and return JAXProblem logger.info(f"Successfully created JAXProblem for {model_name}.") diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index ce026a6e68..5b8137102c 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -5,7 +5,6 @@ import re import shutil from collections.abc import Callable, Iterable, Sized -from numbers import Number from pathlib import Path import diffrax @@ -1200,19 +1199,16 @@ def _state_needs_reinitialisation( if state_id in self._parameter_mappings["hybrid_map"]: return True - if state_id not in self._petab_problem.condition_df: - return False - for condition in simulation_conditions: - xval = self._petab_problem.condition_df.loc[condition, state_id] - if not (isinstance(xval, Number) and np.isnan(xval)): - return True - return False + return any( + state_id + in self._parameter_mappings["targets_map"].get(condition, {}) + for condition in simulation_conditions + ) def _state_reinitialisation_value( self, simulation_conditions: tuple[str, ...], state_id: str, - p: jt.Float[jt.Array, "np"], ) -> jt.Float[jt.Scalar, ""] | float: # noqa: F722 """ Get the reinitialisation value for a state. @@ -1224,8 +1220,6 @@ def _state_reinitialisation_value( ``state_id``) :param state_id: state id to get reinitialisation value for - :param p: - parameters for the simulation condition :return: reinitialisation value for the state """ @@ -1238,45 +1232,16 @@ def _state_reinitialisation_value( simulation_conditions[0], ) - if state_id not in self._petab_problem.condition_df: - # no reinitialisation, return dummy value - return 0.0 - - xval = None for condition in simulation_conditions: - candidate = self._petab_problem.condition_df.loc[ - condition, state_id - ] - if not (isinstance(candidate, Number) and np.isnan(candidate)): - xval = candidate - break - if xval is None: - # no reinitialisation, return dummy value - return 0.0 - if isinstance(xval, Number): - # numerical value, return as is - return xval - if xval in self.model.parameter_ids: - # model parameter, return value - return p[self.model.parameter_ids.index(xval)] - if xval in self.parameter_ids: - # estimated PEtab parameter, return unscaled value - return jax_unscale( - self.get_petab_parameter_by_id(xval), - self._petab_problem.parameter_df.loc[ - xval, petabv2.PARAMETER_SCALE - ], - ) - # only remaining option is nominal value for PEtab parameter - # that is not estimated, return nominal value - return self._petab_problem.parameter_df.loc[ - xval, petabv2.C.NOMINAL_VALUE - ] + target = self._parameter_mappings["targets_map"].get(condition, {}) + if state_id in target: + return target[state_id] + # no reinitialisation, return dummy value + return 0.0 def load_reinitialisation( self, simulation_conditions: str | tuple[str, ...], - p: jt.Float[jt.Array, "np"], ) -> tuple[jt.Bool[jt.Array, "nx"], jt.Float[jt.Array, "nx"]]: # noqa: F821 """ Load reinitialisation values and mask for the state vector for a simulation condition. @@ -1284,33 +1249,23 @@ def load_reinitialisation( :param simulation_conditions: Condition id(s) simultaneously active for the simulation condition to load reinitialisation for. - :param p: - Parameters for the simulation condition. :return: Tuple of reinitialisation masm and value for states. """ if isinstance(simulation_conditions, str): simulation_conditions = (simulation_conditions,) - if not any( - x_id in self._petab_problem.condition_df - or hasattr(self, "nn_output_ids") - and x_id in self._parameter_mappings["hybrid_map"] + needs_reinit = [ + self._state_needs_reinitialisation(simulation_conditions, x_id) for x_id in self.model.state_ids - ): + ] + if not any(needs_reinit): return jnp.array([]), jnp.array([]) - mask = jnp.array( - [ - self._state_needs_reinitialisation(simulation_conditions, x_id) - for x_id in self.model.state_ids - ] - ) + mask = jnp.array(needs_reinit) reinit_x = jnp.array( [ - self._state_reinitialisation_value( - simulation_conditions, x_id, p - ) + self._state_reinitialisation_value(simulation_conditions, x_id) for x_id in self.model.state_ids ] ) @@ -1439,18 +1394,9 @@ def _prepare_experiments( else: np_array = jnp.zeros((*self._ts_masks.shape[:2], 0)) - mask_reinit_array = jnp.stack( - [ - self.load_reinitialisation(sc, p)[0] - for sc, p in zip(conditions, p_array) - ] - ) - x_reinit_array = jnp.stack( - [ - self.load_reinitialisation(sc, p)[1] - for sc, p in zip(conditions, p_array) - ] - ) + reinit_arrays = [self.load_reinitialisation(sc) for sc in conditions] + mask_reinit_array = jnp.stack([m for m, _ in reinit_arrays]) + x_reinit_array = jnp.stack([x for _, x in reinit_arrays]) return ( p_array, mask_reinit_array, diff --git a/tests/petab_test_suite/test_petab_suite.py b/tests/petab_test_suite/test_petab_suite.py index d971109cc8..f4f87eb904 100755 --- a/tests/petab_test_suite/test_petab_suite.py +++ b/tests/petab_test_suite/test_petab_suite.py @@ -38,6 +38,26 @@ def test_case(case, model_type, version, jax): """Wrapper for _test_case for handling test outcomes""" + if jax and petabtests.test_id_str(case) == "0007": + # Case 0007 uses the (v1-only) `log10-normal` noise distribution. + # PEtab v2 has no log10-based distributions, so `petab1to2` + # substitutes `log-normal` for it (with a warning, filtered in + # pytest.ini). Since the noise formula (sigma) is carried over + # unchanged, chi2 -- computed as ((log(m) - log(y)) / sigma) ** 2 + # instead of ((log10(m) - log10(y)) / sigma) ** 2 -- differs from + # the v1-format ground-truth solution file by an inherent scaling + # factor; recomputing chi2 with log10 in place of log reproduces + # the expected value exactly. This is a genuine limitation of the + # v1->v2 upgrade, not an AMICI bug. LLH is unaffected, since + # AMICI's log-likelihood code is generated directly from the + # pristine v1 problem, not from the (lossy) v2 upgrade. + pytest.xfail( + "Case 0007 requires the v1-only `log10-normal` noise " + "distribution; PEtab v2 lacks an equivalent, so the " + "v1->v2 upgrade substitutes `log-normal`, giving a chi2 " + "value that is inherently different from (but consistent " + "with LLH matching) the v1-format ground truth." + ) try: _test_case(case, model_type, version, jax) except Exception as e: From b42de7b00810cd674d9e12232c582a9a9f82cb0b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 9 Jul 2026 19:25:07 +0100 Subject: [PATCH 4/9] Resolve JAX PEtab reinitialisation values live so gradients flow 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 --- python/sdist/amici/sim/jax/petab.py | 58 ++++++++++++++++++++++++---- python/tests/test_jax.py | 59 +++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 8 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 5b8137102c..8bcc0594b6 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -1199,12 +1199,47 @@ def _state_needs_reinitialisation( if state_id in self._parameter_mappings["hybrid_map"]: return True - return any( - state_id - in self._parameter_mappings["targets_map"].get(condition, {}) - for condition in simulation_conditions + return ( + self._condition_reinit_target_value( + simulation_conditions, state_id + ) + is not None ) + def _condition_reinit_target_value( + self, simulation_conditions: tuple[str, ...], state_id: str + ): + """Return the *raw* condition-table target value initialising a state. + + Looks up the (unresolved) ``target_value`` that (re)initialises + ``state_id`` for the given simultaneously-active conditions, reading + it straight from the condition-table changes. Returns ``None`` if no + active condition sets ``state_id`` (PEtab v2 requires the targets of + simultaneously-active conditions to be disjoint, so at most one does). + + The raw value is returned deliberately -- callers resolve it *live* + against :attr:`parameters` (see + :meth:`_state_reinitialisation_value`), rather than reusing the + ``targets_map`` value cached at construction time, so that gradients + w.r.t. parameters used as initial values are not silently dropped. + """ + for condition in simulation_conditions: + for c in self._petab_problem.conditions: + if c.id != condition: + continue + for change in c.changes: + if change.target_id != state_id: + continue + # NaN targets (e.g. "use the preequilibration/SBML value") + # are dropped during v1->v2 conversion, but guard anyway + if ( + change.target_value.is_number + and change.target_value.is_finite is False + ): + return None + return change.target_value + return None + def _state_reinitialisation_value( self, simulation_conditions: tuple[str, ...], @@ -1232,10 +1267,17 @@ def _state_reinitialisation_value( simulation_conditions[0], ) - for condition in simulation_conditions: - target = self._parameter_mappings["targets_map"].get(condition, {}) - if state_id in target: - return target[state_id] + target_value = self._condition_reinit_target_value( + simulation_conditions, state_id + ) + if target_value is not None: + # Resolve *live* against ``self.parameters`` (not via the + # construction-time ``targets_map`` cache): this method runs inside + # the traced/differentiated region (via ``_prepare_experiments``), + # so reading ``self.parameters`` here keeps the reinitialisation + # value a function of the current parameters -- gradients flow and + # re-simulating after ``update_parameters`` reflects the new value. + return self._resolve_condition_target_value(target_value) # no reinitialisation, return dummy value return 0.0 diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index a8efc0b38d..953af24e51 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -326,6 +326,65 @@ def test_serialisation(lotka_volterra): # noqa: F811 ) +@skip_on_valgrind +def test_condition_table_initial_value_is_differentiable(tmp_path): + """A parameter used as a species initial value via the condition table + must stay a live function of ``JAXProblem.parameters``. + + Regression test for a bug where the reinitialisation value was resolved + through the ``targets_map`` cached at construction time + (see ``JAXProblem._state_reinitialisation_value``): the initial value was + then frozen at the nominal parameter value, so ``update_parameters`` had + no effect on it and its gradient silently leaked into the cache instead of + ``grad.parameters``. + """ + import equinox as eqx + import petab.v1 as petab + from petab.v1.models.sbml_model import SbmlModel + + problem = petab.Problem() + problem.model = SbmlModel.from_antimony( + "compartment_ = 1;\n" + "species A in compartment_, B in compartment_;\n" + "A = 3; B = 0;\n" + "k1 = 0.8; k2 = 0.6;\n" + "fwd: A -> B; k1 * A;\n" + "rev: B -> A; k2 * B;\n" + ) + # `a0` initialises species `A` via the condition table (not in the model) + problem.add_parameter( + "a0", estimate=True, nominal_value=2.0, scale="lin", lb=0.1, ub=10 + ) + problem.add_observable("obs_a", "A", noise_formula="0.5") + problem.add_condition("c0", A="a0") + problem.add_measurement("obs_a", "c0", 0.0, 0.7) + problem.add_measurement("obs_a", "c0", 10.0, 0.1) + + jax_problem = import_petab_problem( + problem, jax=True, output_dir=str(tmp_path) + ) + ia = jax_problem.parameter_ids.index("a0") + + def llh(p): + return run_simulations(jax_problem.update_parameters(p))[0] + + p0 = jax_problem.parameters + # `a0` enters the likelihood only through A(0); updating it must change llh + assert abs(float(llh(p0.at[ia].add(1.0))) - float(llh(p0))) > 1e-6, ( + "initial value is frozen w.r.t. update_parameters" + ) + + # autodiff w.r.t. `a0` (read from grad.parameters) must match finite diff + eps = 1e-6 + fd = (float(llh(p0.at[ia].add(eps))) - float(llh(p0.at[ia].add(-eps)))) / ( + 2 * eps + ) + grad = eqx.filter_grad(lambda m: run_simulations(m)[0])( + jax_problem.update_parameters(p0) + ) + assert_allclose(float(grad.parameters[ia]), fd, rtol=1e-4, atol=1e-4) + + @skip_on_valgrind def test_steady_state_event_no_recompile_across_conditions( tmp_path, monkeypatch From 5411a21ff41dc9fdd80f283a0f40158b9556c679 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 9 Jul 2026 19:35:16 +0100 Subject: [PATCH 5/9] Resolve condition-table parameter overrides live too 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 --- python/sdist/amici/sim/jax/petab.py | 25 ++++++++----- python/tests/test_jax.py | 55 +++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 9 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 8bcc0594b6..e84bed24f3 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -1126,18 +1126,25 @@ def _map_experiment_model_parameter_value( else: init_val = self.model.parameters[p_index] - targets_filtered = { - param: value - for condition, target in self._parameter_mappings[ - "targets_map" - ].items() - for param, value in target.items() - if condition in condition_ids + # Resolve condition-table overrides *live* from the raw changes rather + # than reusing the ``targets_map`` values cached at construction time. + # This method runs inside the traced/differentiated region (via + # ``_prepare_experiments``), so reading ``self.parameters`` (through + # ``_resolve_condition_target_value``) keeps the value a function of + # the current parameters -- gradients w.r.t. an (e.g. condition- + # specific) estimated parameter that a condition maps this one to flow, + # and re-simulating after ``update_parameters`` reflects the new value. + raw_targets = { + change.target_id: change.target_value + for c in self._petab_problem.conditions + if c.id in condition_ids + for change in c.changes } - if pname in targets_filtered: + if pname in raw_targets: return jnp.asarray( - targets_filtered[pname], dtype=self.model.parameters.dtype + self._resolve_condition_target_value(raw_targets[pname]), + dtype=self.model.parameters.dtype, ) elif pname in self._parameter_mappings["hybrid_map"]: return jnp.asarray( diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 953af24e51..31815bd713 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -385,6 +385,61 @@ def llh(p): assert_allclose(float(grad.parameters[ia]), fd, rtol=1e-4, atol=1e-4) +@skip_on_valgrind +def test_condition_table_parameter_override_is_differentiable(tmp_path): + """A model parameter mapped to an estimated parameter via the condition + table (the standard PEtab pattern for condition-specific estimated + parameters) must stay a live function of ``JAXProblem.parameters``. + + Regression test for the same construction-time freezing bug as + ``test_condition_table_initial_value_is_differentiable``, on the parameter + mapping path (``JAXProblem._map_experiment_model_parameter_value``). + """ + import equinox as eqx + import petab.v1 as petab + from petab.v1.models.sbml_model import SbmlModel + + problem = petab.Problem() + problem.model = SbmlModel.from_antimony( + "compartment_ = 1;\n" + "species A in compartment_, B in compartment_;\n" + "A = 1; B = 0;\n" + "k1 = 0.8; k2 = 0.6;\n" + "fwd: A -> B; k1 * A;\n" + "rev: B -> A; k2 * B;\n" + ) + # the condition maps model parameter `k1` to the estimated `k1_c0` + problem.add_parameter( + "k1_c0", estimate=True, nominal_value=0.8, scale="lin", lb=0.1, ub=10 + ) + problem.add_observable("obs_b", "B", noise_formula="0.5") + problem.add_condition("c0", k1="k1_c0") + problem.add_measurement("obs_b", "c0", 1.0, 0.3) + problem.add_measurement("obs_b", "c0", 5.0, 0.4) + + jax_problem = import_petab_problem( + problem, jax=True, output_dir=str(tmp_path) + ) + ik = jax_problem.parameter_ids.index("k1_c0") + + def llh(p): + return run_simulations(jax_problem.update_parameters(p))[0] + + p0 = jax_problem.parameters + assert abs(float(llh(p0.at[ik].add(0.3))) - float(llh(p0))) > 1e-6, ( + "condition-overridden parameter is frozen w.r.t. update_parameters" + ) + + eps = 1e-6 + fd = (float(llh(p0.at[ik].add(eps))) - float(llh(p0.at[ik].add(-eps)))) / ( + 2 * eps + ) + grad = eqx.filter_grad(lambda m: run_simulations(m)[0])( + jax_problem.update_parameters(p0) + ) + assert_allclose(float(grad.parameters[ik]), fd, rtol=1e-4, atol=1e-4) + + @skip_on_valgrind def test_steady_state_event_no_recompile_across_conditions( tmp_path, monkeypatch From 826d605efd8d7c150b5b81b40992108665e1c5a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 9 Jul 2026 19:45:32 +0100 Subject: [PATCH 6/9] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- python/sdist/amici/sim/jax/petab.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index e84bed24f3..75b16d0477 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -1299,7 +1299,7 @@ def load_reinitialisation( Condition id(s) simultaneously active for the simulation condition to load reinitialisation for. :return: - Tuple of reinitialisation masm and value for states. + Tuple of reinitialisation mask and value for states. """ if isinstance(simulation_conditions, str): simulation_conditions = (simulation_conditions,) @@ -1308,8 +1308,7 @@ def load_reinitialisation( self._state_needs_reinitialisation(simulation_conditions, x_id) for x_id in self.model.state_ids ] - if not any(needs_reinit): - return jnp.array([]), jnp.array([]) + # Always return full-length arrays per condition; callers stack/vmap across conditions and require consistent shapes. mask = jnp.array(needs_reinit) reinit_x = jnp.array( From bf387b3187d14cc1b6fe84e4a55c1337d0176491 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Thu, 9 Jul 2026 21:58:06 +0100 Subject: [PATCH 7/9] Drop petab1to2 noise-distribution workaround; rely on upstream fix 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 https://github.com/PEtab-dev/libpetab-python/pull/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 --- .github/workflows/test_petab_test_suite.yml | 2 +- .../amici/importers/petab/v1/_petab_import.py | 57 ------------------- 2 files changed, 1 insertion(+), 58 deletions(-) diff --git a/.github/workflows/test_petab_test_suite.yml b/.github/workflows/test_petab_test_suite.yml index 644ae8624c..71ed7cbaff 100644 --- a/.github/workflows/test_petab_test_suite.yml +++ b/.github/workflows/test_petab_test_suite.yml @@ -87,7 +87,7 @@ jobs: run: | source ./venv/bin/activate \ && python3 -m pip uninstall -y petab \ - && python3 -m pip install git+https://github.com/petab-dev/libpetab-python.git@44c8062ce1b87a74a0ba1bd2551de0cdc2a13ff1 \ + && python3 -m pip install git+https://github.com/petab-dev/libpetab-python.git@1b8599dd1eb9bda74853255b4cc4baf75b7e4d63 \ && python3 -m pip install git+https://github.com/pysb/pysb@master \ && python3 -m pip install sympy>=1.12.1 diff --git a/python/sdist/amici/importers/petab/v1/_petab_import.py b/python/sdist/amici/importers/petab/v1/_petab_import.py index cea731cce7..e257641610 100644 --- a/python/sdist/amici/importers/petab/v1/_petab_import.py +++ b/python/sdist/amici/importers/petab/v1/_petab_import.py @@ -38,60 +38,6 @@ logger = get_logger(__name__, logging.WARNING) -def _fix_petab1to2_noise_distribution_bug( - pristine_v1_problem: petab.Problem, petab_problem_v2 -) -> None: - """Work around a bug in ``petab.v2.petab1to2.v1v2_observable_df``. - - That function is supposed to merge the v1 ``observableTransformation`` - (lin/log/log10) into the v2 ``noiseDistribution`` (e.g. into - ``log-normal``), but its inner ``update_noise_dist`` helper never - returns the value it computes, so every upgraded observable silently - reverts to ``normal`` regardless of the original transformation. This - corrupts chi2 (and anything else derived from ``noiseDistribution``, - e.g. :func:`amici.sim.jax.petab.JAXProblem._get_measurements`'s - ``iy_trafos``) for any v1 problem that used a non-linear observable - transformation -- even though AMICI's own log-likelihood code - generation is unaffected, since that is derived directly from the - pristine v1 problem, not from this v2 upgrade. - - Recompute the correct v2 ``noiseDistribution`` from the pristine v1 - problem and patch it into the upgraded v2 problem's observables, in - place. - - TODO: remove once fixed upstream (``update_noise_dist`` in - https://github.com/PEtab-dev/libpetab-python/blob/main/petab/v2/petab1to2.py - is missing a ``return new_dist``). - """ - v1_obs_df = pristine_v1_problem.observable_df - if petab.C.OBSERVABLE_TRANSFORMATION not in v1_obs_df: - return - - import petab.v2 as petabv2 - - transformations = v1_obs_df[petab.C.OBSERVABLE_TRANSFORMATION].fillna( - petab.C.LIN - ) - if petab.C.NOISE_DISTRIBUTION in v1_obs_df: - distributions = v1_obs_df[petab.C.NOISE_DISTRIBUTION].fillna( - petab.C.NORMAL - ) - else: - distributions = pd.Series(petab.C.NORMAL, index=v1_obs_df.index) - - observables_by_id = {obs.id: obs for obs in petab_problem_v2.observables} - for obs_id, trans in transformations.items(): - if obs_id not in observables_by_id: - continue - dist = distributions[obs_id] - new_dist = dist if trans == petab.C.LIN else f"{trans}-{dist}" - # mirrors the (already applied, and already warned-about) v1-only - # log10-normal -> log-normal substitution in the real petab1to2 - if new_dist == "log10-normal": - new_dist = petabv2.C.LOG_NORMAL - observables_by_id[obs_id].noise_distribution = new_dist - - def import_petab_problem( petab_problem: petab.Problem, output_dir: str | Path | None = None, @@ -340,9 +286,6 @@ def import_petab_problem( prefix_path=tmp_dir ) petab_problem_v2 = petabv2.Problem.from_yaml(yaml_path) - _fix_petab1to2_noise_distribution_bug( - pristine_petab_problem, petab_problem_v2 - ) # Create and return JAXProblem logger.info(f"Successfully created JAXProblem for {model_name}.") From 72bd865d884f414f0b607e0eebf942607b9c7863 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Fri, 10 Jul 2026 00:13:48 +0100 Subject: [PATCH 8/9] Revert unnecessary _petab_import.py changes 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 --- .../amici/importers/petab/v1/_petab_import.py | 22 +++++-------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/python/sdist/amici/importers/petab/v1/_petab_import.py b/python/sdist/amici/importers/petab/v1/_petab_import.py index e257641610..5c1f8800d1 100644 --- a/python/sdist/amici/importers/petab/v1/_petab_import.py +++ b/python/sdist/amici/importers/petab/v1/_petab_import.py @@ -91,16 +91,6 @@ def import_petab_problem( "Unsupported model type " + petab_problem.model.type_id ) - if jax: - # snapshot the pristine v1 problem for the later v1->v2 upgrade - # (below) before it gets mutated by SBML/PySB model compilation - # (e.g. `_workaround_observable_parameters` adds global SBML - # parameters in place); upgrading the mutated problem can fail v1's - # own linting inside `petab1to2`. - import copy - - pristine_petab_problem = copy.deepcopy(petab_problem) - model_name = model_name or petab_problem.model.model_id if petab_problem.model.type_id == MODEL_TYPE_PYSB and model_name is None: @@ -277,14 +267,12 @@ def import_petab_problem( f"Successfully loaded jax model {model_name} from {output_dir}." ) - # JAXProblem requires a PEtab v2 problem; upgrade the pristine v1 - # problem by serializing it to a temporary PEtab v1 problem on disk - # and letting petab auto-upgrade it (``petab.v2.Problem.from_yaml`` - # upgrades v1 YAML files via ``petab1to2``). + # JAXProblem requires a PEtab v2 problem; upgrade the v1 problem by + # serializing it to a temporary PEtab v1 problem on disk and letting + # petab auto-upgrade it (``petab.v2.Problem.from_yaml`` upgrades v1 + # YAML files via ``petab1to2``). with tempfile.TemporaryDirectory() as tmp_dir: - yaml_path = pristine_petab_problem.to_files_generic( - prefix_path=tmp_dir - ) + yaml_path = petab_problem.to_files_generic(prefix_path=tmp_dir) petab_problem_v2 = petabv2.Problem.from_yaml(yaml_path) # Create and return JAXProblem From ab21bf94266f0f112d5d58fc93453aae3f7f4367 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Fri, 10 Jul 2026 01:14:47 +0100 Subject: [PATCH 9/9] Fix _build_simulation_df_v2: consistent masking, drop fragile sc[0] lookup 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 --- python/sdist/amici/sim/jax/petab.py | 60 ++++++++++++++--------------- python/tests/test_jax.py | 53 +++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 32 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 75b16d0477..3f481d4a2d 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -1850,6 +1850,10 @@ def run_simulations( ] conditions = { "dynamic_conditions": dynamic_conditions, + # experiment ids aligned with `dynamic_conditions` and the rows of + # `_iys`/`_ts_masks`, so result-building need not reverse-map a + # condition id back to its experiment + "experiment_ids": [exp.id for exp in experiments], } has_preeq = any(exp.periods[0].is_preequilibration for exp in experiments) @@ -1931,7 +1935,7 @@ def petab_simulate( ret=ReturnValue.y, ) if isinstance(problem._petab_problem, petabv2.Problem): - return _build_simulation_df_v2(problem, y, r["dynamic_conditions"]) + return _build_simulation_df_v2(problem, y, r["experiment_ids"]) else: dfs = [] for ic, sc in enumerate(r["dynamic_conditions"]): @@ -2063,41 +2067,42 @@ def get_simulation_conditions_v2(petab_problem) -> pd.DataFrame: return experiment_df -def _build_simulation_df_v2(problem, y, dyn_conditions): - """Build petab simulation DataFrame of similation results from a PEtab v2 problem.""" - dfs = [] - for ic, sc in enumerate(dyn_conditions): - # all condition ids of a period share the same experiment id, so any - # one of them (here, the first) resolves the lookup - experiment_id = _conditions_to_experiment_map( - problem._petab_problem.experiment_df - )[sc[0]] +def _build_simulation_df_v2(problem, y, experiment_ids): + """Build a PEtab simulation DataFrame from PEtab v2 simulation results. + ``experiment_ids`` is aligned with the rows of ``y`` / + ``problem._iys`` / ``problem._ts_masks`` (one entry per simulated + experiment). + """ + dfs = [] + for ic, experiment_id in enumerate(experiment_ids): # the synthetic default experiment id is reported as NaN, but the # original id is still needed below to query the measurement table reported_experiment_id = ( jnp.nan if experiment_id == "__default__" else experiment_id ) + # `_get_measurements` pads every experiment's arrays to a common + # length; apply the per-experiment mask consistently to the index + # and every column so experiments with fewer timepoints don't cause + # length mismatches or leak padded/duplicated measurement indices. + mask = problem._ts_masks[ic, :] obs = [ - problem.model.observable_ids[io] - for io in problem._iys[ic, problem._ts_masks[ic, :]] + problem.model.observable_ids[io] for io in problem._iys[ic, mask] ] - t = jnp.concat( - ( - problem._ts_dyn[ic, :], - problem._ts_posteq[ic, :], - ) - ) + t = jnp.concat((problem._ts_dyn[ic, :], problem._ts_posteq[ic, :]))[ + mask + ] + n = len(t) df_sc = pd.DataFrame( { - petabv2.C.MODEL_ID: [float("nan")] * len(t), + petabv2.C.MODEL_ID: [float("nan")] * n, petabv2.C.OBSERVABLE_ID: obs, - petabv2.C.EXPERIMENT_ID: [reported_experiment_id] * len(t), - petabv2.C.TIME: t[problem._ts_masks[ic, :]], - petabv2.C.SIMULATION: y[ic, problem._ts_masks[ic, :]], + petabv2.C.EXPERIMENT_ID: [reported_experiment_id] * n, + petabv2.C.TIME: t, + petabv2.C.SIMULATION: y[ic, mask], }, - index=problem._petab_measurement_indices[ic, :], + index=problem._petab_measurement_indices[ic, mask], ) if ( petabv2.C.OBSERVABLE_PARAMETERS @@ -2118,15 +2123,6 @@ def _build_simulation_df_v2(problem, y, dyn_conditions): return pd.concat(dfs).sort_index() -def _conditions_to_experiment_map( - experiment_df: pd.DataFrame, -) -> dict[str, str]: - condition_to_experiment = { - row.conditionId: row.experimentId for row in experiment_df.itertuples() - } - return condition_to_experiment - - def _parse_model_entity_id( model_entity_id: str, nn: dict ) -> list[tuple[str, str]]: diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 31815bd713..66af120e98 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -440,6 +440,59 @@ def llh(p): assert_allclose(float(grad.parameters[ik]), fd, rtol=1e-4, atol=1e-4) +@skip_on_valgrind +def test_petab_simulate_ragged_experiments(tmp_path): + """``petab_simulate`` must handle experiments with different numbers of + measurement timepoints. + + Regression test for ``_build_simulation_df_v2``: ``_get_measurements`` + pads every experiment's arrays to a common length, so an experiment with + fewer timepoints has masked-out padding. If the padding mask is not + applied consistently to the index and all columns, building the + simulation DataFrame raises ``ValueError: arrays must all be same + length`` (or leaks padded/duplicated indices). + """ + import petab.v1 as petab + from amici.sim.jax import petab_simulate + from petab.v1.models.sbml_model import SbmlModel + + problem = petab.Problem() + problem.model = SbmlModel.from_antimony( + "compartment_ = 1;\n" + "species A in compartment_, B in compartment_;\n" + "A = 1; B = 0;\n" + "k1 = 0.8; k2 = 0.6;\n" + "fwd: A -> B; k1 * A;\n" + "rev: B -> A; k2 * B;\n" + ) + # `k1` is set per condition below, so only the free `k2` goes in the + # parameter table + problem.add_parameter( + "k2", estimate=False, nominal_value=0.6, scale="lin", lb=0.1, ub=10 + ) + problem.add_observable("obs_b", "B", noise_formula="0.5") + # two conditions -> two experiments, with DIFFERENT numbers of + # timepoints so `_ts_masks` has genuine padding + problem.add_condition("c0", k1=0.8) + problem.add_condition("c1", k1=0.5) + problem.add_measurement("obs_b", "c0", 1.0, 0.3) + problem.add_measurement("obs_b", "c1", 1.0, 0.4) + problem.add_measurement("obs_b", "c1", 5.0, 0.2) + + jax_problem = import_petab_problem( + problem, jax=True, output_dir=str(tmp_path) + ) + + sim_df = petab_simulate(jax_problem) + + # exactly one simulated row per measurement (1 for c0, 2 for c1) -- no + # length mismatch, no padded/duplicated rows, no missing simulations + assert len(sim_df) == len(problem.measurement_df) == 3 + assert sim_df.index.is_unique + assert sorted(sim_df[petab.TIME].tolist()) == [1.0, 1.0, 5.0] + assert not sim_df[petab.SIMULATION].isna().any() + + @skip_on_valgrind def test_steady_state_event_no_recompile_across_conditions( tmp_path, monkeypatch