From ae9e62a49abb367c0848d0de5b6afcf008e9f49c Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 13:33:35 +0000 Subject: [PATCH 01/22] Chain PEtab v2 experiment periods natively in the JAX simulator 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 Claude-Session: https://claude.ai/code/session_0173CHAQAGCsTtnyNibrFjDv --- .../amici/importers/petab/_petab_importer.py | 48 +- python/sdist/amici/sim/jax/_simulation.py | 6 +- python/sdist/amici/sim/jax/model.py | 338 +++++--- python/sdist/amici/sim/jax/petab.py | 800 +++++++++++++----- .../tests/petab_/test_petab_v2_multiperiod.py | 227 +++++ tests/performance/test_jax_regression.py | 31 +- tests/sbml/testSBMLSuiteJax.py | 19 +- 7 files changed, 1102 insertions(+), 367 deletions(-) create mode 100644 python/tests/petab_/test_petab_v2_multiperiod.py diff --git a/python/sdist/amici/importers/petab/_petab_importer.py b/python/sdist/amici/importers/petab/_petab_importer.py index bc0c991962..94794346fd 100644 --- a/python/sdist/amici/importers/petab/_petab_importer.py +++ b/python/sdist/amici/importers/petab/_petab_importer.py @@ -212,13 +212,12 @@ def _preprocess_sbml(self): if not isinstance(self.petab_problem.model, SbmlModel): raise ValueError("The PEtab problem must contain an SBML model.") - # Convert petab experiments to events, because so far, - # AMICI only supports preequilibration/presimulation/simulation, but - # no arbitrary list of periods. - exp_event_conv = ExperimentsToSbmlConverter(self.petab_problem) - # This will always create a copy of the problem. if self._jax: - self._unconverted_problem = exp_event_conv._original_problem + # The JAX backend natively chains one ODE integration per + # experiment period (see amici.sim.jax.petab), so there is no + # need to convert experiments with more than two periods into + # SBML events. The condition table is left untouched. + self._unconverted_problem = None condition_targets = { change.target_id for condition in self.petab_problem.conditions @@ -232,13 +231,19 @@ def _preprocess_sbml(self): "The JAX backend does not currently support PEtab problems where network " "parameters appear in the conditions table. " ) - self.petab_problem = exp_event_conv.convert() - for experiment in self.petab_problem.experiments: - if len(experiment.periods) > 2: - # This should never happen due to the conversion above - raise NotImplementedError( - "AMICI currently does not support more than two periods." - ) + else: + # Convert petab experiments to events, because so far, the + # sundials backend only supports preequilibration/presimulation/ + # simulation, but no arbitrary list of periods. + exp_event_conv = ExperimentsToSbmlConverter(self.petab_problem) + # This will always create a copy of the problem. + self.petab_problem = exp_event_conv.convert() + for experiment in self.petab_problem.experiments: + if len(experiment.periods) > 2: + # This should never happen due to the conversion above + raise NotImplementedError( + "AMICI currently does not support more than two periods." + ) if self._debug: print("PetabImpoter._preprocess_sbml: petab_problem:") @@ -333,15 +338,22 @@ def _do_import_sbml(self): output_parameter_defaults=self._output_parameter_defaults, ) - # All indicator variables, i.e., all remaining targets after - # experiments-to-event in the PEtab problem must be converted - # to fixed parameters + # All condition-table targets that are not estimated must be + # converted to fixed parameters. For the sundials backend, these are + # only ever the indicator variables introduced by the + # experiments-to-event conversion above. For the JAX backend, which + # keeps the original condition table, this may also contain state + # (species) targets, which must NOT be treated as fixed parameters + # since they are handled via state reinitialisation instead. fixed_parameters = { change.target_id for experiment in self.petab_problem.experiments for period in experiment.periods for condition_id in period.condition_ids for change in self.petab_problem[condition_id].changes + if not self.petab_problem.model.is_state_variable( + change.target_id + ) } from .v1._sbml_import import show_model_info @@ -771,8 +783,6 @@ def create_simulator( Whether to force re-import even if the model module already exists. :return: The created PEtab simulator. """ - from amici.sim.sundials.petab import ExperimentManager, PetabSimulator - if self._jax: model_module = self.import_module(force_import=force_import) model = model_module.Model() @@ -787,6 +797,8 @@ def create_simulator( ), ) + from amici.sim.sundials.petab import ExperimentManager, PetabSimulator + model = self.import_module(force_import=force_import).get_model() em = ExperimentManager(model=model, petab_problem=self.petab_problem) return PetabSimulator(em=em) diff --git a/python/sdist/amici/sim/jax/_simulation.py b/python/sdist/amici/sim/jax/_simulation.py index 95ef034ace..c0db26eef5 100644 --- a/python/sdist/amici/sim/jax/_simulation.py +++ b/python/sdist/amici/sim/jax/_simulation.py @@ -250,7 +250,11 @@ def solve( term, dict(**STARTING_STATS), ) - return sol.ys, jnp.repeat(h[None, :], sol.ys.shape[0]), stats + return ( + sol.ys, + jnp.repeat(h[None, :], sol.ys.shape[0], axis=0), + stats, + ) def cond_fn(carry): _, t_start, y0, _, _, stats = carry diff --git a/python/sdist/amici/sim/jax/model.py b/python/sdist/amici/sim/jax/model.py index 6eb213b252..8c48d7cf40 100644 --- a/python/sdist/amici/sim/jax/model.py +++ b/python/sdist/amici/sim/jax/model.py @@ -439,7 +439,9 @@ def _initialise_heaviside_variables( ) def _x_rdatas( - self, x: jt.Float[jt.Array, "nt nxs"], tcl: jt.Float[jt.Array, "ncl"] + self, + x: jt.Float[jt.Array, "nt nxs"], + tcl: jt.Float[jt.Array, "nt ncl"], ) -> jt.Float[jt.Array, "nt nx"]: """ Compute the full state vector from the reduced state vector and conservation laws. @@ -447,18 +449,19 @@ def _x_rdatas( :param x: reduced state vector :param tcl: - total values for conservation laws + total values for conservation laws, per time point (conservation + laws may change across period boundaries as parameters change) :return: full state vector """ - return jax.vmap(self._x_rdata, in_axes=(0, None))(x, tcl) + return jax.vmap(self._x_rdata, in_axes=(0, 0))(x, tcl) def _nllhs( self, ts: jt.Float[jt.Array, "nt nx"], xs: jt.Float[jt.Array, "nt nxs"], - p: jt.Float[jt.Array, "np"], - tcl: jt.Float[jt.Array, "ncl"], + p: jt.Float[jt.Array, "nt np"], + tcl: jt.Float[jt.Array, "nt ncl"], hs: jt.Float[jt.Array, "nt ne"], mys: jt.Float[jt.Array, "nt"], iys: jt.Int[jt.Array, "nt"], @@ -473,9 +476,10 @@ def _nllhs( :param xs: state vectors :param p: - parameters + parameters, per time point (parameters may differ across period + boundaries) :param tcl: - total values for conservation laws + total values for conservation laws, per time point :param h: heaviside variables :param mys: @@ -489,7 +493,7 @@ def _nllhs( :return: negative log-likelihoods of the observables """ - return jax.vmap(self._nllh, in_axes=(0, 0, None, None, 0, 0, 0, 0, 0))( + return jax.vmap(self._nllh, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, 0))( ts, xs, p, tcl, hs, mys, iys, ops, nps ) @@ -497,8 +501,8 @@ def _ys( self, ts: jt.Float[jt.Array, "nt"], xs: jt.Float[jt.Array, "nt nxs"], - p: jt.Float[jt.Array, "np"], - tcl: jt.Float[jt.Array, "ncl"], + p: jt.Float[jt.Array, "nt np"], + tcl: jt.Float[jt.Array, "nt ncl"], hs: jt.Float[jt.Array, "nt ne"], iys: jt.Float[jt.Array, "nt"], ops: jt.Float[jt.Array, "nt *nop"], @@ -511,9 +515,10 @@ def _ys( :param xs: state vectors :param p: - parameters + parameters, per time point (parameters may differ across period + boundaries) :param tcl: - total values for conservation laws + total values for conservation laws, per time point :param h: heaviside variables :param iys: @@ -527,15 +532,15 @@ def _ys( lambda t, x, p, tcl, h, iy, op: ( self._y(t, x, p, tcl, h, op).at[iy].get() ), - in_axes=(0, 0, None, None, 0, 0, 0), + in_axes=(0, 0, 0, 0, 0, 0, 0), )(ts, xs, p, tcl, hs, iys, ops) def _sigmays( self, ts: jt.Float[jt.Array, "nt"], xs: jt.Float[jt.Array, "nt nxs"], - p: jt.Float[jt.Array, "np"], - tcl: jt.Float[jt.Array, "ncl"], + p: jt.Float[jt.Array, "nt np"], + tcl: jt.Float[jt.Array, "nt ncl"], hs: jt.Float[jt.Array, "nt ne"], iys: jt.Int[jt.Array, "nt"], ops: jt.Float[jt.Array, "nt *nop"], @@ -549,9 +554,10 @@ def _sigmays( :param xs: state vectors :param p: - parameters + parameters, per time point (parameters may differ across period + boundaries) :param tcl: - total values for conservation laws + total values for conservation laws, per time point :param h: heaviside variables :param iys: @@ -567,19 +573,20 @@ def _sigmays( lambda t, x, p, tcl, h, iy, op, np: ( self._sigmay(self._y(t, x, p, tcl, h, op), p, np).at[iy].get() ), - in_axes=(0, 0, None, None, 0, 0, 0, 0), + in_axes=(0, 0, 0, 0, 0, 0, 0, 0), )(ts, xs, p, tcl, hs, iys, ops, nps) - def simulate_condition_unjitted( + def _simulate_period( self, - p: jt.Float[jt.Array, "np"] | None, + p: jt.Float[jt.Array, "np"], + t0: jnp.float_, ts_dyn: jt.Float[jt.Array, "nt_dyn"], ts_posteq: jt.Float[jt.Array, "nt_posteq"], - my: jt.Float[jt.Array, "nt"], - iys: jt.Int[jt.Array, "nt"], - iy_trafos: jt.Int[jt.Array, "nt"], - ops: jt.Float[jt.Array, "nt *nop"], - nps: jt.Float[jt.Array, "nt *nnp"], + tcl: jt.Float[jt.Array, "ncl"], + h: jt.Bool[jt.Array, "ne"], + h_mask: jt.Bool[jt.Array, "ne"], + x_solver: jt.Float[jt.Array, "nxs"], + is_final: bool, solver: diffrax.AbstractSolver, controller: diffrax.AbstractStepSizeController, root_finder: AbstractRootFinder, @@ -588,63 +595,21 @@ def simulate_condition_unjitted( ..., diffrax._custom_types.BoolScalarLike ], max_steps: int | jnp.int_, - x_preeq: jt.Float[jt.Array, "*nx"] = jnp.array([]), - h_preeq: jt.Float[jt.Array, "*ne"] = jnp.array([]), - mask_reinit: jt.Bool[jt.Array, "*nx"] = jnp.array([]), - x_reinit: jt.Float[jt.Array, "*nx"] = jnp.array([]), - init_override: jt.Float[jt.Array, "*nx"] = jnp.array([]), - init_override_mask: jt.Bool[jt.Array, "*nx"] = jnp.array([]), - ts_mask: jt.Bool[jt.Array, "nt"] = jnp.array([]), - h_mask: jt.Bool[jt.Array, "ne"] = jnp.array([]), - t_zero: jnp.float_ = 0.0, - ret: ReturnValue = ReturnValue.llh, - ) -> tuple[jt.Float[jt.Array, "*nt"], dict]: - """ - Unjitted version of simulate_condition. - - See :meth:`simulate_condition` for full documentation. + ): """ - t0 = t_zero - if p is None: - p = self.parameters - - if not h_mask.shape[0]: - h_mask = jnp.ones(self.n_events, dtype=jnp.bool_) - - if x_preeq.shape[0]: - x = x_preeq - elif init_override.shape[0]: - x_def = self._x0(t0, p) - x = jnp.squeeze( - jnp.where(init_override_mask, init_override, x_def) - ) - else: - x = self._x0(t0, p) + Simulate a single experiment period, starting from ``x_solver``/``h`` + at time ``t0`` with parameters ``p``/``tcl``. - if not ts_mask.shape[0]: - ts_mask = jnp.ones_like(my, dtype=jnp.bool_) - - # Re-initialization - if x_reinit.shape[0]: - x = jnp.where(mask_reinit, x_reinit, x) - - x_solver = self._x_solver(x) - tcl = self._tcl(x, p) - - x_solver, _, h, _ = self._handle_t0_event( - t0, - x_solver, - p, - tcl, - root_finder, - self._root_cond_fn, - self._delta_x, - h_mask, - h_preeq, - {}, - ) + Only the final period of an experiment (``is_final=True``) is + post-equilibrated; earlier periods only integrate up to (and + including) the synthetic hand-off time point appended to + ``ts_dyn`` by :meth:`amici.sim.jax.petab.JAXProblem._get_measurements`. - # Dynamic simulation + :return: + Tuple of (state trajectory, heaviside trajectory, ending reduced + state, ending heaviside state, dynamic simulation statistics, + post-equilibration statistics). + """ if ts_dyn.shape[0]: x_dyn, h_dyn, stats_dyn = solve( p, @@ -667,13 +632,16 @@ def simulate_condition_unjitted( self.observable_ids, ) x_solver = x_dyn[-1, :] + h = h_dyn[-1, :] else: x_dyn = jnp.repeat(x_solver[None, :], ts_dyn.shape[0], axis=0) h_dyn = jnp.repeat(h[None, :], ts_dyn.shape[0], axis=0) stats_dyn = None - # Post-equilibration - if ts_posteq.shape[0]: + # Post-equilibration (only ever meaningful for the final period of + # an experiment; earlier periods just hand off state to the next + # period). + if is_final and ts_posteq.shape[0]: x_solver, h, stats_posteq = eq( p, tcl, @@ -702,9 +670,161 @@ def simulate_condition_unjitted( hs = jnp.concatenate((h_dyn, h_posteq), axis=0) else: hs = jnp.zeros((ts.shape[0], h.shape[0])) - x = jnp.concatenate((x_dyn, x_posteq), axis=0) + xs = jnp.concatenate((x_dyn, x_posteq), axis=0) + + return ts, xs, hs, x_solver, h, stats_dyn, stats_posteq + + def simulate_condition_unjitted( + self, + p: jt.Float[jt.Array, "P np"], + ts_dyn: jt.Float[jt.Array, "P nt_dyn"], + ts_posteq: jt.Float[jt.Array, "P nt_posteq"], + my: jt.Float[jt.Array, "P nt"], + iys: jt.Int[jt.Array, "P nt"], + iy_trafos: jt.Int[jt.Array, "P nt"], + ops: jt.Float[jt.Array, "P nt *nop"], + nps: jt.Float[jt.Array, "P nt *nnp"], + solver: diffrax.AbstractSolver, + controller: diffrax.AbstractStepSizeController, + root_finder: AbstractRootFinder, + adjoint: diffrax.AbstractAdjoint, + steady_state_event: Callable[ + ..., diffrax._custom_types.BoolScalarLike + ], + max_steps: int | jnp.int_, + x_preeq: jt.Float[jt.Array, "*nx"] = jnp.array([]), + h_preeq: jt.Float[jt.Array, "*ne"] = jnp.array([]), + mask_reinit: jt.Bool[jt.Array, "P *nx"] = jnp.array([]), + x_reinit: jt.Float[jt.Array, "P *nx"] = jnp.array([]), + init_override: jt.Float[jt.Array, "*nx"] = jnp.array([]), + init_override_mask: jt.Bool[jt.Array, "*nx"] = jnp.array([]), + ts_mask: jt.Bool[jt.Array, "P nt"] = jnp.array([]), + h_mask: jt.Bool[jt.Array, "ne"] = jnp.array([]), + t_zero: jt.Float[jt.Array, "P"] = jnp.array([0.0]), + ret: ReturnValue = ReturnValue.llh, + ) -> tuple[jt.Float[jt.Array, "*nt"], dict]: + """ + Unjitted version of simulate_condition. + + Chains one ODE integration per experiment period (the leading axis, + of static size ``P``, of ``p``/``ts_dyn``/``ts_posteq``/``my``/ + ``iys``/``iy_trafos``/``ops``/``nps``/``mask_reinit``/``x_reinit``/ + ``ts_mask``/``t_zero``), carrying the ODE state and heaviside/event + state from the end of one period into the start of the next, in + lieu of encoding period transitions as model events. ``P == 1`` + reduces to a single, non-chained simulation. + + See :meth:`simulate_condition` for full documentation. + """ + n_periods = p.shape[0] + + if not h_mask.shape[0]: + h_mask = jnp.ones(self.n_events, dtype=jnp.bool_) - nllhs = self._nllhs(ts, x, p, tcl, hs, my, iys, ops, nps) + if not ts_mask.shape[0]: + ts_mask = jnp.ones_like(my, dtype=jnp.bool_) + + has_reinit = x_reinit.shape[-1] > 0 + + t0_0 = t_zero[0] + if x_preeq.shape[0]: + x = x_preeq + elif init_override.shape[0]: + x_def = self._x0(t0_0, p[0]) + x = jnp.where(init_override_mask, init_override, x_def) + else: + x = self._x0(t0_0, p[0]) + + h = h_preeq + x_solver = None + tcl_prev = None + + ts_list = [] + x_list = [] + h_list = [] + tcl_list = [] + p_list = [] + stats_dyn_list = [] + stats_posteq_final = None + + for i in range(n_periods): + p_i = p[i] + t0_i = t_zero[i] + + if i == 0: + x_i_full = x + else: + # carry reduced state from the end of the previous period + # back to full state space + x_i_full = self._x_rdata(x_solver, tcl_prev) + + if has_reinit: + x_i_full = jnp.where(mask_reinit[i], x_reinit[i], x_i_full) + + x_solver = self._x_solver(x_i_full) + tcl_i = self._tcl(x_i_full, p_i) + + x_solver, _, h, _ = self._handle_t0_event( + t0_i, + x_solver, + p_i, + tcl_i, + root_finder, + self._root_cond_fn, + self._delta_x, + h_mask, + h, + {}, + ) + + is_final = i == n_periods - 1 + ts_i, xs_i, hs_i, x_solver, h, stats_dyn_i, stats_posteq_i = ( + self._simulate_period( + p_i, + t0_i, + ts_dyn[i], + ts_posteq[i], + tcl_i, + h, + h_mask, + x_solver, + is_final, + solver, + controller, + root_finder, + adjoint, + steady_state_event, + max_steps, + ) + ) + if is_final: + stats_posteq_final = stats_posteq_i + + ts_list.append(ts_i) + x_list.append(xs_i) + h_list.append(hs_i) + tcl_list.append(jnp.repeat(tcl_i[None, :], ts_i.shape[0], axis=0)) + p_list.append(jnp.repeat(p_i[None, :], ts_i.shape[0], axis=0)) + stats_dyn_list.append(stats_dyn_i) + + tcl_prev = tcl_i + + ts = jnp.concatenate(ts_list, axis=0) + x = jnp.concatenate(x_list, axis=0) + hs = jnp.concatenate(h_list, axis=0) + tcls = jnp.concatenate(tcl_list, axis=0) + ps = jnp.concatenate(p_list, axis=0) + + my = my.reshape(-1) + iys = iys.reshape(-1) + iy_trafos = iy_trafos.reshape(-1) + # avoid `-1` in reshape: it errors on a `math.prod(...) == 0` + # trailing shape (e.g. no observable/noise parameter overrides) + ops = ops.reshape(ops.shape[0] * ops.shape[1], *ops.shape[2:]) + nps = nps.reshape(nps.shape[0] * nps.shape[1], *nps.shape[2:]) + ts_mask = ts_mask.reshape(-1) + + nllhs = self._nllhs(ts, x, ps, tcls, hs, my, iys, ops, nps) nllhs = jnp.where(ts_mask, nllhs, 0.0) llh = -jnp.sum(nllhs) @@ -713,27 +833,27 @@ def simulate_condition_unjitted( x=x, hs=hs, llh=llh, - stats_dyn=stats_dyn, - stats_posteq=stats_posteq, + stats_dyn=stats_dyn_list, + stats_posteq=stats_posteq_final, ) if ret == ReturnValue.llh: output = llh elif ret == ReturnValue.nllhs: output = nllhs elif ret == ReturnValue.x: - output = self._x_rdatas(x, tcl) + output = self._x_rdatas(x, tcls) elif ret == ReturnValue.x_solver: output = x elif ret == ReturnValue.y: - output = self._ys(ts, x, p, tcl, hs, iys, ops) + output = self._ys(ts, x, ps, tcls, hs, iys, ops) elif ret == ReturnValue.sigmay: - output = self._sigmays(ts, x, p, tcl, hs, iys, ops, nps) + output = self._sigmays(ts, x, ps, tcls, hs, iys, ops, nps) elif ret == ReturnValue.x0: - output = self._x_rdata(x[0, :], tcl) + output = self._x_rdata(x[0, :], tcls[0]) elif ret == ReturnValue.x0_solver: output = x[0, :] elif ret == ReturnValue.tcl: - output = tcl + output = tcls[0] elif ret in (ReturnValue.res, ReturnValue.chi2): obs_trafo = jax.vmap( lambda y, iy_trafo: ( @@ -746,11 +866,11 @@ def simulate_condition_unjitted( ), ) ys_obj = obs_trafo( - self._ys(ts, x, p, tcl, hs, iys, ops), iy_trafos + self._ys(ts, x, ps, tcls, hs, iys, ops), iy_trafos ) m_obj = obs_trafo(my, iy_trafos) if ret == ReturnValue.chi2: - sigma_obj = self._sigmays(ts, x, p, tcl, hs, iys, ops, nps) + sigma_obj = self._sigmays(ts, x, ps, tcls, hs, iys, ops, nps) chi2 = jnp.square((m_obj - ys_obj) / sigma_obj) chi2 = jnp.where(ts_mask, chi2, 0.0) output = jnp.sum(chi2) @@ -765,14 +885,14 @@ def simulate_condition_unjitted( @eqx.filter_jit def simulate_condition( self, - p: jt.Float[jt.Array, "np"] | None, - ts_dyn: jt.Float[jt.Array, "nt_dyn"], - ts_posteq: jt.Float[jt.Array, "nt_posteq"], - my: jt.Float[jt.Array, "nt"], - iys: jt.Int[jt.Array, "nt"], - iy_trafos: jt.Int[jt.Array, "nt"], - ops: jt.Float[jt.Array, "nt *nop"], - nps: jt.Float[jt.Array, "nt *nnp"], + p: jt.Float[jt.Array, "P np"], + ts_dyn: jt.Float[jt.Array, "P nt_dyn"], + ts_posteq: jt.Float[jt.Array, "P nt_posteq"], + my: jt.Float[jt.Array, "P nt"], + iys: jt.Int[jt.Array, "P nt"], + iy_trafos: jt.Int[jt.Array, "P nt"], + ops: jt.Float[jt.Array, "P nt *nop"], + nps: jt.Float[jt.Array, "P nt *nnp"], solver: diffrax.AbstractSolver, controller: diffrax.AbstractStepSizeController, root_finder: AbstractRootFinder, @@ -783,13 +903,13 @@ def simulate_condition( max_steps: int | jnp.int_, x_preeq: jt.Float[jt.Array, "*nx"] = jnp.array([]), h_preeq: jt.Bool[jt.Array, "*ne"] = jnp.array([]), - mask_reinit: jt.Bool[jt.Array, "*nx"] = jnp.array([]), - x_reinit: jt.Float[jt.Array, "*nx"] = jnp.array([]), + mask_reinit: jt.Bool[jt.Array, "P *nx"] = jnp.array([]), + x_reinit: jt.Float[jt.Array, "P *nx"] = jnp.array([]), init_override: jt.Float[jt.Array, "*nx"] = jnp.array([]), init_override_mask: jt.Bool[jt.Array, "*nx"] = jnp.array([]), - ts_mask: jt.Bool[jt.Array, "nt"] = jnp.array([]), + ts_mask: jt.Bool[jt.Array, "P nt"] = jnp.array([]), h_mask: jt.Bool[jt.Array, "ne"] = jnp.array([]), - t_zero: jnp.float_ = 0.0, + t_zero: jt.Float[jt.Array, "P"] = jnp.array([0.0]), ret: ReturnValue = ReturnValue.llh, ) -> tuple[jt.Float[jt.Array, "*nt"], dict]: r""" @@ -798,9 +918,15 @@ def simulate_condition( This is the JIT-compiled version for optimal performance. For runtime type checking with beartype, use :meth:`simulate_condition_unjitted` instead. + Chains one ODE integration per experiment period (the leading axis, + of static size ``P``, of ``p``/``ts_dyn``/``ts_posteq``/``my``/ + ``iys``/``iy_trafos``/``ops``/``nps``/``mask_reinit``/``x_reinit``/ + ``ts_mask``/``t_zero``); ``P == 1`` reduces to a single, non-chained + simulation. + :param p: - parameters for simulation ordered according to ids in :ivar parameter_ids:. If ``None``, - the values stored in :attr:`parameters` are used. + parameters for simulation ordered according to ids in :ivar parameter_ids:, one row per + experiment period. :param ts_dyn: time points for dynamic simulation. Sorted in monotonically increasing order but duplicate time points are allowed to facilitate the evaluation of multiple observables at specific time points. diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 87c2b1357c..7d1258c3d9 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -135,6 +135,7 @@ class JAXProblem(eqx.Module): model: JAXModel simulation_conditions: tuple[tuple[str, ...], ...] _parameter_mappings: dict[str, ...] + _max_periods: int _ts_dyn: np.ndarray _ts_posteq: np.ndarray _my: np.ndarray @@ -181,6 +182,7 @@ def __init__( ) self._parameter_mappings = self._get_parameter_mappings() ( + self._max_periods, self._ts_dyn, self._ts_posteq, self._my, @@ -194,7 +196,7 @@ def __init__( self._np_numeric, self._np_mask, self._np_indices, - ) = self._get_measurements(scs) + ) = self._get_measurements(self._petab_problem.experiments) def save(self, directory: Path): """ @@ -236,8 +238,9 @@ def load(cls, directory: Path): return eqx.tree_deserialise_leaves(f, problem) def _get_measurements( - self, simulation_conditions: pd.DataFrame + self, experiments: list[petabv2.Experiment] ) -> tuple[ + int, np.ndarray, np.ndarray, np.ndarray, @@ -253,11 +256,21 @@ def _get_measurements( np.ndarray, ]: """ - Get measurements for the model based on the provided simulation conditions. + Get measurements for the model, bucketed per experiment and per + (non-pre-equilibration) period, and padded to a common shape + ``(n_experiments, max_periods, n_timepoints, ...)`` suitable for + :func:`eqx.filter_vmap` over experiments with an inner per-period + chain (see :meth:`run_simulation`). + + Periods beyond an experiment's own real period count, and (for + non-terminal real periods) one synthetic time point at the start of + the following period, are added so that state can always be handed + off at exactly the period boundary; those entries are marked + ``False`` in the returned mask and never contribute to the + log-likelihood. - :param simulation_conditions: - Simulation conditions to create parameter mappings for. Same format as returned by - :meth:`petabv1.Problem.get_simulation_conditions_from_measurement_df`. + :param experiments: + Experiments to build measurement arrays for. :return: tuple of padded - dynamic time points @@ -266,7 +279,8 @@ def _get_measurements( - observable indices - observable transformations indices - measurement masks - - data indices (index in petab measurement dataframe). + - data indices (index in petab measurement dataframe, -1 for + synthetic/padding entries). - numeric values for observable parameter overrides - non-numeric mask for observable parameter overrides - parameter indices (problem parameters) for observable parameter overrides @@ -304,81 +318,28 @@ def _get_measurements( .max() ) - for _, simulation_condition in simulation_conditions.iterrows(): + def get_parameter_override(x): if ( - "preequilibration" - in simulation_condition[petabv2.C.CONDITION_ID] - ): - continue - - if isinstance(self._petab_problem, HybridV2Problem): - query = " & ".join( - [ - f"{k} == '{v}'" - if isinstance(v, str) - else f"{k} == {v}" - for k, v in simulation_condition.items() - if k != petabv2.C.CONDITION_ID - ] - ) - else: - query = " & ".join( - [f"{k} == '{v}'" for k, v in simulation_condition.items()] - ) - m = self._petab_problem.measurement_df.query(query).sort_values( - by=petabv2.C.TIME - ) - - ts = m[petabv2.C.TIME] - ts_dyn = ts[np.isfinite(ts)] - ts_posteq = ts[np.logical_not(np.isfinite(ts))] - index = pd.concat([ts_dyn, ts_posteq]).index - ts_dyn = ts_dyn.values - ts_posteq = ts_posteq.values - my = m[petabv2.C.MEASUREMENT].values - iys = np.array( - [ - self.model.observable_ids.index(oid) - for oid in m[petabv2.C.OBSERVABLE_ID].values + x in self._petab_problem.parameter_df.index + and not self._petab_problem.parameter_df.loc[ + x, petabv2.C.ESTIMATE ] - ) - if ( - petabv2.C.NOISE_DISTRIBUTION - in self._petab_problem.observable_df ): - iy_trafos = np.array( - [ - SCALE_TO_INT[petabv2.C.LOG] - if obs.noise_distribution == petabv2.C.LOG_NORMAL - else SCALE_TO_INT[petabv2.C.LIN] - for obs in self._petab_problem.observables - ] - ) - else: - iy_trafos = np.zeros_like(iys) - - parameter_overrides_par_indices = dict() - parameter_overrides_numeric_vals = dict() - parameter_overrides_mask = dict() - - def get_parameter_override(x): - if ( - x in self._petab_problem.parameter_df.index - and not self._petab_problem.parameter_df.loc[ - x, petabv2.C.ESTIMATE - ] - ): - return self._petab_problem.parameter_df.loc[ - x, petabv2.C.NOMINAL_VALUE - ] - return x + return self._petab_problem.parameter_df.loc[ + x, petabv2.C.NOMINAL_VALUE + ] + return x + def get_overrides(m: pd.DataFrame) -> dict[str, tuple]: + """Numeric values, non-numeric mask and parameter indices for + observable/noise parameter overrides of the rows in ``m``.""" + overrides = {} for col in [ petabv2.C.OBSERVABLE_PARAMETERS, petabv2.C.NOISE_PARAMETERS, ]: if col not in m or m[col].isna().all() or all(m[col] == ""): - mat_numeric = jnp.ones((len(m), n_pars[col])) + mat_numeric = np.ones((len(m), n_pars[col])) par_mask = np.zeros_like(mat_numeric, dtype=bool) par_index = np.zeros_like(mat_numeric, dtype=int) elif pd.api.types.is_numeric_dtype(m[col].dtype): @@ -424,32 +385,199 @@ def get_parameter_override(x): # replace dummy index with some valid index par_index[~par_mask] = 0 - parameter_overrides_numeric_vals[col] = mat_numeric - parameter_overrides_mask[col] = par_mask - parameter_overrides_par_indices[col] = par_index - - measurements[tuple(simulation_condition)] = ( - ts_dyn, # 0 - ts_posteq, # 1 - my, # 2 - iys, # 3 - iy_trafos, # 4 - parameter_overrides_numeric_vals[ - petabv2.C.OBSERVABLE_PARAMETERS - ], # 5 - parameter_overrides_mask[petabv2.C.OBSERVABLE_PARAMETERS], # 6 - parameter_overrides_par_indices[ - petabv2.C.OBSERVABLE_PARAMETERS - ], # 7 - parameter_overrides_numeric_vals[ - petabv2.C.NOISE_PARAMETERS - ], # 8 - parameter_overrides_mask[petabv2.C.NOISE_PARAMETERS], # 9 - parameter_overrides_par_indices[ - petabv2.C.NOISE_PARAMETERS - ], # 10 + overrides[col] = (mat_numeric, par_mask, par_index) + return overrides + + def placeholder_row(col: str) -> tuple: + """A single dummy override row for a synthetic/padding entry.""" + return ( + np.ones((1, n_pars[col])), + np.zeros((1, n_pars[col]), dtype=bool), + np.zeros((1, n_pars[col]), dtype=int), ) - petab_indices[tuple(simulation_condition)] = tuple(index.tolist()) + + def get_iy_trafos(iys: np.ndarray) -> np.ndarray: + if ( + petabv2.C.NOISE_DISTRIBUTION + in self._petab_problem.observable_df + ): + return np.array( + [ + SCALE_TO_INT[petabv2.C.LOG] + if obs.noise_distribution == petabv2.C.LOG_NORMAL + else SCALE_TO_INT[petabv2.C.LIN] + for obs in self._petab_problem.observables + ] + ) + return np.zeros_like(iys) + + dyn_periods_by_exp = { + exp.id: [ + period + for period in exp.sorted_periods + if not period.is_preequilibration + ] + for exp in experiments + } + max_periods = max( + (len(periods) for periods in dyn_periods_by_exp.values()), + default=1, + ) + max_periods = max(max_periods, 1) + + for exp in experiments: + dyn_periods = dyn_periods_by_exp[exp.id] + + # All measurements for this experiment. Periods are + # distinguished by time window below since PEtab v2 + # measurements reference an experiment, not an individual + # condition/period. + if isinstance(self._petab_problem, HybridV2Problem): + query = f"{petabv2.C.EXPERIMENT_ID} == '{exp.id}'" + else: + query = f"{petabv2.C.EXPERIMENT_ID} == '{exp.id}'" + m_full = self._petab_problem.measurement_df.query( + query + ).sort_values(by=petabv2.C.TIME) + m_full_times = m_full[petabv2.C.TIME].to_numpy() + + last_period_time = dyn_periods[-1].time if dyn_periods else 0.0 + + for i_period in range(max_periods): + is_real = i_period < len(dyn_periods) + if is_real: + is_own_last = i_period == len(dyn_periods) - 1 + t_lo = dyn_periods[i_period].time + t_hi = ( + np.inf + if is_own_last + else dyn_periods[i_period + 1].time + ) + in_window = ( + (m_full_times >= t_lo) + & (m_full_times < t_hi) + & np.isfinite(m_full_times) + ) + m = m_full[in_window] + + ts_dyn_real = m[petabv2.C.TIME].values + my_real = m[petabv2.C.MEASUREMENT].values + iys_real = np.array( + [ + self.model.observable_ids.index(oid) + for oid in m[petabv2.C.OBSERVABLE_ID].values + ] + ) + iy_trafos_real = get_iy_trafos(iys_real) + overrides_real = get_overrides(m) + index_dyn = list(m.index) + + if is_own_last: + posteq_mask = np.isfinite(m_full_times) == False # noqa: E712 + m_posteq = m_full[posteq_mask] + ts_posteq = m_posteq[petabv2.C.TIME].values + index_posteq = list(m_posteq.index) + else: + ts_posteq = np.array([]) + index_posteq = [] + + if is_own_last: + # No further period to hand off to within this + # experiment's own chain; padding slots (if any) + # after this one are pure no-ops. + if len(ts_dyn_real): + ts_dyn = ts_dyn_real + dyn_valid = np.ones(len(ts_dyn_real), dtype=bool) + my = my_real + iys = iys_real + iy_trafos = iy_trafos_real + overrides = overrides_real + index_dyn_full = index_dyn + else: + # No real dyn measurements in this period (e.g. + # a pure post-equilibration period): still need + # a valid (masked), non-backward-in-time entry + # so that padding never produces a time point + # before this period's start. + ts_dyn = np.array([t_lo]) + dyn_valid = np.array([False]) + my = np.array([0.0]) + iys = np.array([0]) + iy_trafos = np.array([0]) + overrides = { + col: placeholder_row(col) + for col in overrides_real + } + index_dyn_full = [-1] + else: + # Append a synthetic, masked boundary time point at + # the start of the next period so that the ODE + # state is always available at exactly the period + # boundary, regardless of whether there is a real + # measurement there. + ts_dyn = np.append(ts_dyn_real, t_hi) + dyn_valid = np.append( + np.ones(len(ts_dyn_real), dtype=bool), False + ) + my = np.append(my_real, 0.0) + iys = np.append(iys_real, 0) + iy_trafos = np.append(iy_trafos_real, 0) + placeholders = { + col: placeholder_row(col) + for col in overrides_real + } + overrides = { + col: tuple( + np.concatenate( + [overrides_real[col][j], placeholders[col][j]] + ) + for j in range(3) + ) + for col in overrides_real + } + index_dyn_full = [*index_dyn, -1] + + posteq_valid = np.ones(len(ts_posteq), dtype=bool) + index = [*index_dyn_full, *index_posteq] + else: + # Padding slot: this experiment has no period here at + # all. A single masked, zero-duration integration step + # that leaves the carried-over state unchanged. + ts_dyn = np.array([last_period_time]) + dyn_valid = np.array([False]) + my = np.array([0.0]) + iys = np.array([0]) + iy_trafos = np.array([0]) + overrides = { + col: placeholder_row(col) + for col in [ + petabv2.C.OBSERVABLE_PARAMETERS, + petabv2.C.NOISE_PARAMETERS, + ] + } + ts_posteq = np.array([]) + posteq_valid = np.array([]) + index = [-1] + + valid = np.concatenate([dyn_valid, posteq_valid]).astype( + bool + ) + + measurements[(exp.id, i_period)] = ( + ts_dyn, # 0 + ts_posteq, # 1 + my, # 2 + iys, # 3 + iy_trafos, # 4 + overrides[petabv2.C.OBSERVABLE_PARAMETERS][0], # 5 + overrides[petabv2.C.OBSERVABLE_PARAMETERS][1], # 6 + overrides[petabv2.C.OBSERVABLE_PARAMETERS][2], # 7 + overrides[petabv2.C.NOISE_PARAMETERS][0], # 8 + overrides[petabv2.C.NOISE_PARAMETERS][1], # 9 + overrides[petabv2.C.NOISE_PARAMETERS][2], # 10 + valid, # 11 + ) + petab_indices[(exp.id, i_period)] = tuple(index) # compute maximum lengths n_ts_dyn = max(len(mv[0]) for mv in measurements.values()) @@ -459,12 +587,16 @@ def get_parameter_override(x): ts_dyn = np.stack( [ np.pad(mv[0], (0, n_ts_dyn - len(mv[0])), mode="edge") + if len(mv[0]) + else np.zeros(n_ts_dyn, dtype=mv[0].dtype) for mv in measurements.values() ] ) ts_posteq = np.stack( [ np.pad(mv[1], (0, n_ts_posteq - len(mv[1])), mode="edge") + if len(mv[1]) + else np.zeros(n_ts_posteq, dtype=mv[1].dtype) for mv in measurements.values() ] ) @@ -479,8 +611,16 @@ def pad_measurement(x_dyn, x_peq): ) return np.concatenate( ( - np.pad(x_dyn, pad_width_dyn, mode="edge"), - np.pad(x_peq, pad_width_peq, mode="edge"), + np.pad(x_dyn, pad_width_dyn, mode="edge") + if len(x_dyn) + else np.zeros( + (n_ts_dyn, *x_dyn.shape[1:]), dtype=x_dyn.dtype + ), + np.pad(x_peq, pad_width_peq, mode="edge") + if len(x_peq) + else np.zeros( + (n_ts_posteq, *x_peq.shape[1:]), dtype=x_peq.dtype + ), ) ) @@ -509,10 +649,12 @@ def pad_and_stack(output_index: int): np.concatenate( ( np.pad( - np.ones_like(mv[0]), (0, n_ts_dyn - len(mv[0])) + mv[11][: len(mv[0])], + (0, n_ts_dyn - len(mv[0])), ), np.pad( - np.ones_like(mv[1]), (0, n_ts_posteq - len(mv[1])) + mv[11][len(mv[0]) :], + (0, n_ts_posteq - len(mv[1])), ), ) ) @@ -531,31 +673,35 @@ def pad_and_stack(output_index: int): ] ) + n_exp = len(experiments) + + def reshape(arr: np.ndarray) -> np.ndarray: + return arr.reshape(n_exp, max_periods, *arr.shape[1:]) + return ( - ts_dyn, - ts_posteq, - my, - iys, - iy_trafos, - ts_masks, - petab_indices, - op_numeric, - op_mask, - op_indices, - np_numeric, - np_mask, - np_indices, + max_periods, + reshape(ts_dyn), + reshape(ts_posteq), + reshape(my), + reshape(iys), + reshape(iy_trafos), + reshape(ts_masks), + reshape(petab_indices), + reshape(op_numeric), + reshape(op_mask), + reshape(op_indices), + reshape(np_numeric), + reshape(np_mask), + reshape(np_indices), ) 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: _resolve_petab_change_value(ch.target_value) + for ch in c.changes } for c in self._petab_problem.conditions - for ch in c.changes } hybrid_map = ( @@ -566,6 +712,32 @@ def _get_parameter_mappings(self) -> dict[str, ...]: return {"targets_map": targets_map, "hybrid_map": hybrid_map} + def _resolve_parameter_reference( + self, value: float | str + ) -> jt.Float[jt.Scalar, ""]: # noqa: F722 + """ + Resolve a value from ``targets_map`` (see :meth:`_get_parameter_mappings`) + to a JAX scalar. Numeric values pass through; a parameter id + resolves to that parameter's current (estimated) or nominal + (fixed) value. Symbolic references to estimated parameters keep + their dependence on :attr:`parameters` so gradients flow through + correctly. + + :param value: + A numeric literal or PEtab parameter id, as returned by + :func:`_resolve_petab_change_value`. + """ + if isinstance(value, str): + if value in self.parameter_ids: + return self.parameters[self.parameter_ids.index(value)] + return jnp.asarray( + self._petab_problem.parameter_df.loc[ + value, petabv2.C.NOMINAL_VALUE + ], + dtype=self.model.parameters.dtype, + ) + return jnp.asarray(value, dtype=self.model.parameters.dtype) + def get_all_simulation_conditions(self) -> tuple[tuple[str, ...], ...]: if isinstance(self._petab_problem, HybridV2Problem): simulation_conditions = get_simulation_conditions_v2( @@ -1045,8 +1217,21 @@ def _slot_to_array(slot: dict[int | None, object]) -> jnp.ndarray: return nn.forward(net_input)[ind].squeeze() + def _dynamic_periods( + self, experiment: petabv2.Experiment + ) -> list[petabv2.ExperimentPeriod]: + """Non-pre-equilibration periods of ``experiment``, in time order.""" + return [ + period + for period in experiment.sorted_periods + if not period.is_preequilibration + ] + def load_model_parameters( - self, experiment: petabv2.Experiment, is_preeq: bool + self, + experiment: petabv2.Experiment, + is_preeq: bool, + period_index: int | None = None, ) -> jt.Float[jt.Array, "np"]: """ Load parameters for an experiment. @@ -1055,13 +1240,19 @@ def load_model_parameters( Experiment to load parameters for. :param is_preeq: Whether to load preequilibration or simulation parameters. + :param period_index: + Index (within :meth:`_dynamic_periods`) of the period to load + simulation parameters for. Ignored if ``is_preeq`` is ``True``. + Indices beyond the experiment's own period count are clamped to + its last real period (used for padding periods, which are never + actually integrated). :return: Parameters for the experiment. """ p = jnp.stack( [ self._map_experiment_model_parameter_value( - pname, ind, experiment, is_preeq + pname, ind, experiment, is_preeq, period_index ) for ind, pname in enumerate(self.model.parameter_ids) ] @@ -1076,6 +1267,7 @@ def _map_experiment_model_parameter_value( p_index: int, experiment: petabv2.Experiment, is_preeq: bool, + period_index: int | None = None, ): """ Get values for the given parameter `pname` from the relevant petab tables. @@ -1084,14 +1276,20 @@ def _map_experiment_model_parameter_value( :param p_index: Index of the parameter in the model's parameter list :param experiment: PEtab experiment :param is_preeq: Whether to get preequilibration or simulation parameter value + :param period_index: see :meth:`load_model_parameters` :return: Value of the parameter """ - # Find the first period matching the requested phase (preeq vs. sim) - condition_ids = [] - for period in experiment.sorted_periods: - if period.is_preequilibration == is_preeq: - condition_ids = period.condition_ids - break + if period_index is not None: + dyn_periods = self._dynamic_periods(experiment) + clamped_index = min(period_index, len(dyn_periods) - 1) + condition_ids = dyn_periods[clamped_index].condition_ids + else: + # Find the first period matching the requested phase (preeq vs. sim) + condition_ids = [] + for period in experiment.sorted_periods: + if period.is_preequilibration == is_preeq: + condition_ids = period.condition_ids + break _petab_param_map = { param.id: param.nominal_value @@ -1114,9 +1312,7 @@ def _map_experiment_model_parameter_value( } if pname in targets_filtered: - return jnp.asarray( - targets_filtered[pname], dtype=self.model.parameters.dtype - ) + return self._resolve_parameter_reference(targets_filtered[pname]) elif pname in self._parameter_mappings["hybrid_map"]: return jnp.asarray( self._eval_nn( @@ -1154,16 +1350,46 @@ def _find_val(self, param_entry: str): else: return jnp.asarray(param_entry, dtype=self.model.parameters.dtype) + def _first_condition_value( + self, condition_ids: list[str], state_id: str + ): + """ + Find the value assigned to ``state_id`` by the first of + ``condition_ids`` (applied simultaneously, e.g. for one experiment + period) that actually sets it. + + Numeric values are returned as plain Python ``float``s; a + reference to a single other parameter is returned as that + parameter's id (``str``). Compound symbolic expressions are not + supported. + + :param condition_ids: + Condition IDs to check, in priority order. + :param state_id: + State (or parameter) id to look up. + :return: + The assigned value, or ``None`` if none of ``condition_ids`` + sets ``state_id``. + """ + for condition_id in condition_ids: + for change in self._petab_problem[condition_id].changes: + if change.target_id != state_id: + continue + return _resolve_petab_change_value(change.target_value) + return None + def _state_needs_reinitialisation( self, - simulation_condition: str, + condition_ids: list[str], state_id: str, ) -> bool: """ - Check if a state needs reinitialisation for a simulation condition. + Check if a state needs reinitialisation for the given (simultaneous) + conditions. - :param simulation_condition: - simulation condition to check reinitialisation for + :param condition_ids: + condition ids (e.g. of one experiment period) to check + reinitialisation for :param state_id: state id to check reinitialisation for :return: @@ -1175,26 +1401,20 @@ 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 - xval = self._petab_problem.condition_df.loc[ - simulation_condition, state_id - ] - if isinstance(xval, Number) and np.isnan(xval): - return False - return True + return self._first_condition_value(condition_ids, state_id) is not None def _state_reinitialisation_value( self, - simulation_condition: str, + condition_ids: list[str], state_id: str, p: jt.Float[jt.Array, "np"], ) -> jt.Float[jt.Scalar, ""] | float: # noqa: F722 """ Get the reinitialisation value for a state. - :param simulation_condition: - simulation condition to get reinitialisation value for + :param condition_ids: + condition ids (e.g. of one experiment period) to get the + reinitialisation value for :param state_id: state id to get reinitialisation value for :param p: @@ -1203,21 +1423,16 @@ def _state_reinitialisation_value( reinitialisation value for the state """ if state_id in self.nn_output_ids: - return self._eval_nn(state_id, simulation_condition) + return self._eval_nn(state_id, condition_ids[0]) if state_id in self._parameter_mappings["hybrid_map"]: return self._eval_nn( self._parameter_mappings["hybrid_map"][state_id], - simulation_condition, + condition_ids[0], ) - if state_id not in self._petab_problem.condition_df: - # no reinitialisation, return dummy value - return 0.0 - xval = self._petab_problem.condition_df.loc[ - simulation_condition, state_id - ] - if isinstance(xval, Number) and np.isnan(xval): + xval = self._first_condition_value(condition_ids, state_id) + if xval is None: # no reinitialisation, return dummy value return 0.0 if isinstance(xval, Number): @@ -1242,38 +1457,55 @@ def _state_reinitialisation_value( def load_reinitialisation( self, - simulation_condition: str, + condition_ids: list[str] | 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. - - :param simulation_condition: - Simulation condition to load reinitialisation for. + Load reinitialisation values and mask for the state vector for the + given (simultaneous) conditions. + + :param condition_ids: + Condition id(s) (e.g. of one experiment period) to load + reinitialisation for. A bare string is treated as a + single-element list. An empty list means no reinitialisation + (e.g. for a padding period). :param p: Parameters for the simulation condition. :return: Tuple of reinitialisation masm and value for states. """ - if not any( - x_id in self._petab_problem.condition_df + if isinstance(condition_ids, str): + condition_ids = [condition_ids] + + all_condition_targets = { + change.target_id + for condition in self._petab_problem.conditions + for change in condition.changes + } + has_reinitialisable_states = any( + x_id in all_condition_targets or hasattr(self, "nn_output_ids") and x_id in self._parameter_mappings["hybrid_map"] for x_id in self.model.state_ids - ): + ) + if not has_reinitialisable_states: return jnp.array([]), jnp.array([]) + if not condition_ids: + # padding period: no reinitialisation, but still return + # full-shaped, all-False/all-zero arrays for stacking + nx = len(self.model.state_ids) + return jnp.zeros(nx, dtype=bool), jnp.zeros(nx) + mask = jnp.array( [ - self._state_needs_reinitialisation(simulation_condition, x_id) + self._state_needs_reinitialisation(condition_ids, x_id) for x_id in self.model.state_ids ] ) reinit_x = jnp.array( [ - self._state_reinitialisation_value( - simulation_condition, x_id, p - ) + self._state_reinitialisation_value(condition_ids, x_id, p) for x_id in self.model.state_ids ] ) @@ -1300,19 +1532,31 @@ def _prepare_experiments( np_mask: np.ndarray | None = None, np_indices: np.ndarray | None = None, ) -> tuple[ - jt.Float[jt.Array, "nc np"], # noqa: F821, F722 + jt.Float[jt.Array, "nc ... np"], # noqa: F821, F722 jt.Bool[jt.Array, "nx"], # noqa: F821 jt.Float[jt.Array, "nx"], # noqa: F821 - jt.Float[jt.Array, "nc nt nop"], # noqa: F821, F722 - jt.Float[jt.Array, "nc nt nnp"], # noqa: F821, F722 + jt.Float[jt.Array, "nc ... nt nop"], # noqa: F821, F722 + jt.Float[jt.Array, "nc ... nt nnp"], # noqa: F821, F722 ]: """ Prepare experiments for simulation. + For the main simulation (``is_preeq=False``), all returned arrays + (except ``h_mask``) gain a period axis right after the experiment + axis, of static size :attr:`_max_periods`, so that + :meth:`JAXModel.simulate_condition` can chain one ODE integration + per period. Periods beyond an experiment's own period count are + padding periods: they are never reinitialised and never contribute + to the log-likelihood (see :meth:`_get_measurements`), but must + still resolve to *some* valid parameter vector, so they simply + reuse the experiment's own last real period. + :param experiments: Experiments to prepare simulation arrays for. :param conditions: - Simulation conditions to prepare. + Simulation conditions to prepare. Only used for + ``is_preeq=True``, where it is one (pre-equilibration) condition + id per experiment. :param is_preeq: Whether to load preequilibration or simulation parameters. :param op_numeric: @@ -1331,9 +1575,61 @@ def _prepare_experiments( Tuple of parameter arrays, reinitialisation masks and reinitialisation values, observable parameters and noise parameters. """ - p_array = jnp.stack( - [self.load_model_parameters(exp, is_preeq) for exp in experiments] - ) + if is_preeq: + p_array = jnp.stack( + [ + self.load_model_parameters(exp, is_preeq=True) + for exp in experiments + ] + ) + t_zeros = jnp.stack( + [ + exp.periods[0].time if exp.periods[0].time >= 0.0 else 0.0 + for exp in experiments + ] + ) + reinit_condition_ids = conditions + else: + p_array = jnp.stack( + [ + jnp.stack( + [ + self.load_model_parameters( + exp, is_preeq=False, period_index=i + ) + for i in range(self._max_periods) + ] + ) + for exp in experiments + ] + ) + + def period_start_times(exp: petabv2.Experiment) -> jnp.ndarray: + dyn_periods = self._dynamic_periods(exp) + last_time = dyn_periods[-1].time if dyn_periods else 0.0 + return jnp.array( + [ + dyn_periods[i].time + if i < len(dyn_periods) + else last_time + for i in range(self._max_periods) + ] + ) + + t_zeros = jnp.stack( + [period_start_times(exp) for exp in experiments] + ) + + def period_condition_ids( + exp: petabv2.Experiment, i: int + ) -> list[str]: + dyn_periods = self._dynamic_periods(exp) + return dyn_periods[i].condition_ids if i < len(dyn_periods) else [] + + reinit_condition_ids = [ + [period_condition_ids(exp, i) for i in range(self._max_periods)] + for exp in experiments + ] exp_ids = [exp.id for exp in experiments] all_exp_ids = [exp.id for exp in self._petab_problem.experiments] @@ -1347,13 +1643,6 @@ def _prepare_experiments( ] ) - t_zeros = jnp.stack( - [ - exp.periods[0].time if exp.periods[0].time >= 0.0 else 0.0 - for exp in experiments - ] - ) - if self.parameters.size: if isinstance(self._petab_problem, HybridV2Problem): unscaled_parameters = jnp.stack( @@ -1375,43 +1664,78 @@ def _prepare_experiments( ] ) else: - unscaled_parameters = jnp.zeros((*self._ts_masks.shape[:2], 0)) + # No free parameters to gather from. `op_indices`/`np_indices` + # may still contain the dummy index 0 for masked-out (i.e. not + # actually free-parameter-referencing) entries, so keep a + # single dummy slot to gather from -- it is never selected by + # `jnp.where` since the corresponding mask is always False. + unscaled_parameters = jnp.zeros((1,)) + + def gather_unscaled(indices: np.ndarray) -> jnp.ndarray: + fn = lambda ip: unscaled_parameters[ip] # noqa: E731 + for _ in range(indices.ndim): + fn = jax.vmap(fn) + return fn(indices) # placeholder values from sundials code may be needed here if op_numeric is not None and op_numeric.size: op_array = jnp.where( op_mask, - jax.vmap( - jax.vmap(jax.vmap(lambda ip: unscaled_parameters[ip])) - )(op_indices), + gather_unscaled(op_indices), op_numeric, ) else: - op_array = jnp.zeros((*self._ts_masks.shape[:2], 0)) + op_array = jnp.zeros((*self._ts_masks.shape, 0)) if np_numeric is not None and np_numeric.size: np_array = jnp.where( np_mask, - jax.vmap( - jax.vmap(jax.vmap(lambda ip: unscaled_parameters[ip])) - )(np_indices), + gather_unscaled(np_indices), np_numeric, ) else: - np_array = jnp.zeros((*self._ts_masks.shape[:2], 0)) + np_array = jnp.zeros((*self._ts_masks.shape, 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) - ] - ) + if is_preeq: + mask_reinit_array = jnp.stack( + [ + self.load_reinitialisation(cids, p)[0] + for cids, p in zip(reinit_condition_ids, p_array) + ] + ) + x_reinit_array = jnp.stack( + [ + self.load_reinitialisation(cids, p)[1] + for cids, p in zip(reinit_condition_ids, p_array) + ] + ) + else: + mask_reinit_array = jnp.stack( + [ + jnp.stack( + [ + self.load_reinitialisation(cids_i, p_i)[0] + for cids_i, p_i in zip(cids_per_period, p_per_period) + ] + ) + for cids_per_period, p_per_period in zip( + reinit_condition_ids, p_array + ) + ] + ) + x_reinit_array = jnp.stack( + [ + jnp.stack( + [ + self.load_reinitialisation(cids_i, p_i)[1] + for cids_i, p_i in zip(cids_per_period, p_per_period) + ] + ) + for cids_per_period, p_per_period in zip( + reinit_condition_ids, p_array + ) + ] + ) return ( p_array, mask_reinit_array, @@ -1577,17 +1901,6 @@ def run_simulations( Output value and condition specific results and statistics. Results and statistics are returned as a dict with arrays with the leading dimension corresponding to the simulation conditions. """ - simulation_conditions = [ - cid - for exp in experiments - for p in exp.periods - for cid in p.condition_ids - ] - dynamic_conditions = list( - sc for sc in simulation_conditions if "preequilibration" not in sc - ) - dynamic_conditions = list(dict.fromkeys(dynamic_conditions)) - ( p_array, mask_reinit_array, @@ -1598,7 +1911,7 @@ def run_simulations( t_zeros, ) = self._prepare_experiments( experiments, - dynamic_conditions, + [], False, self._op_numeric, self._op_mask, @@ -1624,7 +1937,7 @@ def run_simulations( jnp.array( [ self._eval_nn( - p, exp.periods[-1].condition_ids[0] + p, exp.sorted_periods[-1].condition_ids[0] ) # TODO: Add mapping of p to eval_nn? if p in set(self.model.parameter_ids) else 1.0 @@ -1728,14 +2041,8 @@ def run_preequilibrations( ], max_steps: jnp.int_, ): - simulation_conditions = [ - cid - for exp in experiments - for p in exp.periods - for cid in p.condition_ids - ] preequilibration_conditions = list( - {sc for sc in simulation_conditions if "preequilibration" in sc} + _get_preequilibration_condition_ids(experiments) ) p_array, mask_reinit_array, x_reinit_array, _, _, h_mask, _ = ( @@ -1815,21 +2122,22 @@ def run_simulations( if exp.id in simulation_experiments ] + preeq_condition_ids = _get_preequilibration_condition_ids(experiments) simulation_conditions = [ cid for exp in experiments - for p in exp.periods + for p in exp.sorted_periods for cid in p.condition_ids ] dynamic_conditions = list( - sc for sc in simulation_conditions if "preequilibration" not in sc + sc for sc in simulation_conditions if sc not in preeq_condition_ids ) dynamic_conditions = list(dict.fromkeys(dynamic_conditions)) conditions = { "dynamic_conditions": dynamic_conditions, } - has_preeq = any(exp.periods[0].time < 0.0 for exp in experiments) + has_preeq = any(exp.has_preequilibration for exp in experiments) if has_preeq: preeqs, preresults, h_preeqs = problem.run_preequilibrations( @@ -2085,6 +2393,26 @@ def _conditions_to_experiment_map( return condition_to_experiment +def _get_preequilibration_condition_ids( + experiments: Iterable[petabv2.Experiment], +) -> set[str]: + """Get the condition IDs used by pre-equilibration periods. + + Determined from :attr:`petabv2.ExperimentPeriod.is_preequilibration` + rather than by pattern-matching condition IDs, since condition IDs are no + longer guaranteed to follow the naming convention introduced by + ``ExperimentsToSbmlConverter`` (which is not applied for the JAX + backend). + """ + return { + cid + for experiment in experiments + for period in experiment.sorted_periods + if period.is_preequilibration + for cid in period.condition_ids + } + + def _parse_model_entity_id( model_entity_id: str, nn: dict ) -> list[tuple[str, str]]: @@ -2136,3 +2464,27 @@ def _try_float(value): if isinstance(e, ValueError) and "could not convert" in msg: return value raise + + +def _resolve_petab_change_value(target_value) -> float | str: + """ + Resolve a :class:`petabv2.Change.target_value` (a sympy expression, or + already a plain number) to either a numeric literal or a single + parameter id. + + Only numeric literals and references to a single other parameter are + supported; compound symbolic expressions (e.g. ``"k1 + k2"``) are not. + + :param target_value: + Value to resolve. + :return: + A ``float``, or a ``str`` naming a PEtab/model parameter id. + """ + if getattr(target_value, "is_number", True): + return float(target_value) + if getattr(target_value, "is_Symbol", False): + return str(target_value) + raise NotImplementedError( + "Condition table changes with compound symbolic expressions are " + f"not supported, got {target_value!r}." + ) diff --git a/python/tests/petab_/test_petab_v2_multiperiod.py b/python/tests/petab_/test_petab_v2_multiperiod.py new file mode 100644 index 0000000000..466273f0a8 --- /dev/null +++ b/python/tests/petab_/test_petab_v2_multiperiod.py @@ -0,0 +1,227 @@ +"""Tests for the JAX PEtab v2 simulator's native N-period chaining. + +These exercise PEtab v2 experiments with more than two periods (pre- +equilibration + several sequential dosing/condition-switching periods), +which the JAX backend now simulates by chaining one ODE integration per +period directly, instead of collapsing them into SBML events at import +time (see ``amici.sim.jax.petab.JAXProblem``). +""" + +import jax +import numpy as np +import pytest +from petab.v2 import C, Problem, ProblemConfig +from petab.v2.models.sbml_model import SbmlModel + +from amici.importers.petab import PetabImporter +from amici.sim.jax import ReturnValue, run_simulations + +jax.config.update("jax_enable_x64", True) + + +def _linear_decay_problem() -> Problem: + """A single-species linear-decay model (dx/dt = -k*x).""" + problem = Problem() + problem.config = ProblemConfig() + problem.model = SbmlModel.from_antimony("xx = 1; xx' = -kk*xx;") + problem.add_observable("obs1", "xx", noise_formula="1") + return problem + + +def _import_jax(problem: Problem, module_name: str, tmp_path): + pi = PetabImporter( + petab_problem=problem, + module_name=module_name, + output_dir=tmp_path / module_name, + compile_=True, + jax=True, + verbose=False, + ) + return pi.create_simulator(force_import=True) + + +def test_three_period_chain_matches_analytical_solution(tmp_path): + """A pre-equilibration followed by three sequential dosing periods, + with measurements split across the 2nd and 3rd periods, must match a + closed-form (segment-wise exponential decay) reference solution.""" + problem = _linear_decay_problem() + problem.add_condition("cond_preeq", kk=0.5) + # period 1: t in [0, 2), k=0.3, dosed to xx=3.0 at t=0 + problem.add_condition("cond_p1", kk=0.3, xx=3.0) + # period 2: t in [2, 4), k=0.6, no reinit (state carries over from p1) + problem.add_condition("cond_p2", kk=0.6) + # period 3: t >= 4, k=0.2, dosed again to xx=1.5 + problem.add_condition("cond_p3", kk=0.2, xx=1.5) + problem.add_experiment( + "exp1", + C.TIME_PREEQUILIBRATION, + "cond_preeq", + 0.0, + "cond_p1", + 2.0, + "cond_p2", + 4.0, + "cond_p3", + ) + measurement_times = (0.5, 1.5, 2.5, 3.5, 4.5, 5.5) + for t in measurement_times: + problem.add_measurement( + "obs1", time=t, measurement=0.0, experiment_id="exp1" + ) + + jax_problem = _import_jax(problem, "test_three_period_chain", tmp_path) + assert jax_problem._max_periods == 3 + + x, _ = run_simulations(jax_problem, ret=ReturnValue.x) + ts_mask = np.asarray(jax_problem._ts_masks)[0].reshape(-1) + actual = np.asarray(x)[0].reshape(-1)[ts_mask] + + x2 = 3.0 * np.exp(-0.3 * 2.0) + + def analytical(t): + if t < 2.0: + return 3.0 * np.exp(-0.3 * t) + if t < 4.0: + return x2 * np.exp(-0.6 * (t - 2.0)) + return 1.5 * np.exp(-0.2 * (t - 4.0)) + + expected = np.array([analytical(t) for t in measurement_times]) + np.testing.assert_allclose(actual, expected, rtol=1e-4) + + llh, _ = run_simulations(jax_problem, ret=ReturnValue.llh) + assert np.isfinite(llh) + + +def test_two_period_preequilibration_matches_analytical_solution(tmp_path): + """Regression check for the common (already-supported) pre-equilibration + + one main period case, now driven through the same native-chaining + code path as N>2-period experiments (P=2 special case).""" + problem = _linear_decay_problem() + problem.add_condition("cond_preeq", kk=0.5) + problem.add_condition("cond_main", kk=0.7, xx=2.0) + problem.add_experiment( + "exp1", C.TIME_PREEQUILIBRATION, "cond_preeq", 0.0, "cond_main" + ) + ts = (0.0, 1.0, 2.0, 3.0) + for t in ts: + problem.add_measurement( + "obs1", time=t, measurement=0.0, experiment_id="exp1" + ) + + jax_problem = _import_jax(problem, "test_two_period_preeq", tmp_path) + assert jax_problem._max_periods == 1 + + x, _ = run_simulations(jax_problem, ret=ReturnValue.x) + expected = 2.0 * np.exp(-0.7 * np.array(ts)) + np.testing.assert_allclose( + np.asarray(x)[0, :, 0], expected, rtol=1e-4 + ) + + +def test_single_period_matches_analytical_solution(tmp_path): + """Regression check for the simplest (no pre-equilibration, single + period) case, P=1 with no chaining at all.""" + problem = _linear_decay_problem() + problem.add_condition("cond1", kk=0.7) + problem.add_experiment("exp1", 0.0, "cond1") + ts = (0.0, 1.0, 2.0, 3.0) + for t in ts: + problem.add_measurement( + "obs1", time=t, measurement=0.0, experiment_id="exp1" + ) + + jax_problem = _import_jax(problem, "test_single_period", tmp_path) + + x, _ = run_simulations(jax_problem, ret=ReturnValue.x) + # model's default initial value (xx=1) applies, no reinit here + expected = 1.0 * np.exp(-0.7 * np.array(ts)) + np.testing.assert_allclose( + np.asarray(x)[0, :, 0], expected, rtol=1e-4 + ) + + +def test_gradient_through_multiperiod_chain_matches_finite_differences( + tmp_path, +): + """Gradients of the log-likelihood w.r.t. an estimated parameter that + is referenced from a non-first period's condition table must flow + correctly through the chained periods.""" + problem = _linear_decay_problem() + problem.add_parameter( + "k_free", nominal_value=0.3, estimate=True, lb=0.01, ub=2 + ) + problem.add_condition("cond_preeq", kk=0.5) + problem.add_condition("cond_p1", kk="k_free", xx=3.0) + problem.add_condition("cond_p2", kk=0.6) + problem.add_condition("cond_p3", kk=0.2, xx=1.5) + problem.add_experiment( + "exp1", + C.TIME_PREEQUILIBRATION, + "cond_preeq", + 0.0, + "cond_p1", + 2.0, + "cond_p2", + 4.0, + "cond_p3", + ) + for t in (0.5, 1.5, 2.5, 3.5, 4.5, 5.5): + problem.add_measurement( + "obs1", time=t, measurement=1.0, experiment_id="exp1" + ) + + jax_problem = _import_jax(problem, "test_multiperiod_gradient", tmp_path) + + def llh_fn(p): + return run_simulations(jax_problem.update_parameters(p))[0] + + p0 = jax_problem.parameters + grad = jax.grad(llh_fn)(p0) + + eps = 1e-5 + fd = np.array( + [ + ( + llh_fn(p0.at[i].add(eps)) - llh_fn(p0.at[i].add(-eps)) + ) + / (2 * eps) + for i in range(len(p0)) + ] + ) + np.testing.assert_allclose(np.asarray(grad), fd, rtol=1e-3, atol=1e-4) + + +@pytest.mark.parametrize("jax_flag", [True]) +def test_petab_importer_skips_event_conversion_for_jax(tmp_path, jax_flag): + """The JAX backend must not run ExperimentsToSbmlConverter (the + experiments-to-events conversion), even for experiments with more than + two periods, since it now chains periods natively.""" + problem = _linear_decay_problem() + problem.add_condition("cond_preeq", kk=0.5) + problem.add_condition("cond_p1", kk=0.3, xx=3.0) + problem.add_condition("cond_p2", kk=0.6) + problem.add_condition("cond_p3", kk=0.2, xx=1.5) + problem.add_experiment( + "exp1", + C.TIME_PREEQUILIBRATION, + "cond_preeq", + 0.0, + "cond_p1", + 2.0, + "cond_p2", + 4.0, + "cond_p3", + ) + problem.add_measurement("obs1", time=0.5, measurement=0.0, experiment_id="exp1") + + pi = PetabImporter( + petab_problem=problem, + module_name="test_no_event_conversion", + output_dir=tmp_path / "test_no_event_conversion", + compile_=True, + jax=jax_flag, + verbose=False, + ) + # experiment periods are untouched (still 4: preeq + 3 dosing periods) + assert len(pi.petab_problem.experiments[0].periods) == 4 + assert pi._unconverted_problem is None diff --git a/tests/performance/test_jax_regression.py b/tests/performance/test_jax_regression.py index 0ba530e38a..8b5c7aa874 100644 --- a/tests/performance/test_jax_regression.py +++ b/tests/performance/test_jax_regression.py @@ -89,14 +89,17 @@ def _sim_kwargs(model, solver_kwargs) -> dict: model_name = type(model).__name__ ts_dyn, my, iys, iy_trafos, ops, nps = _MAP[model_name] + # `simulate_condition`/`simulate_condition_unjitted` chain one ODE + # integration per experiment period; add a leading period axis of + # size 1 since these synthetic models are all single-period. return dict( - ts_dyn=ts_dyn, - ts_posteq=jnp.array([]), - my=my, - iys=iys, - iy_trafos=iy_trafos, - ops=ops, - nps=nps, + ts_dyn=ts_dyn[None, :], + ts_posteq=jnp.zeros((1, 0)), + my=my[None, :], + iys=iys[None, :], + iy_trafos=iy_trafos[None, :], + ops=ops[None, :, :], + nps=nps[None, :, :], **solver_kwargs, ) @@ -106,6 +109,11 @@ def _extract_stats(stats: dict) -> dict: out = {} for key in ("stats_dyn", "stats_posteq"): s = stats.get(key) + # `stats_dyn` is a list with one entry per experiment period (see + # JAXModel._simulate_period); these synthetic models are all + # single-period, so use the (only) entry. + if isinstance(s, list): + s = next((entry for entry in reversed(s) if entry is not None), None) if s is None: out[key] = None else: @@ -126,7 +134,8 @@ def test_tier1_fwd_sim( model_id, tier1_models, solver_kwargs, results_collector ): model = tier1_models[model_id] - p = model.parameters + # add a leading period axis of size 1 (single-period simulation) + p = model.parameters[None, :] kwargs = _sim_kwargs(model, solver_kwargs) # Deterministic run (unjitted, for exact step counts) @@ -154,7 +163,8 @@ def test_tier1_fwd_sim( @pytest.mark.parametrize("model_id", TIER1_FWD_CASES) def test_tier1_adj(model_id, tier1_models, solver_kwargs, results_collector): model = tier1_models[model_id] - p = model.parameters + # add a leading period axis of size 1 (single-period simulation) + p = model.parameters[None, :] kwargs = _sim_kwargs(model, solver_kwargs) def _fn(p): @@ -181,7 +191,8 @@ def test_tier1_fwd_sens( model_id, tier1_models, solver_kwargs, results_collector ): model = tier1_models[model_id] - p = model.parameters + # add a leading period axis of size 1 (single-period simulation) + p = model.parameters[None, :] # Forward sensitivity requires DirectAdjoint kwargs = { **_sim_kwargs(model, solver_kwargs), diff --git a/tests/sbml/testSBMLSuiteJax.py b/tests/sbml/testSBMLSuiteJax.py index 4bf9f34e2b..295b3d97a5 100644 --- a/tests/sbml/testSBMLSuiteJax.py +++ b/tests/sbml/testSBMLSuiteJax.py @@ -72,15 +72,18 @@ def run_jax_simulation(model, importer, ts, atol, rtol, tol_factor=1e2): dcoeff=DEFAULT_CONTROLLER_SETTINGS["dcoeff"], ) root_finder = optimistix.Newton(atol=atol, rtol=rtol) + # `simulate_condition` chains one ODE integration per experiment + # period; add a leading period axis of size 1 for this single, + # non-chained simulation. x, stats = model.simulate_condition( - p, - ts_jnp, - jnp.array([]), - zeros, - jnp.zeros_like(ts_jnp, dtype=int), - jnp.zeros_like(ts_jnp, dtype=int), - jnp.zeros((ts_jnp.shape[0], 0)), - jnp.zeros((ts_jnp.shape[0], 0)), + p[None, :], + ts_jnp[None, :], + jnp.zeros((1, 0)), + zeros[None, :], + jnp.zeros_like(ts_jnp, dtype=int)[None, :], + jnp.zeros_like(ts_jnp, dtype=int)[None, :], + jnp.zeros((1, ts_jnp.shape[0], 0)), + jnp.zeros((1, ts_jnp.shape[0], 0)), solver, controller, root_finder, From 9c845160670fadc50de9621404ef3d949ea4b8ba Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 13:57:31 +0000 Subject: [PATCH 02/22] Fix JAX PEtab v2 bugs surfaced by petabtests v2 suite - _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. --- .../amici/importers/petab/_petab_importer.py | 26 ++++++- python/sdist/amici/sim/jax/petab.py | 67 ++++++++++++++----- 2 files changed, 75 insertions(+), 18 deletions(-) diff --git a/python/sdist/amici/importers/petab/_petab_importer.py b/python/sdist/amici/importers/petab/_petab_importer.py index 94794346fd..61ed0fd57f 100644 --- a/python/sdist/amici/importers/petab/_petab_importer.py +++ b/python/sdist/amici/importers/petab/_petab_importer.py @@ -231,6 +231,23 @@ def _preprocess_sbml(self): "The JAX backend does not currently support PEtab problems where network " "parameters appear in the conditions table. " ) + # Condition-table changes are applied directly in Python at + # simulation time (see JAXProblem), by either overriding a + # model parameter or reinitialising a species state. Any other + # target (e.g. a compartment size) has no such mechanism here. + sbml_model = self.petab_problem.model.sbml_model + unsupported_targets = { + target_id + for target_id in condition_targets + if sbml_model.getSpecies(target_id) is None + and sbml_model.getParameter(target_id) is None + } + if unsupported_targets: + raise NotImplementedError( + "The JAX backend only supports condition table changes " + "that target a species or a parameter. Got change(s) " + f"targeting: {sorted(unsupported_targets)}." + ) else: # Convert petab experiments to events, because so far, the # sundials backend only supports preequilibration/presimulation/ @@ -343,8 +360,13 @@ def _do_import_sbml(self): # only ever the indicator variables introduced by the # experiments-to-event conversion above. For the JAX backend, which # keeps the original condition table, this may also contain state - # (species) targets, which must NOT be treated as fixed parameters - # since they are handled via state reinitialisation instead. + # targets (species, or rate-/assignment-rule-governed parameters), + # which must NOT be treated as fixed parameters since they are + # handled via state reinitialisation instead. Compartment targets + # are also excluded here, but are unsupported for the JAX backend + # entirely (see the NotImplementedError raised in + # `_preprocess_sbml`) since AMICI does not support making a + # compartment a runtime-settable fixed parameter either way. fixed_parameters = { change.target_id for experiment in self.petab_problem.experiments diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 7d1258c3d9..6ad1e396f4 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -2291,9 +2291,11 @@ 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 + # NOTE: `condition_df` is long-format (one row per condition + # *change*), so a condition with no changes (e.g. the just-created + # default condition) would not appear in it at all. Read condition + # ids from the `Condition` objects directly instead. + condition_ids = [c.id for c in petab_problem.conditions] condition_ids = [ c for c in condition_ids if "preequilibration" not in c ] @@ -2334,36 +2336,69 @@ def get_simulation_conditions_v2(petab_problem) -> pd.DataFrame: return experiment_df +def _dynamic_condition_index_map( + experiments: list[petabv2.Experiment], +) -> dict[str, tuple[int, int]]: + """Map each non-pre-equilibration condition id to the + ``(experiment_index, period_index)`` position of the period it belongs + to, matching the traversal/de-duplication order used to build + ``dynamic_conditions`` in :func:`run_simulations`. + + :param experiments: + Experiments to build the mapping for, in the same order as used to + build :attr:`JAXProblem._ts_dyn` etc. (i.e. + ``problem._petab_problem.experiments``). + """ + preeq_ids = _get_preequilibration_condition_ids(experiments) + positions: dict[str, tuple[int, int]] = {} + for exp_idx, exp in enumerate(experiments): + period_idx = 0 + for period in exp.sorted_periods: + if period.is_preequilibration: + continue + for cid in period.condition_ids: + if cid not in preeq_ids: + positions.setdefault(cid, (exp_idx, period_idx)) + period_idx += 1 + return positions + + def _build_simulation_df_v2(problem, y, dyn_conditions): """Build petab simulation DataFrame of similation results from a PEtab v2 problem.""" + experiments = problem._petab_problem.experiments + position_map = _dynamic_condition_index_map(experiments) + nt_per_period = problem._ts_masks.shape[-1] + dfs = [] - for ic, sc in enumerate(dyn_conditions): - experiment_id = _conditions_to_experiment_map( - problem._petab_problem.experiment_df - )[sc] + for sc in dyn_conditions: + exp_idx, period_idx = position_map[sc] + experiment_id = experiments[exp_idx].id if experiment_id == "__default__": experiment_id = jnp.nan + mask = problem._ts_masks[exp_idx, period_idx, :] obs = [ problem.model.observable_ids[io] - for io in problem._iys[ic, problem._ts_masks[ic, :]] + for io in problem._iys[exp_idx, period_idx, mask] ] - t = jnp.concat( + t = jnp.concatenate( ( - problem._ts_dyn[ic, :], - problem._ts_posteq[ic, :], + problem._ts_dyn[exp_idx, period_idx, :], + problem._ts_posteq[exp_idx, period_idx, :], ) ) + y_period = y[exp_idx].reshape(-1, nt_per_period)[period_idx, :] + n_real = int(mask.sum()) df_sc = pd.DataFrame( { - petabv2.C.MODEL_ID: [float("nan")] * len(t), + petabv2.C.MODEL_ID: [float("nan")] * n_real, petabv2.C.OBSERVABLE_ID: obs, - petabv2.C.EXPERIMENT_ID: [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: [experiment_id] * n_real, + petabv2.C.TIME: t[mask], + petabv2.C.SIMULATION: y_period[mask], }, - index=problem._petab_measurement_indices[ic, :], + index=problem._petab_measurement_indices[exp_idx, period_idx, mask], ) if ( petabv2.C.OBSERVABLE_PARAMETERS From ff3d96aaa4d54857042509e8a4e707945955e256 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 14:29:08 +0000 Subject: [PATCH 03/22] Fix remaining JAX PEtab v2 bugs found by the full petabtests v2 suite - 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. --- .../amici/importers/petab/_petab_importer.py | 10 ++- python/sdist/amici/sim/jax/petab.py | 68 ++++++++++++------- 2 files changed, 52 insertions(+), 26 deletions(-) diff --git a/python/sdist/amici/importers/petab/_petab_importer.py b/python/sdist/amici/importers/petab/_petab_importer.py index 61ed0fd57f..cfad332b32 100644 --- a/python/sdist/amici/importers/petab/_petab_importer.py +++ b/python/sdist/amici/importers/petab/_petab_importer.py @@ -287,7 +287,15 @@ def _preprocess_pysb(self): pysb.bng.generate_equations(self.petab_problem.model.model) - # Convert PEtab v2 experiments/conditions to events + # Convert PEtab v2 experiments/conditions to events. Unlike for SBML + # (see `_preprocess_sbml`), this is not skipped for the JAX backend: + # PySB condition-table targets are frequently pysb.Observable + # names that alias an underlying pysb.Initial/Expression rather + # than a state or free parameter directly, and applying those + # requires the same model-rewriting this converter already does. + # JAXProblem's native per-period parameter/state resolution has no + # equivalent for that, so PySB models keep going through event + # conversion for both backends. converter = ExperimentsToPySBConverter(self.petab_problem) self.petab_problem, self._events = converter.convert() diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 6ad1e396f4..8169d67fb2 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -1442,12 +1442,10 @@ def _state_reinitialisation_value( # model parameter, return value return p[self.model.parameter_ids.index(xval)] if xval in self.parameter_ids: - # estimated PEtab parameter, return unscaled value + # estimated PEtab parameter, return unscaled value. PEtab v2 + # has no parameterScale column -- all parameters are linear. return jax_unscale( - self.get_petab_parameter_by_id(xval), - self._petab_problem.parameter_df.loc[ - xval, petabv2.PARAMETER_SCALE - ], + self.get_petab_parameter_by_id(xval), petabv2.C.LIN ) # only remaining option is nominal value for PEtab parameter # that is not estimated, return nominal value @@ -1654,12 +1652,7 @@ def period_condition_ids( else: unscaled_parameters = jnp.stack( [ - jax_unscale( - self.parameters[ip], - self._petab_problem.parameter_df.loc[ - p_id, petabv2.C.PARAMETER_SCALE - ], - ) + jax_unscale(self.parameters[ip], petabv2.C.LIN) for ip, p_id in enumerate(self.parameter_ids) ] ) @@ -2123,15 +2116,12 @@ def run_simulations( ] preeq_condition_ids = _get_preequilibration_condition_ids(experiments) - simulation_conditions = [ - cid + dynamic_conditions = [ + _dynamic_period_label(exp, period, preeq_condition_ids) for exp in experiments - for p in exp.sorted_periods - for cid in p.condition_ids + for period in exp.sorted_periods + if not period.is_preequilibration ] - dynamic_conditions = list( - sc for sc in simulation_conditions if sc not in preeq_condition_ids - ) dynamic_conditions = list(dict.fromkeys(dynamic_conditions)) conditions = { "dynamic_conditions": dynamic_conditions, @@ -2336,13 +2326,42 @@ def get_simulation_conditions_v2(petab_problem) -> pd.DataFrame: return experiment_df +def _dynamic_period_label( + experiment: petabv2.Experiment, + period: petabv2.ExperimentPeriod, + preeq_condition_ids: set[str], +) -> str: + """Label identifying a non-pre-equilibration ``period`` for the purposes + of building/looking up the ``dynamic_conditions`` list used by + :func:`run_simulations`/:func:`_build_simulation_df_v2`. + + Each period contributes exactly one label -- even if it has several + simultaneous condition ids (e.g. an indicator-condition encoding, where + a period may carry both an experiment-indicator and a + pre-equilibration-toggle condition id) -- since a period is always one + simulation leg/row-group, regardless of how many condition ids describe + it. The first non-pre-equilibration condition id is used; periods + without any condition table changes at all (e.g. a period that only + fixes a non-zero start time, with all parameters/states left at their + default) have no condition id to key off, so synthesize one unique to + this ``(experiment, period)`` pair. + """ + cids = [ + cid for cid in period.condition_ids if cid not in preeq_condition_ids + ] + if cids: + return cids[0] + return f"__no_condition__{experiment.id}__{period.time}" + + def _dynamic_condition_index_map( experiments: list[petabv2.Experiment], ) -> dict[str, tuple[int, int]]: - """Map each non-pre-equilibration condition id to the - ``(experiment_index, period_index)`` position of the period it belongs - to, matching the traversal/de-duplication order used to build - ``dynamic_conditions`` in :func:`run_simulations`. + """Map each non-pre-equilibration period's label (see + :func:`_dynamic_period_label`) to the ``(experiment_index, + period_index)`` position of the period it belongs to, matching the + traversal/de-duplication order used to build ``dynamic_conditions`` in + :func:`run_simulations`. :param experiments: Experiments to build the mapping for, in the same order as used to @@ -2356,9 +2375,8 @@ def _dynamic_condition_index_map( for period in exp.sorted_periods: if period.is_preequilibration: continue - for cid in period.condition_ids: - if cid not in preeq_ids: - positions.setdefault(cid, (exp_idx, period_idx)) + label = _dynamic_period_label(exp, period, preeq_ids) + positions.setdefault(label, (exp_idx, period_idx)) period_idx += 1 return positions From 58d7e7e60a669a4f6d3c8f5e1a873f1336c718ec Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 14:48:46 +0000 Subject: [PATCH 04/22] Refactor get_overrides into smaller functions; simplify; analytical gradient 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. --- python/sdist/amici/sim/jax/petab.py | 203 ++++++++++-------- .../tests/petab_/test_petab_v2_multiperiod.py | 50 +++-- 2 files changed, 153 insertions(+), 100 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 8169d67fb2..9c19ea4a89 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -70,6 +70,97 @@ def jax_unscale( raise ValueError(f"Invalid parameter scaling: {scale_str}") +def _override_placeholder( + n_rows: int, n_pars: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """All-numeric-one override triple, used for an absent/empty override + column, or for a synthetic/padding measurement row.""" + mat_numeric = np.ones((n_rows, n_pars)) + par_mask = np.zeros_like(mat_numeric, dtype=bool) + par_index = np.zeros_like(mat_numeric, dtype=int) + return mat_numeric, par_mask, par_index + + +def _resolve_override_symbol(value, parameter_df: pd.DataFrame): + """Replace a reference to a non-estimated PEtab parameter by its + nominal value; pass through anything else (numeric literals, estimated + parameter ids, model entity ids) unchanged.""" + if ( + value in parameter_df.index + and not parameter_df.loc[value, petabv2.C.ESTIMATE] + ): + return parameter_df.loc[value, petabv2.C.NOMINAL_VALUE] + return value + + +def _split_override_column( + col_values: pd.Series, n_pars: int, parameter_df: pd.DataFrame +) -> np.ndarray: + """Split ``;``-separated observable/noise parameter override strings + into a ``(n_rows, n_pars)`` matrix of numeric values / free parameter + ids, resolving non-estimated parameter references to their nominal + value and right-padding with ``1.0``.""" + + def resolve_row(entry: list | float) -> list: + if isinstance(entry, list): + return [_resolve_override_symbol(v, parameter_df) for v in entry] + return [] if pd.isna(entry) else [entry] + + rows = col_values.str.split(petabv2.C.PARAMETER_SEPARATOR).apply( + resolve_row + ) + padded = rows.apply( + lambda row: np.pad( + row, (0, n_pars - len(row)), mode="constant", constant_values=1.0 + ) + ) + return np.stack(padded) + + +def _override_triple_from_matrix( + mat: np.ndarray, parameter_ids: tuple[str, ...] +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Split a raw (numeric-or-parameter-id) override matrix into numeric + values, a free-parameter mask, and free-parameter indices, so that the + two can be recombined via ``jnp.where(mask, p[index], numeric)``.""" + par_index = np.vectorize( + lambda x: parameter_ids.index(x) if x in parameter_ids else -1 + )(mat) + par_mask = par_index != -1 + mat = np.where(par_mask, 0.0, mat).astype(float) + par_index[~par_mask] = 0 + return mat, par_mask, par_index + + +def _column_overrides( + m: pd.DataFrame, + col: str, + n_pars: int, + petab_problem: petabv2.Problem, + parameter_ids: tuple[str, ...], +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Numeric values, non-numeric mask and parameter indices for one + observable/noise parameter override column of the rows in ``m``. + + ``petab_problem.parameter_df`` is only accessed once actually needed + (i.e. for a column of non-numeric, string-encoded overrides) rather + than eagerly on every call, since building it can be expensive and, + for some problems (e.g. array-valued PEtab-SciML parameters), emits + pydantic serialization warnings. + """ + if col not in m or m[col].isna().all() or (m[col] == "").all(): + return _override_placeholder(len(m), n_pars) + if pd.api.types.is_numeric_dtype(m[col].dtype): + mat_numeric = np.expand_dims(m[col].values, axis=1) + return ( + mat_numeric, + np.zeros_like(mat_numeric, dtype=bool), + np.zeros_like(mat_numeric, dtype=int), + ) + mat = _split_override_column(m[col], n_pars, petab_problem.parameter_df) + return _override_triple_from_matrix(mat, parameter_ids) + + # IDEA: Implement this class in petab-sciml instead? class HybridProblem(petabv1.Problem): hybridization_df: pd.DataFrame @@ -318,83 +409,26 @@ def _get_measurements( .max() ) - def get_parameter_override(x): - if ( - x in self._petab_problem.parameter_df.index - and not self._petab_problem.parameter_df.loc[ - x, petabv2.C.ESTIMATE - ] - ): - return self._petab_problem.parameter_df.loc[ - x, petabv2.C.NOMINAL_VALUE - ] - return x - def get_overrides(m: pd.DataFrame) -> dict[str, tuple]: """Numeric values, non-numeric mask and parameter indices for observable/noise parameter overrides of the rows in ``m``.""" - overrides = {} - for col in [ - petabv2.C.OBSERVABLE_PARAMETERS, - petabv2.C.NOISE_PARAMETERS, - ]: - if col not in m or m[col].isna().all() or all(m[col] == ""): - mat_numeric = np.ones((len(m), n_pars[col])) - par_mask = np.zeros_like(mat_numeric, dtype=bool) - par_index = np.zeros_like(mat_numeric, dtype=int) - elif pd.api.types.is_numeric_dtype(m[col].dtype): - mat_numeric = np.expand_dims(m[col].values, axis=1) - par_mask = np.zeros_like(mat_numeric, dtype=bool) - par_index = np.zeros_like(mat_numeric, dtype=int) - else: - split_vals = m[col].str.split( - petabv2.C.PARAMETER_SEPARATOR - ) - list_vals = split_vals.apply( - lambda x: ( - [get_parameter_override(y) for y in x] - if isinstance(x, list) - else [] - if pd.isna(x) - else [x] - ) # every string gets transformed to lists, so this is already a float - ) - vals = list_vals.apply( - lambda x: np.pad( - x, - (0, n_pars[col] - len(x)), - mode="constant", - constant_values=1.0, - ) - ) - mat = np.stack(vals) - # deconstruct such that we can reconstruct mapped parameter overrides via vectorized operations - # mat = np.where(par_mask, map(lambda ip: p.at[ip], par_index), mat_numeric) - par_index = np.vectorize( - lambda x: ( - self.parameter_ids.index(x) - if x in self.parameter_ids - else -1 - ) - )(mat) - # map out numeric values - par_mask = par_index != -1 - # remove non-numeric values - mat[par_mask] = 0.0 - mat_numeric = mat.astype(float) - # replace dummy index with some valid index - par_index[~par_mask] = 0 - - overrides[col] = (mat_numeric, par_mask, par_index) - return overrides + return { + col: _column_overrides( + m, + col, + n_pars[col], + self._petab_problem, + self.parameter_ids, + ) + for col in ( + petabv2.C.OBSERVABLE_PARAMETERS, + petabv2.C.NOISE_PARAMETERS, + ) + } def placeholder_row(col: str) -> tuple: """A single dummy override row for a synthetic/padding entry.""" - return ( - np.ones((1, n_pars[col])), - np.zeros((1, n_pars[col]), dtype=bool), - np.zeros((1, n_pars[col]), dtype=int), - ) + return _override_placeholder(1, n_pars[col]) def get_iy_trafos(iys: np.ndarray) -> np.ndarray: if ( @@ -421,7 +455,7 @@ def get_iy_trafos(iys: np.ndarray) -> np.ndarray: } max_periods = max( (len(periods) for periods in dyn_periods_by_exp.values()), - default=1, + default=0, ) max_periods = max(max_periods, 1) @@ -432,10 +466,7 @@ def get_iy_trafos(iys: np.ndarray) -> np.ndarray: # distinguished by time window below since PEtab v2 # measurements reference an experiment, not an individual # condition/period. - if isinstance(self._petab_problem, HybridV2Problem): - query = f"{petabv2.C.EXPERIMENT_ID} == '{exp.id}'" - else: - query = f"{petabv2.C.EXPERIMENT_ID} == '{exp.id}'" + query = f"{petabv2.C.EXPERIMENT_ID} == '{exp.id}'" m_full = self._petab_problem.measurement_df.query( query ).sort_values(by=petabv2.C.TIME) @@ -1431,8 +1462,9 @@ def _state_reinitialisation_value( condition_ids[0], ) - xval = self._first_condition_value(condition_ids, state_id) - if xval is None: + if ( + xval := self._first_condition_value(condition_ids, state_id) + ) is None: # no reinitialisation, return dummy value return 0.0 if isinstance(xval, Number): @@ -2127,9 +2159,7 @@ def run_simulations( "dynamic_conditions": dynamic_conditions, } - has_preeq = any(exp.has_preequilibration for exp in experiments) - - if has_preeq: + if has_preeq := any(exp.has_preequilibration for exp in experiments): preeqs, preresults, h_preeqs = problem.run_preequilibrations( experiments, solver, @@ -2285,9 +2315,10 @@ def add_default_experiment_names_to_v2_problem(petab_problem: petabv2.Problem): # *change*), so a condition with no changes (e.g. the just-created # default condition) would not appear in it at all. Read condition # ids from the `Condition` objects directly instead. - condition_ids = [c.id for c in petab_problem.conditions] condition_ids = [ - c for c in condition_ids if "preequilibration" not in c + c.id + for c in petab_problem.conditions + if "preequilibration" not in c.id ] default_experiment = petabv2.core.Experiment( id="__default__", @@ -2440,10 +2471,10 @@ def _build_simulation_df_v2(problem, y, dyn_conditions): 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 { + row.conditionId: row.experimentId + for row in experiment_df.itertuples() } - return condition_to_experiment def _get_preequilibration_condition_ids( diff --git a/python/tests/petab_/test_petab_v2_multiperiod.py b/python/tests/petab_/test_petab_v2_multiperiod.py index 466273f0a8..ea28d59a90 100644 --- a/python/tests/petab_/test_petab_v2_multiperiod.py +++ b/python/tests/petab_/test_petab_v2_multiperiod.py @@ -140,12 +140,21 @@ def test_single_period_matches_analytical_solution(tmp_path): ) -def test_gradient_through_multiperiod_chain_matches_finite_differences( +def test_gradient_through_multiperiod_chain_matches_analytical_derivative( tmp_path, ): """Gradients of the log-likelihood w.r.t. an estimated parameter that is referenced from a non-first period's condition table must flow - correctly through the chained periods.""" + correctly through the chained periods. + + The reference gradient is derived analytically rather than by finite + differences: within a segment, ``xx(t) = xx0 * exp(-k*(t - t0))``, and + with a unit-sigma Gaussian noise model the per-measurement + log-likelihood contribution is ``-0.5*log(2*pi) - 0.5*(m-y)^2``, so + ``d(llh)/dk = sum_i (m_i - y_i) * dy_i/dk``. Period 3 reinitialises + ``xx`` to a literal value, so it -- and its measurements -- no longer + depend on ``k_free`` at all, giving a zero gradient contribution. + """ problem = _linear_decay_problem() problem.add_parameter( "k_free", nominal_value=0.3, estimate=True, lb=0.01, ub=2 @@ -165,9 +174,14 @@ def test_gradient_through_multiperiod_chain_matches_finite_differences( 4.0, "cond_p3", ) - for t in (0.5, 1.5, 2.5, 3.5, 4.5, 5.5): + measurement_times = (0.5, 1.5, 2.5, 3.5, 4.5, 5.5) + measurement_value = 1.0 + for t in measurement_times: problem.add_measurement( - "obs1", time=t, measurement=1.0, experiment_id="exp1" + "obs1", + time=t, + measurement=measurement_value, + experiment_id="exp1", ) jax_problem = _import_jax(problem, "test_multiperiod_gradient", tmp_path) @@ -178,17 +192,25 @@ def llh_fn(p): p0 = jax_problem.parameters grad = jax.grad(llh_fn)(p0) - eps = 1e-5 - fd = np.array( - [ - ( - llh_fn(p0.at[i].add(eps)) - llh_fn(p0.at[i].add(-eps)) - ) - / (2 * eps) - for i in range(len(p0)) - ] + k = float(p0[0]) + x_p1_end = 3.0 * np.exp(-k * 2.0) # xx at t=2, end of period 1 + dx_p1_end_dk = -2.0 * x_p1_end + + def y_and_dy_dk(t: float) -> tuple[float, float]: + if t < 2.0: + y = 3.0 * np.exp(-k * t) + return y, -t * y + if t < 4.0: + y = x_p1_end * np.exp(-0.6 * (t - 2.0)) + return y, dx_p1_end_dk * np.exp(-0.6 * (t - 2.0)) + # period 3 reinitialises xx to a literal, independent of k_free + return 1.5 * np.exp(-0.2 * (t - 4.0)), 0.0 + + expected_grad = sum( + (measurement_value - y) * dy_dk + for y, dy_dk in map(y_and_dy_dk, measurement_times) ) - np.testing.assert_allclose(np.asarray(grad), fd, rtol=1e-3, atol=1e-4) + np.testing.assert_allclose(float(grad[0]), expected_grad, rtol=1e-4) @pytest.mark.parametrize("jax_flag", [True]) From 0dea65e669fc3e8b2bcf618e2fae7bf5dcf3539b Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 16:56:47 +0000 Subject: [PATCH 05/22] Fix DTypePromotionError regression in override matrix masking 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). --- python/sdist/amici/sim/jax/petab.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 9c19ea4a89..a23adc6fae 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -127,7 +127,13 @@ def _override_triple_from_matrix( lambda x: parameter_ids.index(x) if x in parameter_ids else -1 )(mat) par_mask = par_index != -1 - mat = np.where(par_mask, 0.0, mat).astype(float) + # in-place assignment (rather than e.g. `np.where`) is required here: + # `mat` may be a fixed-width numpy string array (not `object` dtype) if + # every entry is a parameter reference, and `np.where(mask, 0.0, mat)` + # then fails to find a common dtype for the float/string mix. + mat = mat.copy() + mat[par_mask] = 0.0 + mat = mat.astype(float) par_index[~par_mask] = 0 return mat, par_mask, par_index From d2ca2b9d22accf1444dff833494c78db2c1f3ec3 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 17:20:32 +0000 Subject: [PATCH 06/22] Introduce OverrideColumn/_PeriodMeasurements types; de-nest _get_measurements 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. --- python/sdist/amici/sim/jax/petab.py | 481 +++++++++++++++++----------- 1 file changed, 293 insertions(+), 188 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index a23adc6fae..ac0e3ffdab 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Iterable, Sized from numbers import Number from pathlib import Path +from typing import NamedTuple import diffrax import equinox as eqx @@ -70,15 +71,35 @@ def jax_unscale( raise ValueError(f"Invalid parameter scaling: {scale_str}") -def _override_placeholder( - n_rows: int, n_pars: int -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """All-numeric-one override triple, used for an absent/empty override - column, or for a synthetic/padding measurement row.""" - mat_numeric = np.ones((n_rows, n_pars)) - par_mask = np.zeros_like(mat_numeric, dtype=bool) - par_index = np.zeros_like(mat_numeric, dtype=int) - return mat_numeric, par_mask, par_index +class OverrideColumn(NamedTuple): + """Numeric values, free-parameter mask, and free-parameter indices for + one observable/noise parameter override column (each of shape + ``(n_rows, n_pars)``). Recombine via + ``jnp.where(mask, p[index], numeric)``. + """ + + numeric: np.ndarray + mask: np.ndarray + index: np.ndarray + + @classmethod + def placeholder(cls, n_rows: int, n_pars: int) -> "OverrideColumn": + """All-numeric-one column, used for an absent/empty override + column, or for a synthetic/padding measurement row.""" + numeric = np.ones((n_rows, n_pars)) + return cls( + numeric, + np.zeros_like(numeric, dtype=bool), + np.zeros_like(numeric, dtype=int), + ) + + @classmethod + def concatenate(cls, *columns: "OverrideColumn") -> "OverrideColumn": + return cls( + np.concatenate([c.numeric for c in columns]), + np.concatenate([c.mask for c in columns]), + np.concatenate([c.index for c in columns]), + ) def _resolve_override_symbol(value, parameter_df: pd.DataFrame): @@ -119,10 +140,9 @@ def resolve_row(entry: list | float) -> list: def _override_triple_from_matrix( mat: np.ndarray, parameter_ids: tuple[str, ...] -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Split a raw (numeric-or-parameter-id) override matrix into numeric - values, a free-parameter mask, and free-parameter indices, so that the - two can be recombined via ``jnp.where(mask, p[index], numeric)``.""" +) -> OverrideColumn: + """Split a raw (numeric-or-parameter-id) override matrix into an + :class:`OverrideColumn`.""" par_index = np.vectorize( lambda x: parameter_ids.index(x) if x in parameter_ids else -1 )(mat) @@ -135,7 +155,7 @@ def _override_triple_from_matrix( mat[par_mask] = 0.0 mat = mat.astype(float) par_index[~par_mask] = 0 - return mat, par_mask, par_index + return OverrideColumn(mat, par_mask, par_index) def _column_overrides( @@ -144,7 +164,7 @@ def _column_overrides( n_pars: int, petab_problem: petabv2.Problem, parameter_ids: tuple[str, ...], -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +) -> OverrideColumn: """Numeric values, non-numeric mask and parameter indices for one observable/noise parameter override column of the rows in ``m``. @@ -155,10 +175,10 @@ def _column_overrides( pydantic serialization warnings. """ if col not in m or m[col].isna().all() or (m[col] == "").all(): - return _override_placeholder(len(m), n_pars) + return OverrideColumn.placeholder(len(m), n_pars) if pd.api.types.is_numeric_dtype(m[col].dtype): mat_numeric = np.expand_dims(m[col].values, axis=1) - return ( + return OverrideColumn( mat_numeric, np.zeros_like(mat_numeric, dtype=bool), np.zeros_like(mat_numeric, dtype=int), @@ -167,6 +187,132 @@ def _column_overrides( return _override_triple_from_matrix(mat, parameter_ids) +def _get_overrides( + m: pd.DataFrame, + n_pars: dict[str, int], + petab_problem: petabv2.Problem, + parameter_ids: tuple[str, ...], +) -> dict[str, OverrideColumn]: + """Numeric values, non-numeric mask and parameter indices for + observable/noise parameter overrides of the rows in ``m``.""" + return { + col: _column_overrides(m, col, n_pars[col], petab_problem, parameter_ids) + for col in (petabv2.C.OBSERVABLE_PARAMETERS, petabv2.C.NOISE_PARAMETERS) + } + + +def _get_iy_trafos( + iys: np.ndarray, petab_problem: petabv2.Problem +) -> np.ndarray: + """Observable transformation index (see ``SCALE_TO_INT``) for each + observable index in ``iys``.""" + if petabv2.C.NOISE_DISTRIBUTION in petab_problem.observable_df: + return np.array( + [ + SCALE_TO_INT[petabv2.C.LOG] + if obs.noise_distribution == petabv2.C.LOG_NORMAL + else SCALE_TO_INT[petabv2.C.LIN] + for obs in petab_problem.observables + ] + ) + return np.zeros_like(iys) + + +class _PeriodMeasurements(NamedTuple): + """One experiment period's bucketed measurement data, as built by + :meth:`JAXProblem._get_measurements` and consumed by its + padding/stacking step.""" + + ts_dyn: np.ndarray + ts_posteq: np.ndarray + my: np.ndarray + iys: np.ndarray + iy_trafos: np.ndarray + op_overrides: OverrideColumn + noise_overrides: OverrideColumn + valid: np.ndarray + + +def _masked_placeholder_period( + t: float, n_pars: dict[str, int] +) -> tuple[ + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + dict[str, OverrideColumn], +]: + """A single masked, zero-information timepoint at time ``t``. + + Used both for a real period that has no actual measurements in its + time window (e.g. a pure post-equilibration period), and for a padding + period that doesn't exist for a given experiment -- in both cases the + only requirement is a valid, non-backward-in-time, fully-masked-out + entry so downstream padding/chaining has something to work with. + + :return: + Tuple of ``(ts_dyn, dyn_valid, my, iys, iy_trafos, overrides)``. + """ + return ( + np.array([t]), + np.array([False]), + np.array([0.0]), + np.array([0]), + np.array([0]), + { + col: OverrideColumn.placeholder(1, n_pars[col]) + for col in (petabv2.C.OBSERVABLE_PARAMETERS, petabv2.C.NOISE_PARAMETERS) + }, + ) + + +def _pad_measurement( + x_dyn: np.ndarray, x_peq: np.ndarray, n_ts_dyn: int, n_ts_posteq: int +) -> np.ndarray: + """Right-pad ``x_dyn``/``x_peq`` (edge mode: repeat the last value) to + ``n_ts_dyn``/``n_ts_posteq`` along the first axis, then concatenate.""" + pad_width_dyn = tuple( + [(0, n_ts_dyn - len(x_dyn))] + [(0, 0)] * (x_dyn.ndim - 1) + ) + pad_width_peq = tuple( + [(0, n_ts_posteq - len(x_peq))] + [(0, 0)] * (x_peq.ndim - 1) + ) + return np.concatenate( + ( + np.pad(x_dyn, pad_width_dyn, mode="edge") + if len(x_dyn) + else np.zeros((n_ts_dyn, *x_dyn.shape[1:]), dtype=x_dyn.dtype), + np.pad(x_peq, pad_width_peq, mode="edge") + if len(x_peq) + else np.zeros( + (n_ts_posteq, *x_peq.shape[1:]), dtype=x_peq.dtype + ), + ) + ) + + +def _pad_and_stack( + measurements: dict[tuple[str, int], _PeriodMeasurements], + extractor: Callable[[_PeriodMeasurements], np.ndarray], + n_ts_dyn: int, + n_ts_posteq: int, +) -> np.ndarray: + """Apply ``extractor`` to every bucketed period, split each result at + its own dynamic/post-equilibrium boundary, pad, and stack.""" + return np.stack( + [ + _pad_measurement( + extractor(mv)[: len(mv.ts_dyn)], + extractor(mv)[len(mv.ts_dyn) :], + n_ts_dyn, + n_ts_posteq, + ) + for mv in measurements.values() + ] + ) + + # IDEA: Implement this class in petab-sciml instead? class HybridProblem(petabv1.Problem): hybridization_df: pd.DataFrame @@ -385,7 +531,7 @@ def _get_measurements( - non-numeric mask for noise parameter overrides - parameter indices (problem parameters) for noise parameter overrides """ - measurements = dict() + measurements: dict[tuple[str, int], _PeriodMeasurements] = {} petab_indices = dict() n_pars = dict() @@ -415,42 +561,6 @@ def _get_measurements( .max() ) - def get_overrides(m: pd.DataFrame) -> dict[str, tuple]: - """Numeric values, non-numeric mask and parameter indices for - observable/noise parameter overrides of the rows in ``m``.""" - return { - col: _column_overrides( - m, - col, - n_pars[col], - self._petab_problem, - self.parameter_ids, - ) - for col in ( - petabv2.C.OBSERVABLE_PARAMETERS, - petabv2.C.NOISE_PARAMETERS, - ) - } - - def placeholder_row(col: str) -> tuple: - """A single dummy override row for a synthetic/padding entry.""" - return _override_placeholder(1, n_pars[col]) - - def get_iy_trafos(iys: np.ndarray) -> np.ndarray: - if ( - petabv2.C.NOISE_DISTRIBUTION - in self._petab_problem.observable_df - ): - return np.array( - [ - SCALE_TO_INT[petabv2.C.LOG] - if obs.noise_distribution == petabv2.C.LOG_NORMAL - else SCALE_TO_INT[petabv2.C.LIN] - for obs in self._petab_problem.observables - ] - ) - return np.zeros_like(iys) - dyn_periods_by_exp = { exp.id: [ period @@ -481,8 +591,17 @@ def get_iy_trafos(iys: np.ndarray) -> np.ndarray: last_period_time = dyn_periods[-1].time if dyn_periods else 0.0 for i_period in range(max_periods): - is_real = i_period < len(dyn_periods) - if is_real: + if i_period >= len(dyn_periods): + # Padding slot: this experiment has no period here at + # all. A single masked, zero-duration integration step + # that leaves the carried-over state unchanged. + ts_dyn, dyn_valid, my, iys, iy_trafos, overrides = ( + _masked_placeholder_period(last_period_time, n_pars) + ) + ts_posteq = np.array([]) + posteq_valid = np.array([]) + index = [-1] + else: is_own_last = i_period == len(dyn_periods) - 1 t_lo = dyn_periods[i_period].time t_hi = ( @@ -505,8 +624,12 @@ def get_iy_trafos(iys: np.ndarray) -> np.ndarray: for oid in m[petabv2.C.OBSERVABLE_ID].values ] ) - iy_trafos_real = get_iy_trafos(iys_real) - overrides_real = get_overrides(m) + iy_trafos_real = _get_iy_trafos( + iys_real, self._petab_problem + ) + overrides_real = _get_overrides( + m, n_pars, self._petab_problem, self.parameter_ids + ) index_dyn = list(m.index) if is_own_last: @@ -514,11 +637,7 @@ def get_iy_trafos(iys: np.ndarray) -> np.ndarray: m_posteq = m_full[posteq_mask] ts_posteq = m_posteq[petabv2.C.TIME].values index_posteq = list(m_posteq.index) - else: - ts_posteq = np.array([]) - index_posteq = [] - if is_own_last: # No further period to hand off to within this # experiment's own chain; padding slots (if any) # after this one are pure no-ops. @@ -536,17 +655,19 @@ def get_iy_trafos(iys: np.ndarray) -> np.ndarray: # a valid (masked), non-backward-in-time entry # so that padding never produces a time point # before this period's start. - ts_dyn = np.array([t_lo]) - dyn_valid = np.array([False]) - my = np.array([0.0]) - iys = np.array([0]) - iy_trafos = np.array([0]) - overrides = { - col: placeholder_row(col) - for col in overrides_real - } + ( + ts_dyn, + dyn_valid, + my, + iys, + iy_trafos, + overrides, + ) = _masked_placeholder_period(t_lo, n_pars) index_dyn_full = [-1] else: + ts_posteq = np.array([]) + index_posteq = [] + # Append a synthetic, masked boundary time point at # the start of the next period so that the ODE # state is always available at exactly the period @@ -559,139 +680,119 @@ def get_iy_trafos(iys: np.ndarray) -> np.ndarray: my = np.append(my_real, 0.0) iys = np.append(iys_real, 0) iy_trafos = np.append(iy_trafos_real, 0) - placeholders = { - col: placeholder_row(col) - for col in overrides_real - } overrides = { - col: tuple( - np.concatenate( - [overrides_real[col][j], placeholders[col][j]] - ) - for j in range(3) + col: OverrideColumn.concatenate( + overrides_real[col], + OverrideColumn.placeholder(1, n_pars[col]), ) for col in overrides_real } index_dyn_full = [*index_dyn, -1] - posteq_valid = np.ones(len(ts_posteq), dtype=bool) index = [*index_dyn_full, *index_posteq] - else: - # Padding slot: this experiment has no period here at - # all. A single masked, zero-duration integration step - # that leaves the carried-over state unchanged. - ts_dyn = np.array([last_period_time]) - dyn_valid = np.array([False]) - my = np.array([0.0]) - iys = np.array([0]) - iy_trafos = np.array([0]) - overrides = { - col: placeholder_row(col) - for col in [ - petabv2.C.OBSERVABLE_PARAMETERS, - petabv2.C.NOISE_PARAMETERS, - ] - } - ts_posteq = np.array([]) - posteq_valid = np.array([]) - index = [-1] + posteq_valid = np.ones(len(ts_posteq), dtype=bool) valid = np.concatenate([dyn_valid, posteq_valid]).astype( bool ) - measurements[(exp.id, i_period)] = ( - ts_dyn, # 0 - ts_posteq, # 1 - my, # 2 - iys, # 3 - iy_trafos, # 4 - overrides[petabv2.C.OBSERVABLE_PARAMETERS][0], # 5 - overrides[petabv2.C.OBSERVABLE_PARAMETERS][1], # 6 - overrides[petabv2.C.OBSERVABLE_PARAMETERS][2], # 7 - overrides[petabv2.C.NOISE_PARAMETERS][0], # 8 - overrides[petabv2.C.NOISE_PARAMETERS][1], # 9 - overrides[petabv2.C.NOISE_PARAMETERS][2], # 10 - valid, # 11 + measurements[(exp.id, i_period)] = _PeriodMeasurements( + ts_dyn=ts_dyn, + ts_posteq=ts_posteq, + my=my, + iys=iys, + iy_trafos=iy_trafos, + op_overrides=overrides[petabv2.C.OBSERVABLE_PARAMETERS], + noise_overrides=overrides[petabv2.C.NOISE_PARAMETERS], + valid=valid, ) petab_indices[(exp.id, i_period)] = tuple(index) # compute maximum lengths - n_ts_dyn = max(len(mv[0]) for mv in measurements.values()) - n_ts_posteq = max(len(mv[1]) for mv in measurements.values()) + n_ts_dyn = max(len(mv.ts_dyn) for mv in measurements.values()) + n_ts_posteq = max(len(mv.ts_posteq) for mv in measurements.values()) # pad with last value and stack ts_dyn = np.stack( [ - np.pad(mv[0], (0, n_ts_dyn - len(mv[0])), mode="edge") - if len(mv[0]) - else np.zeros(n_ts_dyn, dtype=mv[0].dtype) + np.pad(mv.ts_dyn, (0, n_ts_dyn - len(mv.ts_dyn)), mode="edge") + if len(mv.ts_dyn) + else np.zeros(n_ts_dyn, dtype=mv.ts_dyn.dtype) for mv in measurements.values() ] ) ts_posteq = np.stack( [ - np.pad(mv[1], (0, n_ts_posteq - len(mv[1])), mode="edge") - if len(mv[1]) - else np.zeros(n_ts_posteq, dtype=mv[1].dtype) + np.pad( + mv.ts_posteq, + (0, n_ts_posteq - len(mv.ts_posteq)), + mode="edge", + ) + if len(mv.ts_posteq) + else np.zeros(n_ts_posteq, dtype=mv.ts_posteq.dtype) for mv in measurements.values() ] ) - def pad_measurement(x_dyn, x_peq): - # only pad first axis - pad_width_dyn = tuple( - [(0, n_ts_dyn - len(x_dyn))] + [(0, 0)] * (x_dyn.ndim - 1) - ) - pad_width_peq = tuple( - [(0, n_ts_posteq - len(x_peq))] + [(0, 0)] * (x_peq.ndim - 1) - ) - return np.concatenate( - ( - np.pad(x_dyn, pad_width_dyn, mode="edge") - if len(x_dyn) - else np.zeros( - (n_ts_dyn, *x_dyn.shape[1:]), dtype=x_dyn.dtype - ), - np.pad(x_peq, pad_width_peq, mode="edge") - if len(x_peq) - else np.zeros( - (n_ts_posteq, *x_peq.shape[1:]), dtype=x_peq.dtype - ), - ) - ) - - def pad_and_stack(output_index: int): - return np.stack( - [ - pad_measurement( - mv[output_index][: len(mv[0])], - mv[output_index][len(mv[0]) :], - ) - for mv in measurements.values() - ] - ) - - my = pad_and_stack(2) - iys = pad_and_stack(3) - iy_trafos = pad_and_stack(4) - op_numeric = pad_and_stack(5) - op_mask = pad_and_stack(6) - op_indices = pad_and_stack(7) - np_numeric = pad_and_stack(8) - np_mask = pad_and_stack(9) - np_indices = pad_and_stack(10) + my = _pad_and_stack( + measurements, lambda mv: mv.my, n_ts_dyn, n_ts_posteq + ) + iys = _pad_and_stack( + measurements, lambda mv: mv.iys, n_ts_dyn, n_ts_posteq + ) + iy_trafos = _pad_and_stack( + measurements, lambda mv: mv.iy_trafos, n_ts_dyn, n_ts_posteq + ) + op_numeric = _pad_and_stack( + measurements, + lambda mv: mv.op_overrides.numeric, + n_ts_dyn, + n_ts_posteq, + ) + op_mask = _pad_and_stack( + measurements, + lambda mv: mv.op_overrides.mask, + n_ts_dyn, + n_ts_posteq, + ) + op_indices = _pad_and_stack( + measurements, + lambda mv: mv.op_overrides.index, + n_ts_dyn, + n_ts_posteq, + ) + np_numeric = _pad_and_stack( + measurements, + lambda mv: mv.noise_overrides.numeric, + n_ts_dyn, + n_ts_posteq, + ) + np_mask = _pad_and_stack( + measurements, + lambda mv: mv.noise_overrides.mask, + n_ts_dyn, + n_ts_posteq, + ) + np_indices = _pad_and_stack( + measurements, + lambda mv: mv.noise_overrides.index, + n_ts_dyn, + n_ts_posteq, + ) + # mask padding must stay `False` (not repeat the last real mask + # value), so this is stacked directly rather than via + # `_pad_and_stack` (which pads in "edge" mode). ts_masks = np.stack( [ np.concatenate( ( np.pad( - mv[11][: len(mv[0])], - (0, n_ts_dyn - len(mv[0])), + mv.valid[: len(mv.ts_dyn)], + (0, n_ts_dyn - len(mv.ts_dyn)), ), np.pad( - mv[11][len(mv[0]) :], - (0, n_ts_posteq - len(mv[1])), + mv.valid[len(mv.ts_dyn) :], + (0, n_ts_posteq - len(mv.ts_posteq)), ), ) ) @@ -700,9 +801,11 @@ def pad_and_stack(output_index: int): ).astype(bool) petab_indices = np.stack( [ - pad_measurement( - np.array(idx[: len(mv[0])]), - np.array(idx[len(mv[0]) :]), + _pad_measurement( + np.array(idx[: len(mv.ts_dyn)]), + np.array(idx[len(mv.ts_dyn) :]), + n_ts_dyn, + n_ts_posteq, ) for mv, idx in zip( measurements.values(), petab_indices.values() @@ -711,25 +814,27 @@ def pad_and_stack(output_index: int): ) n_exp = len(experiments) - - def reshape(arr: np.ndarray) -> np.ndarray: - return arr.reshape(n_exp, max_periods, *arr.shape[1:]) - + outputs = ( + ts_dyn, + ts_posteq, + my, + iys, + iy_trafos, + ts_masks, + petab_indices, + op_numeric, + op_mask, + op_indices, + np_numeric, + np_mask, + np_indices, + ) return ( max_periods, - reshape(ts_dyn), - reshape(ts_posteq), - reshape(my), - reshape(iys), - reshape(iy_trafos), - reshape(ts_masks), - reshape(petab_indices), - reshape(op_numeric), - reshape(op_mask), - reshape(op_indices), - reshape(np_numeric), - reshape(np_mask), - reshape(np_indices), + *( + arr.reshape(n_exp, max_periods, *arr.shape[1:]) + for arr in outputs + ), ) def _get_parameter_mappings(self) -> dict[str, ...]: From 07c3691f3f4959bf03cda2d4bd6b258376667bd5 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 22:49:17 +0000 Subject: [PATCH 07/22] Fix dtype-freezing default arguments and a missed period-axis call site 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. --- python/sdist/amici/sim/jax/model.py | 61 ++++++++++++++++++++--------- python/tests/test_jax.py | 25 ++++++++---- 2 files changed, 60 insertions(+), 26 deletions(-) diff --git a/python/sdist/amici/sim/jax/model.py b/python/sdist/amici/sim/jax/model.py index 8c48d7cf40..99eadb2f67 100644 --- a/python/sdist/amici/sim/jax/model.py +++ b/python/sdist/amici/sim/jax/model.py @@ -692,15 +692,15 @@ def simulate_condition_unjitted( ..., diffrax._custom_types.BoolScalarLike ], max_steps: int | jnp.int_, - x_preeq: jt.Float[jt.Array, "*nx"] = jnp.array([]), - h_preeq: jt.Float[jt.Array, "*ne"] = jnp.array([]), - mask_reinit: jt.Bool[jt.Array, "P *nx"] = jnp.array([]), - x_reinit: jt.Float[jt.Array, "P *nx"] = jnp.array([]), - init_override: jt.Float[jt.Array, "*nx"] = jnp.array([]), - init_override_mask: jt.Bool[jt.Array, "*nx"] = jnp.array([]), - ts_mask: jt.Bool[jt.Array, "P nt"] = jnp.array([]), - h_mask: jt.Bool[jt.Array, "ne"] = jnp.array([]), - t_zero: jt.Float[jt.Array, "P"] = jnp.array([0.0]), + x_preeq: jt.Float[jt.Array, "*nx"] | None = None, + h_preeq: jt.Float[jt.Array, "*ne"] | None = None, + mask_reinit: jt.Bool[jt.Array, "P *nx"] | None = None, + x_reinit: jt.Float[jt.Array, "P *nx"] | None = None, + init_override: jt.Float[jt.Array, "*nx"] | None = None, + init_override_mask: jt.Bool[jt.Array, "*nx"] | None = None, + ts_mask: jt.Bool[jt.Array, "P nt"] | None = None, + h_mask: jt.Bool[jt.Array, "ne"] | None = None, + t_zero: jt.Float[jt.Array, "P"] | None = None, ret: ReturnValue = ReturnValue.llh, ) -> tuple[jt.Float[jt.Array, "*nt"], dict]: """ @@ -718,6 +718,31 @@ def simulate_condition_unjitted( """ n_periods = p.shape[0] + # Normalize omitted optional arrays here, at call time, rather than + # via eager `jnp.array(...)`-valued default arguments: a default + # constructed once at function-definition time freezes to + # float32 if `jax_enable_x64` is enabled only after this module is + # first imported, silently diverging in dtype from every other + # (call-time-constructed) array flowing through the same call. + if x_preeq is None: + x_preeq = jnp.array([]) + if h_preeq is None: + h_preeq = jnp.array([]) + if mask_reinit is None: + mask_reinit = jnp.array([]) + if x_reinit is None: + x_reinit = jnp.array([]) + if init_override is None: + init_override = jnp.array([]) + if init_override_mask is None: + init_override_mask = jnp.array([]) + if ts_mask is None: + ts_mask = jnp.array([]) + if h_mask is None: + h_mask = jnp.array([]) + if t_zero is None: + t_zero = jnp.zeros(n_periods) + if not h_mask.shape[0]: h_mask = jnp.ones(self.n_events, dtype=jnp.bool_) @@ -901,15 +926,15 @@ def simulate_condition( ..., diffrax._custom_types.BoolScalarLike ], max_steps: int | jnp.int_, - x_preeq: jt.Float[jt.Array, "*nx"] = jnp.array([]), - h_preeq: jt.Bool[jt.Array, "*ne"] = jnp.array([]), - mask_reinit: jt.Bool[jt.Array, "P *nx"] = jnp.array([]), - x_reinit: jt.Float[jt.Array, "P *nx"] = jnp.array([]), - init_override: jt.Float[jt.Array, "*nx"] = jnp.array([]), - init_override_mask: jt.Bool[jt.Array, "*nx"] = jnp.array([]), - ts_mask: jt.Bool[jt.Array, "P nt"] = jnp.array([]), - h_mask: jt.Bool[jt.Array, "ne"] = jnp.array([]), - t_zero: jt.Float[jt.Array, "P"] = jnp.array([0.0]), + x_preeq: jt.Float[jt.Array, "*nx"] | None = None, + h_preeq: jt.Bool[jt.Array, "*ne"] | None = None, + mask_reinit: jt.Bool[jt.Array, "P *nx"] | None = None, + x_reinit: jt.Float[jt.Array, "P *nx"] | None = None, + init_override: jt.Float[jt.Array, "*nx"] | None = None, + init_override_mask: jt.Bool[jt.Array, "*nx"] | None = None, + ts_mask: jt.Bool[jt.Array, "P nt"] | None = None, + h_mask: jt.Bool[jt.Array, "ne"] | None = None, + t_zero: jt.Float[jt.Array, "P"] | None = None, ret: ReturnValue = ReturnValue.llh, ) -> tuple[jt.Float[jt.Array, "*nt"], dict]: r""" diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 238e2bc89c..9ba12ef25a 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -194,14 +194,20 @@ def check_fields_jax( } p = jnp.array([par_dict[par_id] for par_id in jax_model.parameter_ids]) + # `simulate_condition[_unjitted]` chains one ODE integration per + # experiment period; add a leading period axis of size 1 for this + # single, non-chained simulation. `p` itself is kept 1-D here (and + # the period axis added inside `fun` below) so that `jax.grad`/ + # `jax.jacfwd` differentiate w.r.t. the original 1-D parameter vector, + # matching the shapes the rest of this function already expects. kwargs = { - "ts_dyn": jnp.array(ts_dyn), - "ts_posteq": jnp.array(ts_posteq), - "my": jnp.array(my), - "iys": jnp.array(iys), - "ops": jnp.zeros((*my.shape[:2], 0)), - "nps": jnp.zeros((*my.shape[:2], 0)), - "iy_trafos": jnp.array(iy_trafos), + "ts_dyn": jnp.array(ts_dyn)[None, :], + "ts_posteq": jnp.array(ts_posteq)[None, :], + "my": jnp.array(my)[None, :], + "iys": jnp.array(iys)[None, :], + "ops": jnp.zeros((1, *my.shape[:2], 0)), + "nps": jnp.zeros((1, *my.shape[:2], 0)), + "iy_trafos": jnp.array(iy_trafos)[None, :], "x_preeq": jnp.array([]), "solver": diffrax.Kvaerno5(), "controller": diffrax.PIDController(atol=1e-8, rtol=1e-8), @@ -212,7 +218,10 @@ def check_fields_jax( } # Use beartype-wrapped unjitted version for type checking # (beartype cannot introspect jitted functions, so we wrap the unjitted version) - fun = beartype(jax_model.simulate_condition_unjitted) + fun_periodic = beartype(jax_model.simulate_condition_unjitted) + + def fun(p, **kw): + return fun_periodic(p[None, :], **kw) for output in ["llh", "x0", "x", "y", "res"]: okwargs = kwargs | { From 1059a4daee776eb0945181dec8c17eb85bc9f578 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 22:49:25 +0000 Subject: [PATCH 08/22] Update example notebook for native JAX period chaining 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. --- .../example_jax_petab/ExampleJaxPEtab.ipynb | 58 ++----------------- 1 file changed, 4 insertions(+), 54 deletions(-) diff --git a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb index b4fe05dac4..6a92e7cb36 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 = \"__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 = \"__default__\"\n", "\n", "\n", "def plot_simulation(results):\n", @@ -382,57 +382,7 @@ "id": "1a91aff44b93157", "metadata": {}, "outputs": [], - "source": [ - "import diffrax\n", - "import jax.numpy as jnp\n", - "import optimistix\n", - "from amici.sim.jax import ReturnValue\n", - "\n", - "# Define the simulation condition\n", - "experiment_condition = \"_petab_experiment_condition___default__\"\n", - "ic = 0\n", - "\n", - "# Load condition-specific data\n", - "ts_dyn = jax_problem._ts_dyn[ic, :]\n", - "ts_posteq = jax_problem._ts_posteq[ic, :]\n", - "my = jax_problem._my[ic, :]\n", - "iys = jax_problem._iys[ic, :]\n", - "iy_trafos = jax_problem._iy_trafos[ic, :]\n", - "ops = jax_problem._op_numeric[ic, :]\n", - "nps = jax_problem._np_numeric[ic, :]\n", - "\n", - "# Load parameters for the specified condition\n", - "p = jax_problem.load_model_parameters(\n", - " jax_problem._petab_problem.experiments[0], is_preeq=False\n", - ")\n", - "\n", - "\n", - "# Define a function to compute the gradient with respect to dynamic timepoints\n", - "@eqx.filter_jacfwd\n", - "def grad_ts_dyn(tt):\n", - " return jax_problem.model.simulate_condition(\n", - " p=p,\n", - " ts_dyn=tt,\n", - " ts_posteq=ts_posteq,\n", - " my=jnp.array(my),\n", - " iys=jnp.array(iys),\n", - " iy_trafos=jnp.array(iy_trafos),\n", - " ops=jnp.array(ops),\n", - " nps=jnp.array(nps),\n", - " solver=diffrax.Kvaerno5(),\n", - " controller=diffrax.PIDController(atol=1e-8, rtol=1e-8),\n", - " root_finder=optimistix.Newton(atol=1e-8, rtol=1e-8),\n", - " steady_state_event=diffrax.steady_state_event(),\n", - " max_steps=2**10,\n", - " adjoint=diffrax.DirectAdjoint(),\n", - " ret=ReturnValue.y, # Return observables\n", - " )[0]\n", - "\n", - "\n", - "# Compute the gradient with respect to `ts_dyn`\n", - "g = grad_ts_dyn(ts_dyn)\n", - "g" - ] + "source": "import diffrax\nimport jax.numpy as jnp\nimport optimistix\nfrom amici.sim.jax import ReturnValue\n\n# Define the simulation condition\nexperiment_condition = \"__default__\"\nic = 0\n\n# Load condition-specific data\nts_dyn = jax_problem._ts_dyn[ic, :]\nts_posteq = jax_problem._ts_posteq[ic, :]\nmy = jax_problem._my[ic, :]\niys = jax_problem._iys[ic, :]\niy_trafos = jax_problem._iy_trafos[ic, :]\nops = jax_problem._op_numeric[ic, :]\nnps = jax_problem._np_numeric[ic, :]\n\n# Load parameters for the specified condition\np = jax_problem.load_model_parameters(\n jax_problem._petab_problem.experiments[0], is_preeq=False\n)\n\n\n# Define a function to compute the gradient with respect to dynamic timepoints\n@eqx.filter_jacfwd\ndef grad_ts_dyn(tt):\n return jax_problem.model.simulate_condition(\n p=p[None, :],\n ts_dyn=tt,\n ts_posteq=ts_posteq,\n my=jnp.array(my),\n iys=jnp.array(iys),\n iy_trafos=jnp.array(iy_trafos),\n ops=jnp.array(ops),\n nps=jnp.array(nps),\n solver=diffrax.Kvaerno5(),\n controller=diffrax.PIDController(atol=1e-8, rtol=1e-8),\n root_finder=optimistix.Newton(atol=1e-8, rtol=1e-8),\n steady_state_event=diffrax.steady_state_event(),\n max_steps=2**10,\n adjoint=diffrax.DirectAdjoint(),\n ret=ReturnValue.y, # Return observables\n )[0]\n\n\n# Compute the gradient with respect to `ts_dyn`\ng = grad_ts_dyn(ts_dyn)\ng" }, { "cell_type": "markdown", @@ -690,4 +640,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file From 3e58c657dbb917d0641a5a2734441ab1ee7ec573 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 5 Jul 2026 07:36:01 +0000 Subject: [PATCH 09/22] Address PEtab v2 JAX review feedback: override parsing, caching, event 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. --- python/sdist/amici/sim/jax/petab.py | 36 +++++--- .../tests/petab_/test_petab_v2_multiperiod.py | 83 +++++++++++++++++++ 2 files changed, 106 insertions(+), 13 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index ac0e3ffdab..54ab40ace9 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -122,14 +122,23 @@ def _split_override_column( ids, resolving non-estimated parameter references to their nominal value and right-padding with ``1.0``.""" - def resolve_row(entry: list | float) -> list: - if isinstance(entry, list): - return [_resolve_override_symbol(v, parameter_df) for v in entry] - return [] if pd.isna(entry) else [entry] + def resolve_row(entry) -> list: + # `col_values` may have `object` dtype with a mix of `;`-separated + # override strings and already-numeric entries (e.g. a column + # where only some rows use string-encoded overrides); splitting + # via the `.str` accessor on the whole column would silently turn + # every non-string entry into NaN; handle each entry's type here + # instead. + if pd.isna(entry): + return [] + values = ( + entry.split(petabv2.C.PARAMETER_SEPARATOR) + if isinstance(entry, str) + else [entry] + ) + return [_resolve_override_symbol(v, parameter_df) for v in values] - rows = col_values.str.split(petabv2.C.PARAMETER_SEPARATOR).apply( - resolve_row - ) + rows = col_values.apply(resolve_row) padded = rows.apply( lambda row: np.pad( row, (0, n_pars - len(row)), mode="constant", constant_values=1.0 @@ -394,6 +403,7 @@ class JAXProblem(eqx.Module): _petab_measurement_indices: np.ndarray _petab_problem: HybridProblem | petabv2.Problem _unconverted_problem: petabv2.Problem | None + _all_condition_targets: frozenset[str] def __init__( self, @@ -420,6 +430,11 @@ def __init__( self.simulation_conditions = scs.conditionId.to_list() self._petab_problem = _get_hybrid_petab_problem(petab_problem) self._unconverted_problem = unconverted_problem + self._all_condition_targets = frozenset( + change.target_id + for condition in self._petab_problem.conditions + for change in condition.changes + ) self.parameters, self.model = ( self._initialize_model_with_nominal_values(model) ) @@ -1618,13 +1633,8 @@ def load_reinitialisation( if isinstance(condition_ids, str): condition_ids = [condition_ids] - all_condition_targets = { - change.target_id - for condition in self._petab_problem.conditions - for change in condition.changes - } has_reinitialisable_states = any( - x_id in all_condition_targets + x_id in self._all_condition_targets or hasattr(self, "nn_output_ids") and x_id in self._parameter_mappings["hybrid_map"] for x_id in self.model.state_ids diff --git a/python/tests/petab_/test_petab_v2_multiperiod.py b/python/tests/petab_/test_petab_v2_multiperiod.py index ea28d59a90..22ee6a87bd 100644 --- a/python/tests/petab_/test_petab_v2_multiperiod.py +++ b/python/tests/petab_/test_petab_v2_multiperiod.py @@ -213,6 +213,89 @@ def y_and_dy_dk(t: float) -> tuple[float, float]: np.testing.assert_allclose(float(grad[0]), expected_grad, rtol=1e-4) +def _threshold_piecewise_decay_problem() -> Problem: + """A single-species model whose decay rate depends on whether the + species concentration is above or below a threshold, via a + ``piecewise`` rate law. The JAX backend compiles this to a + root-finding/heaviside event, rather than a state-assignment event.""" + problem = Problem() + problem.config = ProblemConfig() + problem.model = SbmlModel.from_antimony( + "xx = 1; kfast = 1.0; " + "xx' = piecewise(-kfast*xx, xx > 2, -0.1*xx);" + ) + problem.add_observable("obs1", "xx", noise_formula="1") + return problem + + +def test_event_heaviside_state_not_reevaluated_after_period_reinit( + tmp_path, +): + """Documents a known limitation of the native per-period chaining: the + heaviside/event state ``h`` from the *end* of one period is carried + unconditionally into ``JAXModel._handle_t0_event`` for the next period + (mirroring the existing pre-equilibration -> main-period handoff), + instead of being re-evaluated against the *actual* (possibly + reinitialised) state at the new period's t0. When a reinitialisation + crosses the threshold of a ``piecewise`` rate law, the branch selected + for the whole following period is determined by where the *previous* + period ended, not by the reinitialised state -- until/unless the ODE + integrator happens to cross the threshold again during that period. + + This pins the *current* behaviour; it does not assert that behaviour + is analytically "correct" for this scenario (see the review discussion + on PR #3198 re: ``JAXModel._handle_t0_event``). + """ + problem = _threshold_piecewise_decay_problem() + # period 1: xx starts at 5 (above the threshold of 2) and decays past + # it, ending (at t=1) below the threshold. + problem.add_condition("cond_p1", xx=5.0) + # period 2 reinitialises xx to 3 (back above the threshold), but the + # heaviside carried over from period 1's end still says "below". + problem.add_condition("cond_p2", xx=3.0) + problem.add_experiment("exp1", 0.0, "cond_p1", 1.0, "cond_p2") + measurement_times = (0.3, 0.8, 1.3, 1.8) + for t in measurement_times: + problem.add_measurement( + "obs1", time=t, measurement=0.0, experiment_id="exp1" + ) + + jax_problem = _import_jax( + problem, "test_event_reinit_heaviside_carryover", tmp_path + ) + assert jax_problem._max_periods == 2 + + x, _ = run_simulations(jax_problem, ret=ReturnValue.x) + ts_mask = np.asarray(jax_problem._ts_masks)[0].reshape(-1) + actual = np.asarray(x)[0].reshape(-1)[ts_mask] + + # period 1 correctly crosses the threshold via root-finding during + # integration (heaviside starts "above" since xx=5 > 2 at t=0). + t_cross = np.log(5.0 / 2.0) + + def period1(t): + if t < t_cross: + return 5.0 * np.exp(-1.0 * t) + return 2.0 * np.exp(-0.1 * (t - t_cross)) + + def period2_current_behavior(t_local): + # the carried-over heaviside (still "below threshold", from + # period 1's end) is reused for the entire period, so the + # fast-decay branch is never selected here even though the + # reinitialised xx=3.0 is above the threshold. + return 3.0 * np.exp(-0.1 * t_local) + + expected = np.array( + [ + period1(0.3), + period1(0.8), + period2_current_behavior(0.3), + period2_current_behavior(0.8), + ] + ) + np.testing.assert_allclose(actual, expected, rtol=1e-4) + + @pytest.mark.parametrize("jax_flag", [True]) def test_petab_importer_skips_event_conversion_for_jax(tmp_path, jax_flag): """The JAX backend must not run ExperimentsToSbmlConverter (the From 2932169975a95206e3e642f7f85c85b4b99debfd Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 5 Jul 2026 08:24:15 +0000 Subject: [PATCH 10/22] Re-evaluate heaviside/event state at every period boundary in the JAX 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. --- python/sdist/amici/sim/jax/model.py | 27 +++++--- .../tests/petab_/test_petab_v2_multiperiod.py | 62 +++++++++---------- 2 files changed, 50 insertions(+), 39 deletions(-) diff --git a/python/sdist/amici/sim/jax/model.py b/python/sdist/amici/sim/jax/model.py index 99eadb2f67..40c998fd06 100644 --- a/python/sdist/amici/sim/jax/model.py +++ b/python/sdist/amici/sim/jax/model.py @@ -1123,17 +1123,28 @@ def _handle_t0_event( root_cond_fn: Callable, delta_x: Callable, h_mask: jt.Bool[jt.Array, "ne"], - h_preeq: jt.Bool[jt.Array, "ne"], + h_prev: jt.Bool[jt.Array, "ne"], stats: dict, ): - y0 = y0_next.copy() - rf0 = self.event_initial_values - 0.5 - - if h_preeq.shape[0]: - # return immediately because preequilibration is equivalent to handling t0 event? - return y0, t0_next, h_preeq, stats + if h_prev.shape[0]: + # `h_prev` is the heaviside state at wherever `y0_next` came + # from (a preceding preequilibration, or the end of the + # previous experiment period), used here only as the + # pre-transition reference value. It is not necessarily the + # trigger state *at* `(t0_next, y0_next)`: a period boundary + # (or the reinitialisation applied after preequilibration) may + # have crossed an event's trigger threshold without the ODE + # integrator ever seeing it, so the trigger condition is always + # re-evaluated below against the actual incoming state, exactly + # as for a genuine t=0. + h = h_prev else: - h = jnp.where(h_mask, jnp.heaviside(rf0, 0.0), jnp.ones_like(rf0)) + h = jnp.where( + h_mask, + jnp.heaviside(self.event_initial_values - 0.5, 0.0), + jnp.ones_like(self.event_initial_values), + ) + rf0 = h - 0.5 args = (p, tcl, h) rfx = root_cond_fn(t0_next, y0_next, args) roots_dir = jnp.sign(rfx - rf0) diff --git a/python/tests/petab_/test_petab_v2_multiperiod.py b/python/tests/petab_/test_petab_v2_multiperiod.py index 22ee6a87bd..a946095adf 100644 --- a/python/tests/petab_/test_petab_v2_multiperiod.py +++ b/python/tests/petab_/test_petab_v2_multiperiod.py @@ -228,30 +228,27 @@ def _threshold_piecewise_decay_problem() -> Problem: return problem -def test_event_heaviside_state_not_reevaluated_after_period_reinit( +def test_event_heaviside_state_reevaluated_after_period_reinit( tmp_path, ): - """Documents a known limitation of the native per-period chaining: the - heaviside/event state ``h`` from the *end* of one period is carried - unconditionally into ``JAXModel._handle_t0_event`` for the next period - (mirroring the existing pre-equilibration -> main-period handoff), - instead of being re-evaluated against the *actual* (possibly - reinitialised) state at the new period's t0. When a reinitialisation - crosses the threshold of a ``piecewise`` rate law, the branch selected - for the whole following period is determined by where the *previous* - period ended, not by the reinitialised state -- until/unless the ODE - integrator happens to cross the threshold again during that period. - - This pins the *current* behaviour; it does not assert that behaviour - is analytically "correct" for this scenario (see the review discussion - on PR #3198 re: ``JAXModel._handle_t0_event``). + """A period-boundary reinitialisation that crosses the threshold of a + ``piecewise`` rate law must select the branch matching the + *reinitialised* state at the new period's t0, not whatever heaviside + state the previous period happened to end with. + + ``JAXModel._handle_t0_event`` re-evaluates the trigger condition + against the actual incoming state at every period boundary (and after + preequilibration), using the previous heaviside state only as the + pre-transition reference for detecting a crossing -- it does not carry + it over unconditionally. """ problem = _threshold_piecewise_decay_problem() # period 1: xx starts at 5 (above the threshold of 2) and decays past # it, ending (at t=1) below the threshold. problem.add_condition("cond_p1", xx=5.0) - # period 2 reinitialises xx to 3 (back above the threshold), but the - # heaviside carried over from period 1's end still says "below". + # period 2 reinitialises xx to 3, back above the threshold: the + # trigger must be re-evaluated at t0 of period 2 rather than reusing + # period 1's ending ("below threshold") heaviside state. problem.add_condition("cond_p2", xx=3.0) problem.add_experiment("exp1", 0.0, "cond_p1", 1.0, "cond_p2") measurement_times = (0.3, 0.8, 1.3, 1.8) @@ -261,7 +258,7 @@ def test_event_heaviside_state_not_reevaluated_after_period_reinit( ) jax_problem = _import_jax( - problem, "test_event_reinit_heaviside_carryover", tmp_path + problem, "test_event_reinit_heaviside_reevaluated", tmp_path ) assert jax_problem._max_periods == 2 @@ -269,28 +266,31 @@ def test_event_heaviside_state_not_reevaluated_after_period_reinit( ts_mask = np.asarray(jax_problem._ts_masks)[0].reshape(-1) actual = np.asarray(x)[0].reshape(-1)[ts_mask] - # period 1 correctly crosses the threshold via root-finding during - # integration (heaviside starts "above" since xx=5 > 2 at t=0). - t_cross = np.log(5.0 / 2.0) + # period 1 crosses the threshold via root-finding during integration + # (heaviside starts "above" since xx=5 > 2 at t=0). + t_cross1 = np.log(5.0 / 2.0) def period1(t): - if t < t_cross: + if t < t_cross1: return 5.0 * np.exp(-1.0 * t) - return 2.0 * np.exp(-0.1 * (t - t_cross)) + return 2.0 * np.exp(-0.1 * (t - t_cross1)) - def period2_current_behavior(t_local): - # the carried-over heaviside (still "below threshold", from - # period 1's end) is reused for the entire period, so the - # fast-decay branch is never selected here even though the - # reinitialised xx=3.0 is above the threshold. - return 3.0 * np.exp(-0.1 * t_local) + # period 2 must also start "above" (heaviside re-evaluated against the + # reinitialised xx=3.0 at period 2's t0), decaying fast until it + # crosses the threshold again, then slow. + t_cross2 = np.log(3.0 / 2.0) + + def period2(t_local): + if t_local < t_cross2: + return 3.0 * np.exp(-1.0 * t_local) + return 2.0 * np.exp(-0.1 * (t_local - t_cross2)) expected = np.array( [ period1(0.3), period1(0.8), - period2_current_behavior(0.3), - period2_current_behavior(0.8), + period2(0.3), + period2(0.8), ] ) np.testing.assert_allclose(actual, expected, rtol=1e-4) From 414783b8a679e55dd5a93319b46c28f20ba36e42 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 10:02:11 +0000 Subject: [PATCH 11/22] Fix per-experiment preequilibration reinit condition mismatch and stale 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. --- python/sdist/amici/sim/jax/petab.py | 67 ++++++++++++++++------------- python/tests/test_jax.py | 16 +++---- 2 files changed, 45 insertions(+), 38 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index a3af2dd4b4..29dea6abc1 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -178,7 +178,10 @@ def resolve_row(entry) -> list: if pd.isna(entry): return [] values = ( - entry.split(petabv2.C.PARAMETER_SEPARATOR) + # an empty cell (as opposed to a missing/NaN one) splits to a + # single empty-string token; drop it rather than trying to + # resolve/convert "" as if it were a real override + [v for v in entry.split(petabv2.C.PARAMETER_SEPARATOR) if v] if isinstance(entry, str) else [entry] ) @@ -1710,7 +1713,6 @@ def update_parameters(self, p: jt.Float[jt.Array, "np"]) -> "JAXProblem": def _prepare_experiments( self, experiments: list[petabv2.Experiment], - conditions: list[str], is_preeq: bool, op_numeric: np.ndarray | None = None, op_mask: np.ndarray | None = None, @@ -1740,10 +1742,6 @@ def _prepare_experiments( :param experiments: Experiments to prepare simulation arrays for. - :param conditions: - Simulation conditions to prepare. Only used for - ``is_preeq=True``, where it is one (pre-equilibration) condition - id per experiment. :param is_preeq: Whether to load preequilibration or simulation parameters. :param op_numeric: @@ -1775,7 +1773,16 @@ def _prepare_experiments( for exp in experiments ] ) - reinit_condition_ids = conditions + + def preeq_condition_ids(exp: petabv2.Experiment) -> list[str]: + for period in exp.sorted_periods: + if period.is_preequilibration: + return period.condition_ids + return [] + + reinit_condition_ids = [ + preeq_condition_ids(exp) for exp in experiments + ] else: p_array = jnp.stack( [ @@ -2098,7 +2105,6 @@ def run_simulations( t_zeros, ) = self._prepare_experiments( experiments, - [], False, self._op_numeric, self._op_mask, @@ -2228,14 +2234,8 @@ def run_preequilibrations( ], max_steps: jnp.int_, ): - preequilibration_conditions = list( - _get_preequilibration_condition_ids(experiments) - ) - p_array, mask_reinit_array, x_reinit_array, _, _, h_mask, _ = ( - self._prepare_experiments( - experiments, preequilibration_conditions, True, None, None - ) + self._prepare_experiments(experiments, True, None, None) ) return self.run_preequilibration( p_array, @@ -2606,21 +2606,28 @@ def _build_simulation_df_v2(problem, y, dyn_conditions): }, index=problem._petab_measurement_indices[exp_idx, period_idx, mask], ) - if ( - petabv2.C.OBSERVABLE_PARAMETERS - in problem._petab_problem.measurement_df - ): - df_sc[petabv2.C.OBSERVABLE_PARAMETERS] = ( - problem._petab_problem.measurement_df.query( - f"{petabv2.C.EXPERIMENT_ID} == '{experiment_id}'" - )[petabv2.C.OBSERVABLE_PARAMETERS] - ) - if petabv2.C.NOISE_PARAMETERS in problem._petab_problem.measurement_df: - df_sc[petabv2.C.NOISE_PARAMETERS] = ( - problem._petab_problem.measurement_df.query( - f"{petabv2.C.EXPERIMENT_ID} == '{experiment_id}'" - )[petabv2.C.NOISE_PARAMETERS] - ) + measurement_df = problem._petab_problem.measurement_df + # `experiment_id` may be `jnp.nan` (the "__default__" experiment + # sentinel): a string `.query()` (`== 'nan'`) does not match actual + # NaN values in the column, silently selecting zero rows and + # leaving the assigned column all-NaN below. Select via a boolean + # mask instead, which handles both the NaN and real-id cases. + if isinstance(experiment_id, float): # NaN sentinel + exp_rows = measurement_df[ + measurement_df[petabv2.C.EXPERIMENT_ID].isna() + ] + else: + exp_rows = measurement_df[ + measurement_df[petabv2.C.EXPERIMENT_ID] == experiment_id + ] + if petabv2.C.OBSERVABLE_PARAMETERS in measurement_df: + df_sc[petabv2.C.OBSERVABLE_PARAMETERS] = exp_rows[ + petabv2.C.OBSERVABLE_PARAMETERS + ] + if petabv2.C.NOISE_PARAMETERS in measurement_df: + df_sc[petabv2.C.NOISE_PARAMETERS] = exp_rows[ + petabv2.C.NOISE_PARAMETERS + ] dfs.append(df_sc) return pd.concat(dfs).sort_index() diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 545a58f81b..47920df435 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -575,14 +575,14 @@ def dispatch(self, *args, **kwargs): for k_val in conditions: kwargs = fresh_solver_kwargs() model.simulate_condition( - jnp.array([k_val]), - ts, - jnp.array([]), - my, - iys, - iy_trafos, - jnp.zeros((3, 0)), - jnp.zeros((3, 0)), + jnp.array([[k_val]]), + ts[None, :], + jnp.zeros((1, 0)), + my[None, :], + iys[None, :], + iy_trafos[None, :], + jnp.zeros((1, 3, 0)), + jnp.zeros((1, 3, 0)), kwargs["solver"], kwargs["controller"], kwargs["root_finder"], From 81d44901518b5de964f5c2104caca5c2050bda49 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 11:03:46 +0000 Subject: [PATCH 12/22] Fix per-measurement-row observable-transformation length mismatch and 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. --- python/sdist/amici/sim/jax/petab.py | 59 ++++++++++++++-------- tests/petab_test_suite/test_petab_suite.py | 9 +++- 2 files changed, 47 insertions(+), 21 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 29dea6abc1..68651d6e63 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -258,19 +258,34 @@ def _get_overrides( def _get_iy_trafos( - iys: np.ndarray, petab_problem: petabv2.Problem + iys: np.ndarray, + petab_problem: petabv2.Problem, + observable_ids: list[str], ) -> np.ndarray: """Observable transformation index (see ``SCALE_TO_INT``) for each - observable index in ``iys``.""" + (per-measurement-row) observable index in ``iys``. + + ``iys`` indexes into ``observable_ids`` (i.e. ``self.model.observable_ids``), + not into ``petab_problem.observables`` directly, so the per-observable + transformation is first resolved by id, then gathered onto ``iys``'s own + (measurement-row) length -- returning one entry per row in + ``petab_problem.observables`` order would silently produce an array of + the wrong length whenever the number of measurement rows differs from + the number of observables. + """ if petabv2.C.NOISE_DISTRIBUTION in petab_problem.observable_df: - return np.array( - [ + trafo_by_id = { + obs.id: ( SCALE_TO_INT[petabv2.C.LOG] if obs.noise_distribution == petabv2.C.LOG_NORMAL else SCALE_TO_INT[petabv2.C.LIN] - for obs in petab_problem.observables - ] + ) + for obs in petab_problem.observables + } + trafo_by_index = np.array( + [trafo_by_id[oid] for oid in observable_ids] ) + return trafo_by_index[iys] return np.zeros_like(iys) @@ -647,7 +662,9 @@ def _get_measurements( ] ) iy_trafos_real = _get_iy_trafos( - iys_real, self._petab_problem + iys_real, + self._petab_problem, + self.model.observable_ids, ) overrides_real = _get_overrides( m, n_pars, fixed_parameter_values, self.parameter_ids @@ -2607,19 +2624,21 @@ def _build_simulation_df_v2(problem, y, dyn_conditions): index=problem._petab_measurement_indices[exp_idx, period_idx, mask], ) measurement_df = problem._petab_problem.measurement_df - # `experiment_id` may be `jnp.nan` (the "__default__" experiment - # sentinel): a string `.query()` (`== 'nan'`) does not match actual - # NaN values in the column, silently selecting zero rows and - # leaving the assigned column all-NaN below. Select via a boolean - # mask instead, which handles both the NaN and real-id cases. - if isinstance(experiment_id, float): # NaN sentinel - exp_rows = measurement_df[ - measurement_df[petabv2.C.EXPERIMENT_ID].isna() - ] - else: - exp_rows = measurement_df[ - measurement_df[petabv2.C.EXPERIMENT_ID] == experiment_id - ] + # `experiment_id` is coerced to `jnp.nan` above for the "__default__" + # experiment sentinel, but `measurement_df`'s own `experimentId` + # column still stores the literal string `"__default__"` (never a + # real NaN): a string `.query()` (`== 'nan'`) or an `.isna()` mask + # both silently select zero rows in that case, leaving the assigned + # column all-NaN below. Match against the literal sentinel string + # instead. + match_id = ( + "__default__" + if isinstance(experiment_id, float) + else experiment_id + ) + exp_rows = measurement_df[ + measurement_df[petabv2.C.EXPERIMENT_ID] == match_id + ] if petabv2.C.OBSERVABLE_PARAMETERS in measurement_df: df_sc[petabv2.C.OBSERVABLE_PARAMETERS] = exp_rows[ petabv2.C.OBSERVABLE_PARAMETERS diff --git a/tests/petab_test_suite/test_petab_suite.py b/tests/petab_test_suite/test_petab_suite.py index 152ea0511a..7dcf2a1783 100755 --- a/tests/petab_test_suite/test_petab_suite.py +++ b/tests/petab_test_suite/test_petab_suite.py @@ -7,6 +7,7 @@ import diffrax import pandas as pd import petab.v1 as petab +import petab.v2 as petabv2 import petabtests import pytest from _pytest.outcomes import Skipped @@ -143,12 +144,18 @@ def _test_case(case, model_type, version, jax): # `"__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. + # recover the original column by index alignment. Drop `experimentId` + # afterwards: `petabtests.evaluate_simulations` determines the PEtab + # version from column presence and errors out if both id columns + # are present at once. simulation_df[petab.SIMULATION_CONDITION_ID] = ( problem.measurement_df.loc[ simulation_df.index, petab.SIMULATION_CONDITION_ID ].values ) + simulation_df = simulation_df.drop( + columns=[petabv2.C.EXPERIMENT_ID] + ) else: model = imported # import_petab_problem returns Model when jax=False solver = model.create_solver() From 8491f6b10ae97e4df380a1fb5997bcc53e782ac0 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 11:09:49 +0000 Subject: [PATCH 13/22] Gitignore the JAX SBML test suite's generated model directory 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. --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 75f7ba71a5..1eba24a059 100644 --- a/.gitignore +++ b/.gitignore @@ -141,6 +141,7 @@ tests/sbml/sbml-test-suite/* tests/sbml/sbml-test-suite/ */sbml-semantic-test-cases/* tests/sbml/SBMLTestModels/ +tests/sbml/SBMLTestModelsJax/ tests/benchmark_models/test_bmc */tests/BIOMD0000000529/* From f9bdca75c1c85d60212a8d5f95632e1c53b16f27 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 11:44:50 +0000 Subject: [PATCH 14/22] Fix integer dtype for empty per-period observable-index arrays 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). --- python/sdist/amici/sim/jax/petab.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 68651d6e63..286594dc2f 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -659,7 +659,8 @@ def _get_measurements( [ self.model.observable_ids.index(oid) for oid in m[petabv2.C.OBSERVABLE_ID].values - ] + ], + dtype=int, ) iy_trafos_real = _get_iy_trafos( iys_real, From c4841dc401b6e8a7fb685b175a515f70dc3f9c49 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 12:05:18 +0000 Subject: [PATCH 15/22] Fix missing observable/override data for post-equilibrium measurements _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. --- python/sdist/amici/sim/jax/petab.py | 64 +++++++++++++++++++++++++---- 1 file changed, 56 insertions(+), 8 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 286594dc2f..7fd3570acb 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -678,16 +678,52 @@ def _get_measurements( ts_posteq = m_posteq[petabv2.C.TIME].values index_posteq = list(m_posteq.index) + # Post-equilibrium measurements (e.g. steady-state + # comparisons) have their own observable/override + # data, distinct from the dynamic-phase measurements + # above; this must be concatenated onto the dyn + # portion so that `mv.my`/`iys`/`iy_trafos`/overrides + # cover the full `len(ts_dyn) + len(ts_posteq)` range + # that `_pad_and_stack` later splits at + # `len(mv.ts_dyn)` -- otherwise post-equilibrium rows + # silently fall back to zero-filled placeholders + # (wrong measurement value, wrong observable index, + # wrong noise override) while still being marked + # valid. + iys_posteq = np.array( + [ + self.model.observable_ids.index(oid) + for oid in m_posteq[ + petabv2.C.OBSERVABLE_ID + ].values + ], + dtype=int, + ) + my_posteq = m_posteq[petabv2.C.MEASUREMENT].values + iy_trafos_posteq = _get_iy_trafos( + iys_posteq, + self._petab_problem, + self.model.observable_ids, + ) + overrides_posteq = _get_overrides( + m_posteq, + n_pars, + fixed_parameter_values, + self.parameter_ids, + ) + # No further period to hand off to within this # experiment's own chain; padding slots (if any) # after this one are pure no-ops. if len(ts_dyn_real): ts_dyn = ts_dyn_real dyn_valid = np.ones(len(ts_dyn_real), dtype=bool) - my = my_real - iys = iys_real - iy_trafos = iy_trafos_real - overrides = overrides_real + my_dyn, iys_dyn, iy_trafos_dyn = ( + my_real, + iys_real, + iy_trafos_real, + ) + overrides_dyn = overrides_real index_dyn_full = index_dyn else: # No real dyn measurements in this period (e.g. @@ -698,12 +734,24 @@ def _get_measurements( ( ts_dyn, dyn_valid, - my, - iys, - iy_trafos, - overrides, + my_dyn, + iys_dyn, + iy_trafos_dyn, + overrides_dyn, ) = _masked_placeholder_period(t_lo, n_pars) index_dyn_full = [-1] + + my = np.concatenate([my_dyn, my_posteq]) + iys = np.concatenate([iys_dyn, iys_posteq]) + iy_trafos = np.concatenate( + [iy_trafos_dyn, iy_trafos_posteq] + ) + overrides = { + col: OverrideColumn.concatenate( + overrides_dyn[col], overrides_posteq[col] + ) + for col in overrides_dyn + } else: ts_posteq = np.array([]) index_posteq = [] From 82dba9ff1b48683b9816b2132cf833b7692a28b5 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 18:41:22 +0000 Subject: [PATCH 16/22] Rename JAXModel.simulate_condition[_unjitted] to simulate_experiment, 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. --- .gitignore | 6 + .../__init__.py | 178 ------------------ 1.0.1/test_noise_params_jax_jax/__init__.py | 178 ------------------ .../example_jax_petab/ExampleJaxPEtab.ipynb | 56 +++++- python/sdist/amici/sim/jax/model.py | 12 +- python/sdist/amici/sim/jax/petab.py | 8 +- python/tests/test_jax.py | 10 +- tests/performance/test_jax_regression.py | 16 +- tests/sbml/testSBMLSuite.py | 2 +- tests/sbml/testSBMLSuiteJax.py | 4 +- 10 files changed, 85 insertions(+), 385 deletions(-) delete mode 100644 1.0.1/test_arbitrary_placeholder_jax_jax/__init__.py delete mode 100644 1.0.1/test_noise_params_jax_jax/__init__.py diff --git a/.gitignore b/.gitignore index 1eba24a059..cc260cab9f 100644 --- a/.gitignore +++ b/.gitignore @@ -39,6 +39,12 @@ models/model_calvetti/build/* amici_models/ +# PetabImporter-generated test models: python/tests/conftest.py points +# AMICI_MODELS_ROOT at the repo root for the test session, so these land +# directly in a bare / directory instead of under +# amici_models//. +/[0-9]*.[0-9]*.[0-9]*/ + # PEtab SciML test suite (downloaded dynamically) tests/sciml/testsuite/ diff --git a/1.0.1/test_arbitrary_placeholder_jax_jax/__init__.py b/1.0.1/test_arbitrary_placeholder_jax_jax/__init__.py deleted file mode 100644 index 75093f0904..0000000000 --- a/1.0.1/test_arbitrary_placeholder_jax_jax/__init__.py +++ /dev/null @@ -1,178 +0,0 @@ -# ruff: noqa: F401, F821, F841 -from pathlib import Path - -import equinox as eqx # noqa: F401 -import jax.numpy as jnp -import jax.random as jr # noqa: F401 -import jaxtyping as jt # noqa: F401 -from jax.numpy import inf as oo # noqa: F401 -from jax.numpy import nan as nan # noqa: F401 - -from amici import _module_from_path # noqa: F401 -from amici.sim.jax.model import JAXModel, safe_div, safe_log # noqa: F401 - - - - -class JAXModel_test_arbitrary_placeholder_jax(JAXModel): - api_version = '0.0.4' - - def __init__(self): - self.jax_py_file = Path(__file__).resolve() - self.nns = {} - self.parameters = jnp.array([]) - self._array_inputs = {} - self._array_input_index = jnp.int32(0) - super().__init__() - - def _xdot(self, t, x, args): - p, tcl, h = args - - xa, xb, = x - _ = p - _ = tcl - _ = h - _ = self._w(t, x, p, tcl, h) - - dxadt = safe_div(-xa, 10) - dxbdt = safe_div(xb, 20) - - return jnp.array([dxadt, dxbdt]) - - def _w(self, t, x, p, tcl, h): - xa, xb, = x - _ = p - _ = tcl - _ = h - - - - return jnp.array([]) - - def _x0(self, t, p): - _ = p - - x00 = 1.00000000000000 - x01 = 2.00000000000000 - - return jnp.array([x00, x01]) - - def _x_solver(self, x): - xa, xb, = x - - x_solver0 = xa - x_solver1 = xb - - return jnp.array([x_solver0, x_solver1]) - - def _x_rdata(self, x, tcl): - xa, xb, = x - _ = tcl - - xa = xa - xb = xb - - return jnp.array([xa, xb]) - - def _tcl(self, x, p): - xa, xb, = x - _ = p - - - - return jnp.array([]) - - def _y(self, t, x, p, tcl, h, op): - xa, xb, = x - _ = p - _ = self._w(t, x, p, tcl, h) - _ = op - - obs_a = xa - obs_b = xb - - return jnp.array([obs_a, obs_b]) - - def _sigmay(self, y, p, np): - _ = p - - obs_a, obs_b, = y - noiseParameter1, = np - - sigma_obs_a = noiseParameter1*obs_a - sigma_obs_b = noiseParameter1 - - return jnp.array([sigma_obs_a, sigma_obs_b]) - - def _nllh(self, t, x, p, tcl, h, my, iy, op, np): - y = self._y(t, x, p, tcl, h, op) - if not y.size: - return jnp.array(0.0) - - obs_a, obs_b, = y - sigma_obs_a, sigma_obs_b, = self._sigmay(y, p, np) - - _amici_cse_0 = sigma_obs_a**2 - _amici_cse_1 = 2*jnp.pi - _amici_cse_2 = sigma_obs_b**2 - Jy0 = 0.5*safe_log(_amici_cse_0*_amici_cse_1) + safe_div(0.5*(-my + obs_a)**2, _amici_cse_0) - Jy1 = 0.5*safe_log(_amici_cse_1*_amici_cse_2) + safe_div(0.5*(-my + obs_b)**2, _amici_cse_2) - - return jnp.array([Jy0, Jy1]).at[iy].get() - - def _known_discs(self, p): - _ = p - - return jnp.array([]) - - def _root_cond_fn(self, t, y, args, **_): - p, tcl, h = args - - xa, xb, = y - _ = p - _ = tcl - _ = h - _ = self._w(t, y, p, tcl, h) - - - - - return jnp.hstack((jnp.array([]), jnp.array([]))) - - def _delta_x(self, y, p, tcl): - xa, xb, = y - _ = p - _ = tcl - # FIXME: workaround until state from event time is properly passed - x_old0, x_old1, = y - - - - return jnp.array([]) - - @property - def event_initial_values(self): - return jnp.array([]) - - @property - def n_events(self): - return 0 + 0 - - @property - def observable_ids(self): - return "obs_a", "obs_b", - - @property - def state_ids(self): - return "xa", "xb", - - @property - def parameter_ids(self): - return tuple() - - @property - def expression_ids(self): - return tuple() - - -Model = JAXModel_test_arbitrary_placeholder_jax diff --git a/1.0.1/test_noise_params_jax_jax/__init__.py b/1.0.1/test_noise_params_jax_jax/__init__.py deleted file mode 100644 index b9238c3223..0000000000 --- a/1.0.1/test_noise_params_jax_jax/__init__.py +++ /dev/null @@ -1,178 +0,0 @@ -# ruff: noqa: F401, F821, F841 -from pathlib import Path - -import equinox as eqx # noqa: F401 -import jax.numpy as jnp -import jax.random as jr # noqa: F401 -import jaxtyping as jt # noqa: F401 -from jax.numpy import inf as oo # noqa: F401 -from jax.numpy import nan as nan # noqa: F401 - -from amici import _module_from_path # noqa: F401 -from amici.sim.jax.model import JAXModel, safe_div, safe_log # noqa: F401 - - - - -class JAXModel_test_noise_params_jax(JAXModel): - api_version = '0.0.4' - - def __init__(self): - self.jax_py_file = Path(__file__).resolve() - self.nns = {} - self.parameters = jnp.array([]) - self._array_inputs = {} - self._array_input_index = jnp.int32(0) - super().__init__() - - def _xdot(self, t, x, args): - p, tcl, h = args - - xa, xb, = x - _ = p - _ = tcl - _ = h - _ = self._w(t, x, p, tcl, h) - - dxadt = safe_div(-xa, 10) - dxbdt = safe_div(xb, 5) - - return jnp.array([dxadt, dxbdt]) - - def _w(self, t, x, p, tcl, h): - xa, xb, = x - _ = p - _ = tcl - _ = h - - - - return jnp.array([]) - - def _x0(self, t, p): - _ = p - - x00 = 1.00000000000000 - x01 = 2.00000000000000 - - return jnp.array([x00, x01]) - - def _x_solver(self, x): - xa, xb, = x - - x_solver0 = xa - x_solver1 = xb - - return jnp.array([x_solver0, x_solver1]) - - def _x_rdata(self, x, tcl): - xa, xb, = x - _ = tcl - - xa = xa - xb = xb - - return jnp.array([xa, xb]) - - def _tcl(self, x, p): - xa, xb, = x - _ = p - - - - return jnp.array([]) - - def _y(self, t, x, p, tcl, h, op): - xa, xb, = x - _ = p - _ = self._w(t, x, p, tcl, h) - _ = op - - obsA = xa - obsB = xb - - return jnp.array([obsA, obsB]) - - def _sigmay(self, y, p, np): - _ = p - - obsA, obsB, = y - noiseParameter1, = np - - sigma_obsA = noiseParameter1 - sigma_obsB = noiseParameter1 - - return jnp.array([sigma_obsA, sigma_obsB]) - - def _nllh(self, t, x, p, tcl, h, my, iy, op, np): - y = self._y(t, x, p, tcl, h, op) - if not y.size: - return jnp.array(0.0) - - obsA, obsB, = y - sigma_obsA, sigma_obsB, = self._sigmay(y, p, np) - - _amici_cse_0 = sigma_obsA**2 - _amici_cse_1 = 2*jnp.pi - _amici_cse_2 = sigma_obsB**2 - Jy0 = 0.5*safe_log(_amici_cse_0*_amici_cse_1) + safe_div(0.5*(-my + obsA)**2, _amici_cse_0) - Jy1 = 0.5*safe_log(_amici_cse_1*_amici_cse_2) + safe_div(0.5*(-my + obsB)**2, _amici_cse_2) - - return jnp.array([Jy0, Jy1]).at[iy].get() - - def _known_discs(self, p): - _ = p - - return jnp.array([]) - - def _root_cond_fn(self, t, y, args, **_): - p, tcl, h = args - - xa, xb, = y - _ = p - _ = tcl - _ = h - _ = self._w(t, y, p, tcl, h) - - - - - return jnp.hstack((jnp.array([]), jnp.array([]))) - - def _delta_x(self, y, p, tcl): - xa, xb, = y - _ = p - _ = tcl - # FIXME: workaround until state from event time is properly passed - x_old0, x_old1, = y - - - - return jnp.array([]) - - @property - def event_initial_values(self): - return jnp.array([]) - - @property - def n_events(self): - return 0 + 0 - - @property - def observable_ids(self): - return "obsA", "obsB", - - @property - def state_ids(self): - return "xa", "xb", - - @property - def parameter_ids(self): - return tuple() - - @property - def expression_ids(self): - return tuple() - - -Model = JAXModel_test_noise_params_jax diff --git a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb index 6a92e7cb36..ad827625a7 100644 --- a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb +++ b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb @@ -373,7 +373,7 @@ "id": "58eb04393a1463d", "metadata": {}, "source": [ - "However, we can compute derivatives with respect to data elements using `JAXModel.simulate_condition`. In the example below, we differentiate the observables `y` (specified by passing `y` to the `ret` argument) with respect to the timepoints at which the model outputs are computed after the solving the differential equation. While this might not be particularly practical, it serves as an nice illustration of the power of automatic differentiation." + "However, we can compute derivatives with respect to data elements using `JAXModel.simulate_experiment`. In the example below, we differentiate the observables `y` (specified by passing `y` to the `ret` argument) with respect to the timepoints at which the model outputs are computed after the solving the differential equation. While this might not be particularly practical, it serves as an nice illustration of the power of automatic differentiation." ] }, { @@ -382,7 +382,57 @@ "id": "1a91aff44b93157", "metadata": {}, "outputs": [], - "source": "import diffrax\nimport jax.numpy as jnp\nimport optimistix\nfrom amici.sim.jax import ReturnValue\n\n# Define the simulation condition\nexperiment_condition = \"__default__\"\nic = 0\n\n# Load condition-specific data\nts_dyn = jax_problem._ts_dyn[ic, :]\nts_posteq = jax_problem._ts_posteq[ic, :]\nmy = jax_problem._my[ic, :]\niys = jax_problem._iys[ic, :]\niy_trafos = jax_problem._iy_trafos[ic, :]\nops = jax_problem._op_numeric[ic, :]\nnps = jax_problem._np_numeric[ic, :]\n\n# Load parameters for the specified condition\np = jax_problem.load_model_parameters(\n jax_problem._petab_problem.experiments[0], is_preeq=False\n)\n\n\n# Define a function to compute the gradient with respect to dynamic timepoints\n@eqx.filter_jacfwd\ndef grad_ts_dyn(tt):\n return jax_problem.model.simulate_condition(\n p=p[None, :],\n ts_dyn=tt,\n ts_posteq=ts_posteq,\n my=jnp.array(my),\n iys=jnp.array(iys),\n iy_trafos=jnp.array(iy_trafos),\n ops=jnp.array(ops),\n nps=jnp.array(nps),\n solver=diffrax.Kvaerno5(),\n controller=diffrax.PIDController(atol=1e-8, rtol=1e-8),\n root_finder=optimistix.Newton(atol=1e-8, rtol=1e-8),\n steady_state_event=diffrax.steady_state_event(),\n max_steps=2**10,\n adjoint=diffrax.DirectAdjoint(),\n ret=ReturnValue.y, # Return observables\n )[0]\n\n\n# Compute the gradient with respect to `ts_dyn`\ng = grad_ts_dyn(ts_dyn)\ng" + "source": [ + "import diffrax\n", + "import jax.numpy as jnp\n", + "import optimistix\n", + "from amici.sim.jax import ReturnValue\n", + "\n", + "# Define the simulation condition\n", + "experiment_condition = \"__default__\"\n", + "ic = 0\n", + "\n", + "# Load condition-specific data\n", + "ts_dyn = jax_problem._ts_dyn[ic, :]\n", + "ts_posteq = jax_problem._ts_posteq[ic, :]\n", + "my = jax_problem._my[ic, :]\n", + "iys = jax_problem._iys[ic, :]\n", + "iy_trafos = jax_problem._iy_trafos[ic, :]\n", + "ops = jax_problem._op_numeric[ic, :]\n", + "nps = jax_problem._np_numeric[ic, :]\n", + "\n", + "# Load parameters for the specified condition\n", + "p = jax_problem.load_model_parameters(\n", + " jax_problem._petab_problem.experiments[0], is_preeq=False\n", + ")\n", + "\n", + "\n", + "# Define a function to compute the gradient with respect to dynamic timepoints\n", + "@eqx.filter_jacfwd\n", + "def grad_ts_dyn(tt):\n", + " return jax_problem.model.simulate_experiment(\n", + " p=p[None, :],\n", + " ts_dyn=tt,\n", + " ts_posteq=ts_posteq,\n", + " my=jnp.array(my),\n", + " iys=jnp.array(iys),\n", + " iy_trafos=jnp.array(iy_trafos),\n", + " ops=jnp.array(ops),\n", + " nps=jnp.array(nps),\n", + " solver=diffrax.Kvaerno5(),\n", + " controller=diffrax.PIDController(atol=1e-8, rtol=1e-8),\n", + " root_finder=optimistix.Newton(atol=1e-8, rtol=1e-8),\n", + " steady_state_event=diffrax.steady_state_event(),\n", + " max_steps=2**10,\n", + " adjoint=diffrax.DirectAdjoint(),\n", + " ret=ReturnValue.y, # Return observables\n", + " )[0]\n", + "\n", + "\n", + "# Compute the gradient with respect to `ts_dyn`\n", + "g = grad_ts_dyn(ts_dyn)\n", + "g" + ] }, { "cell_type": "markdown", @@ -640,4 +690,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/python/sdist/amici/sim/jax/model.py b/python/sdist/amici/sim/jax/model.py index db09dc3291..c64a015f85 100644 --- a/python/sdist/amici/sim/jax/model.py +++ b/python/sdist/amici/sim/jax/model.py @@ -674,7 +674,7 @@ def _simulate_period( return ts, xs, hs, x_solver, h, stats_dyn, stats_posteq - def simulate_condition_unjitted( + def simulate_experiment_unjitted( self, p: jt.Float[jt.Array, "P np"], ts_dyn: jt.Float[jt.Array, "P nt_dyn"], @@ -704,7 +704,7 @@ def simulate_condition_unjitted( ret: ReturnValue = ReturnValue.llh, ) -> tuple[jt.Float[jt.Array, "*nt"], dict]: """ - Unjitted version of simulate_condition. + Unjitted version of simulate_experiment. Chains one ODE integration per experiment period (the leading axis, of static size ``P``, of ``p``/``ts_dyn``/``ts_posteq``/``my``/ @@ -714,7 +714,7 @@ def simulate_condition_unjitted( lieu of encoding period transitions as model events. ``P == 1`` reduces to a single, non-chained simulation. - See :meth:`simulate_condition` for full documentation. + See :meth:`simulate_experiment` for full documentation. """ n_periods = p.shape[0] @@ -908,7 +908,7 @@ def simulate_condition_unjitted( return output, stats @eqx.filter_jit - def simulate_condition( + def simulate_experiment( self, p: jt.Float[jt.Array, "P np"], ts_dyn: jt.Float[jt.Array, "P nt_dyn"], @@ -941,7 +941,7 @@ def simulate_condition( Simulate a condition (JIT-compiled version). This is the JIT-compiled version for optimal performance. For runtime type checking - with beartype, use :meth:`simulate_condition_unjitted` instead. + with beartype, use :meth:`simulate_experiment_unjitted` instead. Chains one ODE integration per experiment period (the leading axis, of static size ``P``, of ``p``/``ts_dyn``/``ts_posteq``/``my``/ @@ -1001,7 +1001,7 @@ def simulate_condition( :return: output according to `ret` and general results/statistics """ - return self.simulate_condition_unjitted( + return self.simulate_experiment_unjitted( p, ts_dyn, ts_posteq, diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 7fd3570acb..af35cf6f0d 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -53,7 +53,7 @@ class SteadyStateEvent(eqx.Module): :func:`diffrax.steady_state_event` returns a fresh closure on every call, which Python compares by identity. Since ``steady_state_event`` is passed as a static argument into :func:`equinox.filter_jit`-compiled - functions (:meth:`JAXModel.simulate_condition`, + functions (:meth:`JAXModel.simulate_experiment`, :meth:`JAXModel.preequilibrate_condition`), constructing a new closure with identical settings on every call (e.g. once per iteration of an optimization loop) silently defeats the JIT cache and forces a full @@ -1799,7 +1799,7 @@ def _prepare_experiments( For the main simulation (``is_preeq=False``), all returned arrays (except ``h_mask``) gain a period axis right after the experiment axis, of static size :attr:`_max_periods`, so that - :meth:`JAXModel.simulate_condition` can chain one ODE integration + :meth:`JAXModel.simulate_experiment` can chain one ODE integration per period. Periods beyond an experiment's own period count are padding periods: they are never reinitialised and never contribute to the log-likelihood (see :meth:`_get_measurements`), but must @@ -2081,7 +2081,7 @@ def run_simulation( :param h_preeq: Pre-equilibration event mask. Can be empty if no pre-equilibration is available :param ts_mask: - padding mask, see :meth:`JAXModel.simulate_condition` for details. + padding mask, see :meth:`JAXModel.simulate_experiment` for details. :param t_zeros: simulation start time for the current experiment. :param ret: @@ -2092,7 +2092,7 @@ def run_simulation( model = eqx.tree_at( lambda m: m._array_input_index, self.model, experiment_index ) - return model.simulate_condition( + return model.simulate_experiment( p=p, ts_dyn=jax.lax.stop_gradient(jnp.array(ts_dyn)), ts_posteq=jax.lax.stop_gradient(jnp.array(ts_posteq)), diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 47920df435..75c2f07c18 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -194,7 +194,7 @@ def check_fields_jax( } p = jnp.array([par_dict[par_id] for par_id in jax_model.parameter_ids]) - # `simulate_condition[_unjitted]` chains one ODE integration per + # `simulate_experiment[_unjitted]` chains one ODE integration per # experiment period; add a leading period axis of size 1 for this # single, non-chained simulation. `p` itself is kept 1-D here (and # the period axis added inside `fun` below) so that `jax.grad`/ @@ -218,7 +218,7 @@ def check_fields_jax( } # Use beartype-wrapped unjitted version for type checking # (beartype cannot introspect jitted functions, so we wrap the unjitted version) - fun_periodic = beartype(jax_model.simulate_condition_unjitted) + fun_periodic = beartype(jax_model.simulate_experiment_unjitted) def fun(p, **kw): return fun_periodic(p[None, :], **kw) @@ -570,11 +570,11 @@ def dispatch(self, *args, **kwargs): iy_trafos = jnp.zeros_like(ts, dtype=int) simulate_traces = patch_trace_counter( - JAXModel, "simulate_condition_unjitted" + JAXModel, "simulate_experiment_unjitted" ) for k_val in conditions: kwargs = fresh_solver_kwargs() - model.simulate_condition( + model.simulate_experiment( jnp.array([[k_val]]), ts[None, :], jnp.zeros((1, 0)), @@ -592,7 +592,7 @@ def dispatch(self, *args, **kwargs): ) assert eqx.debug.get_num_traces(simulate_traces) == 1, ( - "simulate_condition was retraced across conditions with only " + "simulate_experiment was retraced across conditions with only " "numeric differences" ) diff --git a/tests/performance/test_jax_regression.py b/tests/performance/test_jax_regression.py index 8b5c7aa874..13dce0b79b 100644 --- a/tests/performance/test_jax_regression.py +++ b/tests/performance/test_jax_regression.py @@ -30,11 +30,11 @@ "MultiEvent", ] -# ── Helper: build simulate_condition_unjitted kwargs for each model ───────── +# ── Helper: build simulate_experiment_unjitted kwargs for each model ───────── def _sim_kwargs(model, solver_kwargs) -> dict: - """Return keyword arguments for simulate_condition_unjitted.""" + """Return keyword arguments for simulate_experiment_unjitted.""" import tests.performance.synthetic_models.conservation_law as cl import tests.performance.synthetic_models.linear_decay as ld import tests.performance.synthetic_models.lotka_volterra as lv @@ -89,7 +89,7 @@ def _sim_kwargs(model, solver_kwargs) -> dict: model_name = type(model).__name__ ts_dyn, my, iys, iy_trafos, ops, nps = _MAP[model_name] - # `simulate_condition`/`simulate_condition_unjitted` chain one ODE + # `simulate_experiment`/`simulate_experiment_unjitted` chain one ODE # integration per experiment period; add a leading period axis of # size 1 since these synthetic models are all single-period. return dict( @@ -105,7 +105,7 @@ def _sim_kwargs(model, solver_kwargs) -> dict: def _extract_stats(stats: dict) -> dict: - """Pull step counts out of the stats dict returned by simulate_condition.""" + """Pull step counts out of the stats dict returned by simulate_experiment.""" out = {} for key in ("stats_dyn", "stats_posteq"): s = stats.get(key) @@ -139,10 +139,10 @@ def test_tier1_fwd_sim( kwargs = _sim_kwargs(model, solver_kwargs) # Deterministic run (unjitted, for exact step counts) - llh, stats = model.simulate_condition_unjitted(p, **kwargs) + llh, stats = model.simulate_experiment_unjitted(p, **kwargs) # Timing (JIT-compiled path) - sim_fn = model.simulate_condition + sim_fn = model.simulate_experiment t_first, t_exec = measure_exec_time(sim_fn, p, **kwargs) results_collector.add( @@ -168,7 +168,7 @@ def test_tier1_adj(model_id, tier1_models, solver_kwargs, results_collector): kwargs = _sim_kwargs(model, solver_kwargs) def _fn(p): - return model.simulate_condition(p, **kwargs) + return model.simulate_experiment(p, **kwargs) t_first, t_exec = measure_exec_time( eqx.filter_value_and_grad(_fn, has_aux=True), p @@ -200,7 +200,7 @@ def test_tier1_fwd_sens( } def _fn(p): - return model.simulate_condition(p, **kwargs) + return model.simulate_experiment(p, **kwargs) t_first, t_exec = measure_exec_time(jax.jacfwd(_fn, has_aux=True), p) diff --git a/tests/sbml/testSBMLSuite.py b/tests/sbml/testSBMLSuite.py index ca74e1dfe4..16a555c0c0 100755 --- a/tests/sbml/testSBMLSuite.py +++ b/tests/sbml/testSBMLSuite.py @@ -220,7 +220,7 @@ def jax_sensitivity_check( root_finder = optimistix.Newton(**DEFAULT_ROOT_FINDER_SETTINGS) def simulate(pars): - x, _ = jax_model.simulate_condition( + x, _ = jax_model.simulate_experiment( pars, ts_jnp, jnp.array([]), diff --git a/tests/sbml/testSBMLSuiteJax.py b/tests/sbml/testSBMLSuiteJax.py index 295b3d97a5..15b00ecce3 100644 --- a/tests/sbml/testSBMLSuiteJax.py +++ b/tests/sbml/testSBMLSuiteJax.py @@ -72,10 +72,10 @@ def run_jax_simulation(model, importer, ts, atol, rtol, tol_factor=1e2): dcoeff=DEFAULT_CONTROLLER_SETTINGS["dcoeff"], ) root_finder = optimistix.Newton(atol=atol, rtol=rtol) - # `simulate_condition` chains one ODE integration per experiment + # `simulate_experiment` chains one ODE integration per experiment # period; add a leading period axis of size 1 for this single, # non-chained simulation. - x, stats = model.simulate_condition( + x, stats = model.simulate_experiment( p[None, :], ts_jnp[None, :], jnp.zeros((1, 0)), From 6834d3ec2bf72a0f7c91fef4fbc65c7e48e54735 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 18:41:51 +0000 Subject: [PATCH 17/22] Fix stale docstring wording after simulate_experiment rename --- python/sdist/amici/sim/jax/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sdist/amici/sim/jax/model.py b/python/sdist/amici/sim/jax/model.py index c64a015f85..6cf2954c6a 100644 --- a/python/sdist/amici/sim/jax/model.py +++ b/python/sdist/amici/sim/jax/model.py @@ -938,7 +938,7 @@ def simulate_experiment( ret: ReturnValue = ReturnValue.llh, ) -> tuple[jt.Float[jt.Array, "*nt"], dict]: r""" - Simulate a condition (JIT-compiled version). + Simulate an experiment (JIT-compiled version). This is the JIT-compiled version for optimal performance. For runtime type checking with beartype, use :meth:`simulate_experiment_unjitted` instead. From c47d15d328d6176c878c633b6f1d85cad7186082 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 19:08:19 +0000 Subject: [PATCH 18/22] Fix missing period axis in testSBMLSuite.py's jax_sensitivity_check 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. --- tests/sbml/testSBMLSuite.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/sbml/testSBMLSuite.py b/tests/sbml/testSBMLSuite.py index 16a555c0c0..3b4cf4a170 100755 --- a/tests/sbml/testSBMLSuite.py +++ b/tests/sbml/testSBMLSuite.py @@ -221,14 +221,14 @@ def jax_sensitivity_check( def simulate(pars): x, _ = jax_model.simulate_experiment( - pars, - ts_jnp, - jnp.array([]), - zeros, - jnp.zeros_like(ts_jnp, dtype=int), - jnp.zeros_like(ts_jnp, dtype=int), - jnp.zeros((ts_jnp.shape[0], 0)), - jnp.zeros((ts_jnp.shape[0], 0)), + pars[None, :], + ts_jnp[None, :], + jnp.zeros((1, 0)), + zeros[None, :], + jnp.zeros_like(ts_jnp, dtype=int)[None, :], + jnp.zeros_like(ts_jnp, dtype=int)[None, :], + jnp.zeros((1, ts_jnp.shape[0], 0)), + jnp.zeros((1, ts_jnp.shape[0], 0)), solver, controller, root_finder, From 0a7a919cf3a37f1ce2a7cb5a015bb1d354ad436c Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 21:30:06 +0000 Subject: [PATCH 19/22] Simplify per-period measurement bucketing and compile condition-table 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. --- python/sdist/amici/sim/jax/petab.py | 357 +++++++++++++++++----------- python/tests/test_jax.py | 83 +++++++ 2 files changed, 302 insertions(+), 138 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index af35cf6f0d..62d9628e47 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 from typing import NamedTuple @@ -19,11 +18,13 @@ import pandas as pd import petab.v1 as petabv1 import petab.v2 as petabv2 +import sympy as sp from optimistix import AbstractRootFinder from amici import _module_from_path +from amici.exporters.jax.jaxcodeprinter import AmiciJaxCodePrinter from amici.logging import get_logger -from amici.sim.jax.model import JAXModel, ReturnValue +from amici.sim.jax.model import JAXModel, ReturnValue, safe_div, safe_log DEFAULT_CONTROLLER_SETTINGS = { "atol": 1e-8, @@ -292,16 +293,32 @@ def _get_iy_trafos( class _PeriodMeasurements(NamedTuple): """One experiment period's bucketed measurement data, as built by :meth:`JAXProblem._get_measurements` and consumed by its - padding/stacking step.""" + padding/stacking step. + + Dynamic-phase and post-equilibrium-phase data are tracked as separate + fields throughout (mirroring ``ts_dyn``/``ts_posteq``) rather than + concatenated into one array and later re-split by ``len(ts_dyn)``: once + merged, that split point is only recoverable by convention, and a + single missed field silently corrupts that field's post-equilibrium + entries while leaving them marked valid. + """ ts_dyn: np.ndarray ts_posteq: np.ndarray - my: np.ndarray - iys: np.ndarray - iy_trafos: np.ndarray - op_overrides: OverrideColumn - noise_overrides: OverrideColumn - valid: np.ndarray + my_dyn: np.ndarray + my_posteq: np.ndarray + iys_dyn: np.ndarray + iys_posteq: np.ndarray + iy_trafos_dyn: np.ndarray + iy_trafos_posteq: np.ndarray + op_overrides_dyn: OverrideColumn + op_overrides_posteq: OverrideColumn + noise_overrides_dyn: OverrideColumn + noise_overrides_posteq: OverrideColumn + valid_dyn: np.ndarray + valid_posteq: np.ndarray + index_dyn: tuple[int, ...] + index_posteq: tuple[int, ...] def _masked_placeholder_period( @@ -365,17 +382,19 @@ def _pad_measurement( def _pad_and_stack( measurements: dict[tuple[str, int], _PeriodMeasurements], - extractor: Callable[[_PeriodMeasurements], np.ndarray], + extractor_dyn: Callable[[_PeriodMeasurements], np.ndarray], + extractor_posteq: Callable[[_PeriodMeasurements], np.ndarray], n_ts_dyn: int, n_ts_posteq: int, ) -> np.ndarray: - """Apply ``extractor`` to every bucketed period, split each result at - its own dynamic/post-equilibrium boundary, pad, and stack.""" + """Apply ``extractor_dyn``/``extractor_posteq`` to every bucketed + period's independently-tracked dynamic/post-equilibrium portions, pad + each to ``n_ts_dyn``/``n_ts_posteq``, and stack across periods.""" return np.stack( [ _pad_measurement( - extractor(mv)[: len(mv.ts_dyn)], - extractor(mv)[len(mv.ts_dyn) :], + extractor_dyn(mv), + extractor_posteq(mv), n_ts_dyn, n_ts_posteq, ) @@ -561,7 +580,6 @@ def _get_measurements( - parameter indices (problem parameters) for noise parameter overrides """ measurements: dict[tuple[str, int], _PeriodMeasurements] = {} - petab_indices = dict() # Nominal (linear) values of fixed (non-estimated) parameters, used to # resolve observable/noise parameter overrides that reference them. @@ -631,13 +649,25 @@ def _get_measurements( if i_period >= len(dyn_periods): # Padding slot: this experiment has no period here at # all. A single masked, zero-duration integration step - # that leaves the carried-over state unchanged. - ts_dyn, dyn_valid, my, iys, iy_trafos, overrides = ( + # that leaves the carried-over state unchanged. No + # post-equilibrium portion applies to a padding slot. + ts_dyn, dyn_valid, my_dyn, iys_dyn, iy_trafos_dyn, overrides_dyn = ( _masked_placeholder_period(last_period_time, n_pars) ) + index_dyn = (-1,) ts_posteq = np.array([]) - posteq_valid = np.array([]) - index = [-1] + my_posteq = np.array([]) + iys_posteq = np.array([], dtype=int) + iy_trafos_posteq = np.array([], dtype=int) + overrides_posteq = { + col: OverrideColumn.placeholder(0, n_pars[col]) + for col in ( + petabv2.C.OBSERVABLE_PARAMETERS, + petabv2.C.NOISE_PARAMETERS, + ) + } + posteq_valid = np.array([], dtype=bool) + index_posteq = () else: is_own_last = i_period == len(dyn_periods) - 1 t_lo = dyn_periods[i_period].time @@ -670,26 +700,19 @@ def _get_measurements( overrides_real = _get_overrides( m, n_pars, fixed_parameter_values, self.parameter_ids ) - index_dyn = list(m.index) + index_dyn_real = list(m.index) if is_own_last: + # Post-equilibrium measurements (e.g. steady-state + # comparisons) have their own observable/override + # data, distinct from the dynamic-phase measurements + # above, and are tracked as a separate portion of + # this period throughout (see `_PeriodMeasurements`). posteq_mask = np.isfinite(m_full_times) == False # noqa: E712 m_posteq = m_full[posteq_mask] ts_posteq = m_posteq[petabv2.C.TIME].values - index_posteq = list(m_posteq.index) + index_posteq = tuple(m_posteq.index) - # Post-equilibrium measurements (e.g. steady-state - # comparisons) have their own observable/override - # data, distinct from the dynamic-phase measurements - # above; this must be concatenated onto the dyn - # portion so that `mv.my`/`iys`/`iy_trafos`/overrides - # cover the full `len(ts_dyn) + len(ts_posteq)` range - # that `_pad_and_stack` later splits at - # `len(mv.ts_dyn)` -- otherwise post-equilibrium rows - # silently fall back to zero-filled placeholders - # (wrong measurement value, wrong observable index, - # wrong noise override) while still being marked - # valid. iys_posteq = np.array( [ self.model.observable_ids.index(oid) @@ -724,7 +747,7 @@ def _get_measurements( iy_trafos_real, ) overrides_dyn = overrides_real - index_dyn_full = index_dyn + index_dyn = tuple(index_dyn_real) else: # No real dyn measurements in this period (e.g. # a pure post-equilibration period): still need @@ -739,22 +762,17 @@ def _get_measurements( iy_trafos_dyn, overrides_dyn, ) = _masked_placeholder_period(t_lo, n_pars) - index_dyn_full = [-1] - - my = np.concatenate([my_dyn, my_posteq]) - iys = np.concatenate([iys_dyn, iys_posteq]) - iy_trafos = np.concatenate( - [iy_trafos_dyn, iy_trafos_posteq] - ) - overrides = { - col: OverrideColumn.concatenate( - overrides_dyn[col], overrides_posteq[col] - ) - for col in overrides_dyn - } + index_dyn = (-1,) else: ts_posteq = np.array([]) - index_posteq = [] + my_posteq = np.array([]) + iys_posteq = np.array([], dtype=int) + iy_trafos_posteq = np.array([], dtype=int) + overrides_posteq = { + col: OverrideColumn.placeholder(0, n_pars[col]) + for col in overrides_real + } + index_posteq = () # Append a synthetic, masked boundary time point at # the start of the next period so that the ODE @@ -765,36 +783,46 @@ def _get_measurements( dyn_valid = np.append( np.ones(len(ts_dyn_real), dtype=bool), False ) - my = np.append(my_real, 0.0) - iys = np.append(iys_real, 0) - iy_trafos = np.append(iy_trafos_real, 0) - overrides = { + my_dyn = np.append(my_real, 0.0) + iys_dyn = np.append(iys_real, 0) + iy_trafos_dyn = np.append(iy_trafos_real, 0) + overrides_dyn = { col: OverrideColumn.concatenate( overrides_real[col], OverrideColumn.placeholder(1, n_pars[col]), ) for col in overrides_real } - index_dyn_full = [*index_dyn, -1] + index_dyn = (*index_dyn_real, -1) - index = [*index_dyn_full, *index_posteq] posteq_valid = np.ones(len(ts_posteq), dtype=bool) - valid = np.concatenate([dyn_valid, posteq_valid]).astype( - bool - ) - measurements[(exp.id, i_period)] = _PeriodMeasurements( ts_dyn=ts_dyn, ts_posteq=ts_posteq, - my=my, - iys=iys, - iy_trafos=iy_trafos, - op_overrides=overrides[petabv2.C.OBSERVABLE_PARAMETERS], - noise_overrides=overrides[petabv2.C.NOISE_PARAMETERS], - valid=valid, + my_dyn=my_dyn, + my_posteq=my_posteq, + iys_dyn=iys_dyn, + iys_posteq=iys_posteq, + iy_trafos_dyn=iy_trafos_dyn, + iy_trafos_posteq=iy_trafos_posteq, + op_overrides_dyn=overrides_dyn[ + petabv2.C.OBSERVABLE_PARAMETERS + ], + op_overrides_posteq=overrides_posteq[ + petabv2.C.OBSERVABLE_PARAMETERS + ], + noise_overrides_dyn=overrides_dyn[ + petabv2.C.NOISE_PARAMETERS + ], + noise_overrides_posteq=overrides_posteq[ + petabv2.C.NOISE_PARAMETERS + ], + valid_dyn=dyn_valid, + valid_posteq=posteq_valid, + index_dyn=index_dyn, + index_posteq=index_posteq, ) - petab_indices[(exp.id, i_period)] = tuple(index) # compute maximum lengths n_ts_dyn = max(len(mv.ts_dyn) for mv in measurements.values()) @@ -823,47 +851,65 @@ def _get_measurements( ) my = _pad_and_stack( - measurements, lambda mv: mv.my, n_ts_dyn, n_ts_posteq + measurements, + lambda mv: mv.my_dyn, + lambda mv: mv.my_posteq, + n_ts_dyn, + n_ts_posteq, ) iys = _pad_and_stack( - measurements, lambda mv: mv.iys, n_ts_dyn, n_ts_posteq + measurements, + lambda mv: mv.iys_dyn, + lambda mv: mv.iys_posteq, + n_ts_dyn, + n_ts_posteq, ) iy_trafos = _pad_and_stack( - measurements, lambda mv: mv.iy_trafos, n_ts_dyn, n_ts_posteq + measurements, + lambda mv: mv.iy_trafos_dyn, + lambda mv: mv.iy_trafos_posteq, + n_ts_dyn, + n_ts_posteq, ) op_numeric = _pad_and_stack( measurements, - lambda mv: mv.op_overrides.numeric, + lambda mv: mv.op_overrides_dyn.numeric, + lambda mv: mv.op_overrides_posteq.numeric, n_ts_dyn, n_ts_posteq, ) op_mask = _pad_and_stack( measurements, - lambda mv: mv.op_overrides.mask, + lambda mv: mv.op_overrides_dyn.mask, + lambda mv: mv.op_overrides_posteq.mask, n_ts_dyn, n_ts_posteq, ) op_indices = _pad_and_stack( measurements, - lambda mv: mv.op_overrides.index, + lambda mv: mv.op_overrides_dyn.index, + lambda mv: mv.op_overrides_posteq.index, n_ts_dyn, n_ts_posteq, ) np_numeric = _pad_and_stack( measurements, - lambda mv: mv.noise_overrides.numeric, + lambda mv: mv.noise_overrides_dyn.numeric, + lambda mv: mv.noise_overrides_posteq.numeric, n_ts_dyn, n_ts_posteq, ) np_mask = _pad_and_stack( measurements, - lambda mv: mv.noise_overrides.mask, + lambda mv: mv.noise_overrides_dyn.mask, + lambda mv: mv.noise_overrides_posteq.mask, n_ts_dyn, n_ts_posteq, ) np_indices = _pad_and_stack( measurements, - lambda mv: mv.noise_overrides.index, + lambda mv: mv.noise_overrides_dyn.index, + lambda mv: mv.noise_overrides_posteq.index, n_ts_dyn, n_ts_posteq, ) @@ -875,11 +921,11 @@ def _get_measurements( np.concatenate( ( np.pad( - mv.valid[: len(mv.ts_dyn)], + mv.valid_dyn, (0, n_ts_dyn - len(mv.ts_dyn)), ), np.pad( - mv.valid[len(mv.ts_dyn) :], + mv.valid_posteq, (0, n_ts_posteq - len(mv.ts_posteq)), ), ) @@ -890,14 +936,12 @@ def _get_measurements( petab_indices = np.stack( [ _pad_measurement( - np.array(idx[: len(mv.ts_dyn)]), - np.array(idx[len(mv.ts_dyn) :]), + np.array(mv.index_dyn), + np.array(mv.index_posteq), n_ts_dyn, n_ts_posteq, ) - for mv, idx in zip( - measurements.values(), petab_indices.values() - ) + for mv in measurements.values() ] ) @@ -926,16 +970,16 @@ def _get_measurements( ) def _get_parameter_mappings(self) -> dict[str, ...]: - # `targets_map` intentionally stores each value only lightly parsed - # (a numeric literal, or a parameter id `str`, via - # `_resolve_petab_change_value`) rather than eagerly resolved - # against `self.parameters`: this dict is cached once at - # construction time (`self._parameter_mappings`, see `__init__`), - # and `update_parameters` (`eqx.tree_at(lambda p: p.parameters, - # self, p)`) does not recompute it, so an eagerly-resolved value - # would go stale after a parameter update. - # `_resolve_parameter_reference` re-resolves the actual value live - # against `self.parameters` at use time instead. + # `targets_map` intentionally stores each value only as a compiled + # `_CompiledConditionExpr` (see `_resolve_petab_change_value`) + # rather than eagerly resolved against `self.parameters`: this + # dict is cached once at construction time + # (`self._parameter_mappings`, see `__init__`), and + # `update_parameters` (`eqx.tree_at(lambda p: p.parameters, self, + # p)`) does not recompute it, so an eagerly-resolved value would + # go stale after a parameter update. `_resolve_parameter_reference` + # re-resolves the actual value live against `self.parameters` at + # use time instead. targets_map = { c.id: { ch.target_id: _resolve_petab_change_value(ch.target_value) @@ -958,30 +1002,36 @@ def _get_parameter_mappings(self) -> dict[str, ...]: return {"targets_map": targets_map, "hybrid_map": hybrid_map} def _resolve_parameter_reference( - self, value: float | str + self, value: "_CompiledConditionExpr" ) -> jt.Float[jt.Scalar, ""]: # noqa: F722 """ Resolve a value from ``targets_map`` (see :meth:`_get_parameter_mappings`) - to a JAX scalar. Numeric values pass through; a parameter id - resolves to that parameter's current (estimated) or nominal - (fixed) value. Symbolic references to estimated parameters keep - their dependence on :attr:`parameters` so gradients flow through + to a JAX scalar by evaluating its compiled expression with each + free symbol resolved to that parameter's current (estimated) or + nominal (fixed) value. A numeric literal or a single parameter + reference is just the zero/one-free-symbol case of the same + mechanism. Symbolic references to estimated parameters keep their + dependence on :attr:`parameters` so gradients flow through correctly. :param value: - A numeric literal or PEtab parameter id, as returned by + A compiled expression, as returned by :func:`_resolve_petab_change_value`. """ - if isinstance(value, str): - if value in self.parameter_ids: - return self.parameters[self.parameter_ids.index(value)] + + def resolve_symbol(name: str) -> jt.Array: + if name in self.parameter_ids: + return self.parameters[self.parameter_ids.index(name)] return jnp.asarray( self._petab_problem.parameter_df.loc[ - value, petabv2.C.NOMINAL_VALUE + name, petabv2.C.NOMINAL_VALUE ], dtype=self.model.parameters.dtype, ) - return jnp.asarray(value, dtype=self.model.parameters.dtype) + + return jnp.asarray( + value(resolve_symbol), dtype=self.model.parameters.dtype + ) def get_all_simulation_conditions(self) -> tuple[tuple[str, ...], ...]: simulation_conditions = get_simulation_conditions_v2( @@ -1620,10 +1670,9 @@ def _first_condition_value( ``condition_ids`` (applied simultaneously, e.g. for one experiment period) that actually sets it. - Numeric values are returned as plain Python ``float``s; a - reference to a single other parameter is returned as that - parameter's id (``str``). Compound symbolic expressions are not - supported. + The value is returned as a :class:`_CompiledConditionExpr` (a + numeric literal or a reference to a single other parameter is just + its zero/one-free-symbol special case). :param condition_ids: Condition IDs to check, in priority order. @@ -1698,23 +1747,28 @@ def _state_reinitialisation_value( ) 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. PEtab v2 - # has no parameterScale column -- all parameters are linear. - return jax_unscale( - self.get_petab_parameter_by_id(xval), petabv2.C.LIN - ) - # 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 - ] + + def resolve_symbol(name: str): + if name in self.model.parameter_ids: + # model parameter, return value + return p[self.model.parameter_ids.index(name)] + if name in self.parameter_ids: + # estimated PEtab parameter, return unscaled value. PEtab + # v2 has no parameterScale column -- all parameters are + # linear. + return jax_unscale( + self.get_petab_parameter_by_id(name), petabv2.C.LIN + ) + # only remaining option is nominal value for PEtab parameter + # that is not estimated, return nominal value + return self._petab_problem.parameter_df.loc[ + name, petabv2.C.NOMINAL_VALUE + ] + + # a numeric literal or a single parameter reference is just the + # zero/one-free-symbol case of the same compiled-expression + # mechanism + return xval(resolve_symbol) def load_reinitialisation( self, @@ -2782,25 +2836,52 @@ def _try_float(value): raise -def _resolve_petab_change_value(target_value) -> float | str: +class _CompiledConditionExpr(NamedTuple): + """A condition table change's compound symbolic ``target_value`` + expression (e.g. ``k1 + 2 * k2``), compiled to a JAX-callable function + of its free parameter symbols (sorted by name), using the same + sympy-to-JAX code printer (:class:`AmiciJaxCodePrinter`) that the + model's own equations are generated with -- rather than being + restricted to the numeric-literal/single-parameter-reference special + cases that :func:`_resolve_petab_change_value` can resolve without + compilation. + + Each caller resolves the free symbols to JAX values according to its + own rules (e.g. a state reinitialisation value may reference the + model's own parameters directly, while a parameter override may only + reference other PEtab parameters) and supplies them via + :meth:`__call__`. """ - Resolve a :class:`petabv2.Change.target_value` (a sympy expression, or - already a plain number) to either a numeric literal or a single - parameter id. - Only numeric literals and references to a single other parameter are - supported; compound symbolic expressions (e.g. ``"k1 + k2"``) are not. + symbol_names: tuple[str, ...] + fn: Callable[..., jt.Array] + + @classmethod + def compile(cls, expr: sp.Expr) -> "_CompiledConditionExpr": + free_syms = sorted(expr.free_symbols, key=str) + symbol_names = tuple(s.name for s in free_syms) + body = AmiciJaxCodePrinter().doprint(expr) + namespace = {"jnp": jnp, "safe_log": safe_log, "safe_div": safe_div} + src = f"def _expr({', '.join(symbol_names)}):\n return {body}\n" + exec(src, namespace) # noqa: S102 -- code-generated from the PEtab problem's own sympy expressions, not user input + return cls(symbol_names, namespace["_expr"]) + + def __call__(self, resolve_symbol: Callable[[str], jt.Array]) -> jt.Array: + return self.fn(*(resolve_symbol(name) for name in self.symbol_names)) + + +def _resolve_petab_change_value(target_value) -> _CompiledConditionExpr: + """ + Compile a :class:`petabv2.Change.target_value` (a sympy expression, a + plain number, or a single parameter reference) to a + :class:`_CompiledConditionExpr`. A numeric literal or a single + parameter reference is just the zero/one-free-symbol case of the same + general compiled-expression mechanism, so no special-casing is needed + here. :param target_value: Value to resolve. :return: - A ``float``, or a ``str`` naming a PEtab/model parameter id. + A :class:`_CompiledConditionExpr`. """ - if getattr(target_value, "is_number", True): - return float(target_value) - if getattr(target_value, "is_Symbol", False): - return str(target_value) - raise NotImplementedError( - "Condition table changes with compound symbolic expressions are " - f"not supported, got {target_value!r}." - ) + return _CompiledConditionExpr.compile(sp.sympify(target_value)) diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 75c2f07c18..bbf8b6291d 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -449,6 +449,89 @@ def llh(p): assert_allclose(float(grad.parameters[ik]), fd, rtol=1e-4, atol=1e-4) +@skip_on_valgrind +def test_condition_table_compound_expression_is_differentiable(tmp_path): + """A condition table change with a compound symbolic ``target_value`` + (e.g. ``a0 + b0``, referencing two estimated parameters) must be + compiled and evaluated correctly, with gradients flowing through every + free symbol. + + PEtab v1's condition table format only allows numeric literals or a + single parameter reference as a condition value; compound expressions + are a PEtab v2-only feature, so this uses the native v2 API rather + than the v1-upgrade path used by the sibling tests above. + + Regression test for ``_resolve_petab_change_value``, which used to + raise ``NotImplementedError`` for anything beyond a numeric literal or + a single parameter reference; compound expressions are now compiled to + JAX via the same sympy-to-JAX code printer used to generate the + model's own equations (see ``_CompiledConditionExpr``). + """ + import equinox as eqx + import petab.v2 as petabv2 + from amici.importers.petab._petab_importer import PetabImporter + from petab.v2.core import ProblemConfig + from petab.v2.models.sbml_model import SbmlModel + + problem = petabv2.Problem() + problem.config = ProblemConfig() + 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" + ) + problem.add_parameter( + "a0", nominal_value=2.0, estimate=True, lb=0.1, ub=10 + ) + problem.add_parameter( + "b0", nominal_value=1.0, estimate=True, lb=0.1, ub=10 + ) + problem.add_observable("obs_a", "A", noise_formula="0.5") + # compound expression: A's initial value is the *sum* of two estimated + # parameters, not a numeric literal or a single parameter reference + problem.add_condition("c0", A="a0 + b0") + problem.add_experiment("exp0", 0.0, "c0") + for t in [0.0, 10.0]: + problem.add_measurement( + "obs_a", experiment_id="exp0", time=t, measurement=0.5 + ) + + jax_problem = PetabImporter( + problem, + jax=True, + module_name="test_condition_table_compound_expression_jax", + verbose=False, + output_dir=str(tmp_path), + ).create_simulator(force_import=True) + + ia = jax_problem.parameter_ids.index("a0") + ib = jax_problem.parameter_ids.index("b0") + + def llh(p): + return run_simulations(jax_problem.update_parameters(p))[0] + + p0 = jax_problem.parameters + # `a0 + b0` enters the likelihood only through A(0); updating either + # must change llh + assert abs(float(llh(p0.at[ia].add(1.0))) - float(llh(p0))) > 1e-6, ( + "compound-expression condition value is frozen w.r.t. " + "update_parameters" + ) + + eps = 1e-6 + grad = eqx.filter_grad(lambda m: run_simulations(m)[0])( + jax_problem.update_parameters(p0) + ) + for i in (ia, ib): + fd = ( + float(llh(p0.at[i].add(eps))) - float(llh(p0.at[i].add(-eps))) + ) / (2 * eps) + assert_allclose(float(grad.parameters[i]), 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 From 82bece715a412657af324f84f97fe7d770c4c351 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 21:47:36 +0000 Subject: [PATCH 20/22] Reject state-referencing condition expressions explicitly instead of 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. --- python/sdist/amici/sim/jax/petab.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 62d9628e47..f102d0837b 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -1022,6 +1022,17 @@ def _resolve_parameter_reference( def resolve_symbol(name: str) -> jt.Array: if name in self.parameter_ids: return self.parameters[self.parameter_ids.index(name)] + if name in self.model.state_ids: + # a parameter override referencing a state's value would + # require the actual simulated trajectory, which does not + # exist yet at this (static, pre-simulation) precompute + # step -- see the matching check in + # `_state_reinitialisation_value`. + raise NotImplementedError( + "Condition table changes referencing a state's value " + f"(got {name!r}) are not supported for parameter " + "overrides." + ) return jnp.asarray( self._petab_problem.parameter_df.loc[ name, petabv2.C.NOMINAL_VALUE @@ -1759,6 +1770,19 @@ def resolve_symbol(name: str): return jax_unscale( self.get_petab_parameter_by_id(name), petabv2.C.LIN ) + if name in self.model.state_ids: + # a reinitialisation value referencing another state's + # value (e.g. `A = A + 5.0`) would need that state's + # actual simulated value at this period boundary, which + # `load_reinitialisation` cannot provide: `x_reinit` is + # precomputed once per experiment (see + # `_prepare_experiments`), before any period is actually + # integrated, so no simulated state exists yet to + # reference. + raise NotImplementedError( + "Condition table changes referencing another state's " + f"value (got {name!r}) are not supported." + ) # only remaining option is nominal value for PEtab parameter # that is not estimated, return nominal value return self._petab_problem.parameter_df.loc[ From 908d2db600b3baa1e255e967f3c9d23c1f8e442a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 22:13:15 +0000 Subject: [PATCH 21/22] Inline _resolve_override_symbol 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. --- python/sdist/amici/sim/jax/petab.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index f102d0837b..cd340a9481 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -151,14 +151,6 @@ def _get_fixed_parameter_values( } -def _resolve_override_symbol(value, fixed_parameter_values: dict[str, float]): - """Replace a reference to a non-estimated PEtab parameter by its - nominal value; pass through anything else (numeric literals, estimated - parameter ids, model entity ids, array-valued fixed parameters) - unchanged.""" - return fixed_parameter_values.get(value, value) - - def _split_override_column( col_values: pd.Series, n_pars: int, @@ -186,10 +178,11 @@ def resolve_row(entry) -> list: if isinstance(entry, str) else [entry] ) - return [ - _resolve_override_symbol(v, fixed_parameter_values) - for v in values - ] + # replace a reference to a non-estimated PEtab parameter by its + # nominal value; pass through anything else (numeric literals, + # estimated parameter ids, model entity ids, array-valued fixed + # parameters) unchanged + return [fixed_parameter_values.get(v, v) for v in values] rows = col_values.apply(resolve_row) padded = rows.apply( From f7d9aa4e3387de80172e3901b48a805adc4e036b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 23:45:55 +0000 Subject: [PATCH 22/22] Add deprecated simulate_condition alias; trim multi-period tests redundant 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. --- python/sdist/amici/sim/jax/model.py | 22 +++++++ .../tests/petab_/test_petab_v2_multiperiod.py | 48 -------------- python/tests/test_jax.py | 62 +++++++++++++++++++ 3 files changed, 84 insertions(+), 48 deletions(-) diff --git a/python/sdist/amici/sim/jax/model.py b/python/sdist/amici/sim/jax/model.py index 6cf2954c6a..8c4817d968 100644 --- a/python/sdist/amici/sim/jax/model.py +++ b/python/sdist/amici/sim/jax/model.py @@ -4,6 +4,7 @@ import enum import os +import warnings from abc import abstractmethod from collections.abc import Callable from dataclasses import field @@ -907,6 +908,17 @@ def simulate_experiment_unjitted( return output, stats + def simulate_condition_unjitted(self, *args, **kwargs): + """Deprecated alias for :meth:`simulate_experiment_unjitted`.""" + warnings.warn( + "`simulate_condition_unjitted` has been renamed to " + "`simulate_experiment_unjitted` and will be removed in a " + "future release.", + DeprecationWarning, + stacklevel=2, + ) + return self.simulate_experiment_unjitted(*args, **kwargs) + @eqx.filter_jit def simulate_experiment( self, @@ -1028,6 +1040,16 @@ def simulate_experiment( ret, ) + def simulate_condition(self, *args, **kwargs): + """Deprecated alias for :meth:`simulate_experiment`.""" + warnings.warn( + "`simulate_condition` has been renamed to `simulate_experiment` " + "and will be removed in a future release.", + DeprecationWarning, + stacklevel=2, + ) + return self.simulate_experiment(*args, **kwargs) + @eqx.filter_jit def preequilibrate_condition( self, diff --git a/python/tests/petab_/test_petab_v2_multiperiod.py b/python/tests/petab_/test_petab_v2_multiperiod.py index a946095adf..47a2a7ae2f 100644 --- a/python/tests/petab_/test_petab_v2_multiperiod.py +++ b/python/tests/petab_/test_petab_v2_multiperiod.py @@ -92,54 +92,6 @@ def analytical(t): assert np.isfinite(llh) -def test_two_period_preequilibration_matches_analytical_solution(tmp_path): - """Regression check for the common (already-supported) pre-equilibration - + one main period case, now driven through the same native-chaining - code path as N>2-period experiments (P=2 special case).""" - problem = _linear_decay_problem() - problem.add_condition("cond_preeq", kk=0.5) - problem.add_condition("cond_main", kk=0.7, xx=2.0) - problem.add_experiment( - "exp1", C.TIME_PREEQUILIBRATION, "cond_preeq", 0.0, "cond_main" - ) - ts = (0.0, 1.0, 2.0, 3.0) - for t in ts: - problem.add_measurement( - "obs1", time=t, measurement=0.0, experiment_id="exp1" - ) - - jax_problem = _import_jax(problem, "test_two_period_preeq", tmp_path) - assert jax_problem._max_periods == 1 - - x, _ = run_simulations(jax_problem, ret=ReturnValue.x) - expected = 2.0 * np.exp(-0.7 * np.array(ts)) - np.testing.assert_allclose( - np.asarray(x)[0, :, 0], expected, rtol=1e-4 - ) - - -def test_single_period_matches_analytical_solution(tmp_path): - """Regression check for the simplest (no pre-equilibration, single - period) case, P=1 with no chaining at all.""" - problem = _linear_decay_problem() - problem.add_condition("cond1", kk=0.7) - problem.add_experiment("exp1", 0.0, "cond1") - ts = (0.0, 1.0, 2.0, 3.0) - for t in ts: - problem.add_measurement( - "obs1", time=t, measurement=0.0, experiment_id="exp1" - ) - - jax_problem = _import_jax(problem, "test_single_period", tmp_path) - - x, _ = run_simulations(jax_problem, ret=ReturnValue.x) - # model's default initial value (xx=1) applies, no reinit here - expected = 1.0 * np.exp(-0.7 * np.array(ts)) - np.testing.assert_allclose( - np.asarray(x)[0, :, 0], expected, rtol=1e-4 - ) - - def test_gradient_through_multiperiod_chain_matches_analytical_derivative( tmp_path, ): diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index bbf8b6291d..8eee17e1cd 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -700,6 +700,68 @@ def dispatch(self, *args, **kwargs): ) +def test_simulate_condition_is_deprecated_alias_for_simulate_experiment( + tmp_path, +): + """``simulate_condition[_unjitted]`` (pre-rename names) must still work, + with a ``DeprecationWarning``, and produce the same result as the + ``simulate_experiment[_unjitted]`` methods they were renamed to.""" + from amici.importers.antimony import antimony2sbml + from amici.importers.sbml import SbmlImporter + from amici.sim.jax.petab import ( + DEFAULT_CONTROLLER_SETTINGS, + DEFAULT_ROOT_FINDER_SETTINGS, + SteadyStateEvent, + ) + + ant_model = """ + model simulate_condition_alias + x' = -k * x + x = 1 + k = 1 + end + """ + sbml = antimony2sbml(ant_model) + importer = SbmlImporter(sbml, from_file=False) + importer.sbml2jax("simulate_condition_alias", output_dir=tmp_path) + module = amici._module_from_path( + "simulate_condition_alias", tmp_path / "__init__.py" + ) + model = module.Model() + + ts = jnp.array([0.0, 1.0, 2.0]) + my = jnp.zeros_like(ts) + iys = jnp.zeros_like(ts, dtype=int) + iy_trafos = jnp.zeros_like(ts, dtype=int) + args = ( + jnp.array([[2.5]]), + ts[None, :], + jnp.zeros((1, 0)), + my[None, :], + iys[None, :], + iy_trafos[None, :], + jnp.zeros((1, 3, 0)), + jnp.zeros((1, 3, 0)), + diffrax.Kvaerno5(), + diffrax.PIDController(**DEFAULT_CONTROLLER_SETTINGS), + optimistix.Newton(**DEFAULT_ROOT_FINDER_SETTINGS), + diffrax.RecursiveCheckpointAdjoint(), + SteadyStateEvent(), + 1000, + ) + + expected_llh, _ = model.simulate_experiment(*args) + expected_llh_unjitted, _ = model.simulate_experiment_unjitted(*args) + + with pytest.warns(DeprecationWarning, match="simulate_experiment"): + actual_llh, _ = model.simulate_condition(*args) + with pytest.warns(DeprecationWarning, match="simulate_experiment_unjitted"): + actual_llh_unjitted, _ = model.simulate_condition_unjitted(*args) + + assert_allclose(float(actual_llh), float(expected_llh)) + assert_allclose(float(actual_llh_unjitted), float(expected_llh_unjitted)) + + @skip_on_valgrind def test_time_dependent_discontinuity(tmp_path): """Models with time dependent discontinuities are handled."""