diff --git a/.gitignore b/.gitignore index 75f7ba71a5..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/ @@ -141,6 +147,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/* diff --git a/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb index ac6d054124..ad827625a7 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", @@ -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." ] }, { @@ -389,7 +389,7 @@ "from amici.sim.jax import ReturnValue\n", "\n", "# Define the simulation condition\n", - "experiment_condition = (\"_petab_experiment_condition___default__\",)\n", + "experiment_condition = \"__default__\"\n", "ic = 0\n", "\n", "# Load condition-specific data\n", @@ -410,8 +410,8 @@ "# 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", + " 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", diff --git a/python/sdist/amici/importers/petab/_petab_importer.py b/python/sdist/amici/importers/petab/_petab_importer.py index cb60923d63..600750999d 100644 --- a/python/sdist/amici/importers/petab/_petab_importer.py +++ b/python/sdist/amici/importers/petab/_petab_importer.py @@ -287,13 +287,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 @@ -307,13 +306,36 @@ 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 + # 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( - "AMICI currently does not support more than two periods." + "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/ + # 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:") @@ -340,7 +362,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() @@ -412,15 +442,27 @@ 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 + # 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 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 @@ -864,8 +906,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() @@ -880,6 +920,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 3ae12b56ce..d42c22ad36 100644 --- a/python/sdist/amici/sim/jax/_simulation.py +++ b/python/sdist/amici/sim/jax/_simulation.py @@ -266,7 +266,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 7ba1ec76ee..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 @@ -439,7 +440,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 +450,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 +477,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 +494,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 +502,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 +516,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 +533,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 +555,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 +574,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 +596,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. + Simulate a single experiment period, starting from ``x_solver``/``h`` + at time ``t0`` with parameters ``p``/``tcl``. - See :meth:`simulate_condition` for full documentation. - """ - t0 = t_zero - if p is None: - p = self.parameters + 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`. - 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) - - 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, - {}, - ) - - # 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 +633,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 +671,186 @@ 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_experiment_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"] | 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]: + """ + 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``/ + ``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_experiment` for full documentation. + """ + 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_) - 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 +859,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 +892,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) @@ -762,17 +908,28 @@ def simulate_condition_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_condition( + def simulate_experiment( 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, @@ -781,26 +938,32 @@ 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, "*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, + 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""" - 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_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``/ + ``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. @@ -850,7 +1013,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, @@ -877,6 +1040,16 @@ def simulate_condition( 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, @@ -972,21 +1145,25 @@ 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, ): rf0 = self.event_initial_values - 0.5 - if h_preeq.shape[0]: - # Dynamic phase following preequilibration: carry the event state - # out of preequilibration, but re-evaluate the triggers at t0 under - # the dynamic-period parameters, which may differ from the - # preequilibration ones (e.g. a stimulus whose onset time is a - # condition-specific parameter that is inactive during - # preequilibration). Events already active after preequilibration - # keep their state and are not re-fired (no sign change), while a - # trigger that differs under the dynamic parameters is corrected. - h = jnp.where(h_mask, h_preeq, jnp.ones_like(h_preeq)) + 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). 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, or the dynamic + # phase may use different parameters than preequilibration did + # (e.g. a stimulus onset time that is inactive during + # preequilibration), 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_prev` only supplies the pre-transition reference value. + h = jnp.where(h_mask, h_prev, jnp.ones_like(h_prev)) rf0 = jnp.where(h > 0.5, 0.5, -0.5) else: h = jnp.where(h_mask, jnp.heaviside(rf0, 0.0), jnp.ones_like(rf0)) diff --git a/python/sdist/amici/sim/jax/petab.py b/python/sdist/amici/sim/jax/petab.py index 72d49e0374..cd340a9481 100644 --- a/python/sdist/amici/sim/jax/petab.py +++ b/python/sdist/amici/sim/jax/petab.py @@ -6,6 +6,7 @@ import shutil from collections.abc import Callable, Iterable, Sized from pathlib import Path +from typing import NamedTuple import diffrax import equinox as eqx @@ -17,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, @@ -51,7 +54,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 @@ -94,39 +97,302 @@ def jax_unscale( raise ValueError(f"Invalid parameter scaling: {scale_str}") -def _get_period_condition_ids( - exp: petabv2.Experiment, is_preequilibration: bool -) -> tuple[str, ...]: - """Get the condition ids of an experiment's (pre)equilibration period. - - A period may reference multiple condition ids applied simultaneously - (PEtab v2 requires their targets to be disjoint); all of them are - returned so callers can consider each one, rather than collapsing them - into a single id that would not correspond to any row of - :attr:`petab.v2.Problem.condition_df`. - - :param exp: - PEtab v2 experiment. - :param is_preequilibration: - If ``True``, look for the preequilibration period - (:attr:`petabv2.ExperimentPeriod.is_preequilibration`, i.e. - ``time == -inf``). If ``False``, look for the dynamic - (non-preequilibration) period. +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 _get_fixed_parameter_values( + petab_problem: petabv2.Problem, +) -> dict[str, float]: + """Nominal (linear) values of fixed (non-estimated) parameters, used to + resolve observable/noise parameter overrides that reference them. + + Array-valued (PEtab-SciML) fixed parameters are excluded: their nominal + value is the literal string ``"array"``, not a substitutable number. + + Built once (from :attr:`petabv2.Problem.parameter_tables` directly) + rather than via the :attr:`petabv2.Problem.parameter_df` property on + every override column, since building the latter can be expensive and, + for some problems (e.g. array-valued PEtab-SciML parameters), emits + pydantic serialization warnings. + """ + return { + p.id: p.nominal_value + for pt in petab_problem.parameter_tables + for p in pt.elements + if not p.estimate and p.nominal_value != "array" + } + + +def _split_override_column( + col_values: pd.Series, + n_pars: int, + fixed_parameter_values: dict[str, float], +) -> 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: + # `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 = ( + # 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] + ) + # 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( + 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, ...] +) -> 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) + par_mask = par_index != -1 + # 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 OverrideColumn(mat, par_mask, par_index) + + +def _column_overrides( + m: pd.DataFrame, + col: str, + n_pars: int, + fixed_parameter_values: dict[str, float], + parameter_ids: tuple[str, ...], +) -> OverrideColumn: + """Numeric values, non-numeric mask and parameter indices for one + observable/noise parameter override column of the rows in ``m``.""" + if col not in m or m[col].isna().all() or (m[col] == "").all(): + 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 OverrideColumn( + mat_numeric, + np.zeros_like(mat_numeric, dtype=bool), + np.zeros_like(mat_numeric, dtype=int), + ) + mat = _split_override_column(m[col], n_pars, fixed_parameter_values) + return _override_triple_from_matrix(mat, parameter_ids) + + +def _get_overrides( + m: pd.DataFrame, + n_pars: dict[str, int], + fixed_parameter_values: dict[str, float], + 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], fixed_parameter_values, parameter_ids + ) + for col in (petabv2.C.OBSERVABLE_PARAMETERS, petabv2.C.NOISE_PARAMETERS) + } + + +def _get_iy_trafos( + iys: np.ndarray, + petab_problem: petabv2.Problem, + observable_ids: list[str], +) -> np.ndarray: + """Observable transformation index (see ``SCALE_TO_INT``) for each + (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: + 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 + } + trafo_by_index = np.array( + [trafo_by_id[oid] for oid in observable_ids] + ) + return trafo_by_index[iys] + 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. + + 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_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( + 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: - The condition ids of the matching period, in period order. - :raises ValueError: - If ``exp`` has no matching period with a non-empty - ``condition_ids``. + Tuple of ``(ts_dyn, dyn_valid, my, iys, iy_trafos, overrides)``. """ - for period in exp.periods: - if period.is_preequilibration != is_preequilibration: - continue - if period.condition_ids: - return tuple(period.condition_ids) - - kind = "preequilibration" if is_preequilibration else "dynamic" - raise ValueError( - f"Experiment {exp.id!r} has no {kind} period with a condition id." + 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_dyn: Callable[[_PeriodMeasurements], np.ndarray], + extractor_posteq: Callable[[_PeriodMeasurements], np.ndarray], + n_ts_dyn: int, + n_ts_posteq: int, +) -> np.ndarray: + """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_dyn(mv), + extractor_posteq(mv), + n_ts_dyn, + n_ts_posteq, + ) + for mv in measurements.values() + ] ) @@ -150,6 +416,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 @@ -165,6 +432,7 @@ class JAXProblem(eqx.Module): _petab_measurement_indices: np.ndarray _petab_problem: petabv2.Problem _unconverted_problem: petabv2.Problem | None + _all_condition_targets: frozenset[str] def __init__( self, @@ -191,11 +459,17 @@ def __init__( self.simulation_conditions = scs.conditionId.to_list() self._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) ) self._parameter_mappings = self._get_parameter_mappings() ( + self._max_periods, self._ts_dyn, self._ts_posteq, self._my, @@ -209,7 +483,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): """ @@ -248,8 +522,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, @@ -265,11 +540,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 @@ -278,7 +563,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 @@ -286,17 +572,15 @@ def _get_measurements( - non-numeric mask for noise parameter overrides - parameter indices (problem parameters) for noise parameter overrides """ - measurements = dict() - petab_indices = dict() + measurements: dict[tuple[str, int], _PeriodMeasurements] = {} # Nominal (linear) values of fixed (non-estimated) parameters, used to # resolve observable/noise parameter overrides that reference them. - fixed_parameter_values = { - p.id: p.nominal_value - for pt in self._petab_problem.parameter_tables - for p in pt.elements - if not p.estimate and p.nominal_value != "array" - } + # Built once up front (rather than lazily per override column) since + # it is otherwise rebuilt once per period. + fixed_parameter_values = _get_fixed_parameter_values( + self._petab_problem + ) n_pars = dict() for col in [ @@ -325,206 +609,317 @@ def _get_measurements( .max() ) - for _, simulation_condition in simulation_conditions.iterrows(): - 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 - ] - ) - 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 - ] - ) - if ( - petabv2.C.NOISE_DISTRIBUTION - in self._petab_problem.observable_df - ): - # Map each measurement row to its observable's noise trafo so - # ``iy_trafos`` is aligned with the other per-measurement arrays - # (``my``, ``iys``, ...) of length ``len(m)``, as required by the - # dyn/post-eq padding split. - obs_trafo = { - 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 self._petab_problem.observables - } - iy_trafos = np.array( - [ - obs_trafo[oid] - for oid in m[petabv2.C.OBSERVABLE_ID].values - ] - ) - 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): - # Substitute fixed parameters with their nominal value; leave - # estimated-parameter names untouched (mapped to problem - # parameters later via ``par_index``). - return fixed_parameter_values.get(x, x) - - 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])) - 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) + 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=0, + ) + 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. + 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): + 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. 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([]) + 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: - split_vals = m[col].str.split( - petabv2.C.PARAMETER_SEPARATOR + 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 ) - list_vals = split_vals.apply( - lambda x: ( - # drop empty tokens (empty cells split to ``['']``); - # an empty override list is padded to the neutral 1.0 - [get_parameter_override(y) for y in x if y != ""] - if isinstance(x, list) - else [] - if pd.isna(x) - else [x] - ) # every string gets transformed to lists, so this is already a float + in_window = ( + (m_full_times >= t_lo) + & (m_full_times < t_hi) + & np.isfinite(m_full_times) ) - vals = list_vals.apply( - lambda x: np.pad( - x, - (0, n_pars[col] - len(x)), - mode="constant", - constant_values=1.0, - ) + 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 + ], + dtype=int, + ) + iy_trafos_real = _get_iy_trafos( + iys_real, + self._petab_problem, + self.model.observable_ids, ) - 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 + overrides_real = _get_overrides( + m, n_pars, fixed_parameter_values, self.parameter_ids + ) + 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 = tuple(m_posteq.index) + + 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, ) - )(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 - - 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 - ) - petab_indices[tuple(simulation_condition)] = tuple(index.tolist()) + + # 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_dyn, iys_dyn, iy_trafos_dyn = ( + my_real, + iys_real, + iy_trafos_real, + ) + overrides_dyn = overrides_real + index_dyn = tuple(index_dyn_real) + 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, + dyn_valid, + my_dyn, + iys_dyn, + iy_trafos_dyn, + overrides_dyn, + ) = _masked_placeholder_period(t_lo, n_pars) + index_dyn = (-1,) + else: + ts_posteq = np.array([]) + 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 + # 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_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 = (*index_dyn_real, -1) + + posteq_valid = np.ones(len(ts_posteq), dtype=bool) + + measurements[(exp.id, i_period)] = _PeriodMeasurements( + ts_dyn=ts_dyn, + ts_posteq=ts_posteq, + 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, + ) # 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") + 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") + 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"), - np.pad(x_peq, pad_width_peq, mode="edge"), - ) - ) - - 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_dyn, + lambda mv: mv.my_posteq, + n_ts_dyn, + n_ts_posteq, + ) + iys = _pad_and_stack( + 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_dyn, + lambda mv: mv.iy_trafos_posteq, + n_ts_dyn, + n_ts_posteq, + ) + op_numeric = _pad_and_stack( + measurements, + 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_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_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_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_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_dyn.index, + lambda mv: mv.noise_overrides_posteq.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( - np.ones_like(mv[0]), (0, n_ts_dyn - len(mv[0])) + mv.valid_dyn, + (0, n_ts_dyn - len(mv.ts_dyn)), ), np.pad( - np.ones_like(mv[1]), (0, n_ts_posteq - len(mv[1])) + mv.valid_posteq, + (0, n_ts_posteq - len(mv.ts_posteq)), ), ) ) @@ -533,17 +928,18 @@ 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]) :]), - ) - for mv, idx in zip( - measurements.values(), petab_indices.values() + _pad_measurement( + np.array(mv.index_dyn), + np.array(mv.index_posteq), + n_ts_dyn, + n_ts_posteq, ) + for mv in measurements.values() ] ) - return ( + n_exp = len(experiments) + outputs = ( ts_dyn, ts_posteq, my, @@ -558,33 +954,28 @@ def pad_and_stack(output_index: int): np_mask, np_indices, ) - - def _resolve_condition_target_value(self, target_value): - """Resolve a condition change's target value to a number. - - A condition table target value may be a numeric literal, or a - reference to another PEtab parameter id (to be substituted with - that parameter's current value, e.g. to share an estimated - parameter's value across multiple conditions). - """ - if not target_value.is_number: - pname = str(target_value) - if pname in self.parameter_ids: - return self.parameters[self.parameter_ids.index(pname)] - _petab_param_map = { - param.id: param.nominal_value - for param in self._petab_problem.parameters - } - if pname in _petab_param_map: - return _petab_param_map[pname] - return jnp.asarray(target_value, dtype=self.model.parameters.dtype) + return ( + max_periods, + *( + arr.reshape(n_exp, max_periods, *arr.shape[1:]) + for arr in outputs + ), + ) def _get_parameter_mappings(self) -> dict[str, ...]: + # `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: self._resolve_condition_target_value( - ch.target_value - ) + ch.target_id: _resolve_petab_change_value(ch.target_value) for ch in c.changes } for c in self._petab_problem.conditions @@ -603,6 +994,49 @@ def _get_parameter_mappings(self) -> dict[str, ...]: return {"targets_map": targets_map, "hybrid_map": hybrid_map} + def _resolve_parameter_reference( + self, value: "_CompiledConditionExpr" + ) -> jt.Float[jt.Scalar, ""]: # noqa: F722 + """ + Resolve a value from ``targets_map`` (see :meth:`_get_parameter_mappings`) + 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 compiled expression, as returned by + :func:`_resolve_petab_change_value`. + """ + + 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 + ], + 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( self._petab_problem @@ -1095,8 +1529,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. @@ -1105,13 +1552,23 @@ 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. """ + if not self.model.parameter_ids: + # a model with no free SBML parameters (e.g. only literal rate + # constants); `jnp.stack` of an empty sequence raises. + return jnp.array([]) 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) ] @@ -1126,6 +1583,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. @@ -1134,14 +1592,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 @@ -1154,26 +1618,17 @@ def _map_experiment_model_parameter_value( else: init_val = self.model.parameters[p_index] - # Resolve condition-table overrides *live* from the raw changes rather - # than reusing the ``targets_map`` values cached at construction time. - # This method runs inside the traced/differentiated region (via - # ``_prepare_experiments``), so reading ``self.parameters`` (through - # ``_resolve_condition_target_value``) keeps the value a function of - # the current parameters -- gradients w.r.t. an (e.g. condition- - # specific) estimated parameter that a condition maps this one to flow, - # and re-simulating after ``update_parameters`` reflects the new value. - raw_targets = { - change.target_id: change.target_value - for c in self._petab_problem.conditions - if c.id in condition_ids - for change in c.changes + targets_filtered = { + param: value + for condition, target in self._parameter_mappings[ + "targets_map" + ].items() + for param, value in target.items() + if condition in condition_ids } - if pname in raw_targets: - return jnp.asarray( - self._resolve_condition_target_value(raw_targets[pname]), - dtype=self.model.parameters.dtype, - ) + if pname in targets_filtered: + return self._resolve_parameter_reference(targets_filtered[pname]) elif pname in self._parameter_mappings["hybrid_map"]: return jnp.asarray( self._eval_nn( @@ -1211,18 +1666,45 @@ 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. + + 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. + :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_conditions: tuple[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_conditions: - condition ids simultaneously active for the simulation condition - to check reinitialisation for (PEtab v2 requires their targets - to be disjoint, so at most one of them defines ``state_id``) + :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: @@ -1234,114 +1716,123 @@ def _state_needs_reinitialisation( if state_id in self._parameter_mappings["hybrid_map"]: return True - return ( - self._condition_reinit_target_value( - simulation_conditions, state_id - ) - is not None - ) - - def _condition_reinit_target_value( - self, simulation_conditions: tuple[str, ...], state_id: str - ): - """Return the *raw* condition-table target value initialising a state. - - Looks up the (unresolved) ``target_value`` that (re)initialises - ``state_id`` for the given simultaneously-active conditions, reading - it straight from the condition-table changes. Returns ``None`` if no - active condition sets ``state_id`` (PEtab v2 requires the targets of - simultaneously-active conditions to be disjoint, so at most one does). - - The raw value is returned deliberately -- callers resolve it *live* - against :attr:`parameters` (see - :meth:`_state_reinitialisation_value`), rather than reusing the - ``targets_map`` value cached at construction time, so that gradients - w.r.t. parameters used as initial values are not silently dropped. - """ - for condition in simulation_conditions: - for c in self._petab_problem.conditions: - if c.id != condition: - continue - for change in c.changes: - if change.target_id != state_id: - continue - # NaN targets (e.g. "use the preequilibration/SBML value") - # are dropped during v1->v2 conversion, but guard anyway - if ( - change.target_value.is_number - and change.target_value.is_finite is False - ): - return None - return change.target_value - return None + return self._first_condition_value(condition_ids, state_id) is not None def _state_reinitialisation_value( self, - simulation_conditions: tuple[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_conditions: - condition ids simultaneously active for the simulation condition - to get the reinitialisation value for (PEtab v2 requires their - targets to be disjoint, so at most one of them defines - ``state_id``) + :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: + parameters for the simulation condition :return: reinitialisation value for the state """ if state_id in self.nn_output_ids: - return self._eval_nn(state_id, simulation_conditions[0]) + 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_conditions[0], + condition_ids[0], ) - target_value = self._condition_reinit_target_value( - simulation_conditions, state_id - ) - if target_value is not None: - # Resolve *live* against ``self.parameters`` (not via the - # construction-time ``targets_map`` cache): this method runs inside - # the traced/differentiated region (via ``_prepare_experiments``), - # so reading ``self.parameters`` here keeps the reinitialisation - # value a function of the current parameters -- gradients flow and - # re-simulating after ``update_parameters`` reflects the new value. - return self._resolve_condition_target_value(target_value) - # no reinitialisation, return dummy value - return 0.0 + if ( + xval := self._first_condition_value(condition_ids, state_id) + ) is None: + # no reinitialisation, return dummy value + return 0.0 + + 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 + ) + 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[ + 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, - simulation_conditions: str | tuple[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_conditions: - Condition id(s) simultaneously active for the 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 mask and value for states. + Tuple of reinitialisation masm and value for states. """ - if isinstance(simulation_conditions, str): - simulation_conditions = (simulation_conditions,) + if isinstance(condition_ids, str): + condition_ids = [condition_ids] - needs_reinit = [ - self._state_needs_reinitialisation(simulation_conditions, x_id) + has_reinitialisable_states = any( + 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 - ] - # Always return full-length arrays per condition; callers stack/vmap across conditions and require consistent shapes. + ) + 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(needs_reinit) + mask = jnp.array( + [ + 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_conditions, x_id) + self._state_reinitialisation_value(condition_ids, x_id, p) for x_id in self.model.state_ids ] ) @@ -1356,22 +1847,9 @@ def update_parameters(self, p: jt.Float[jt.Array, "np"]) -> "JAXProblem": """ return eqx.tree_at(lambda p: p.parameters, self, p) - def _experiment_indices(self, experiments) -> np.ndarray: - """Positions of ``experiments`` in the full experiment ordering. - - All ``self._*`` per-experiment arrays are built in ``__init__`` keyed by - ``self._petab_problem.experiments`` order, so simulating a subset - (``simulation_experiments``) requires indexing them by these positions. - """ - positions = { - exp.id: i for i, exp in enumerate(self._petab_problem.experiments) - } - return np.array([positions[exp.id] for exp in experiments], dtype=int) - 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, @@ -1380,19 +1858,27 @@ 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_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 + 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. :param is_preeq: Whether to load preequilibration or simulation parameters. :param op_numeric: @@ -1411,23 +1897,81 @@ 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 + ] + ) - # one row per simulated experiment (aligned with ``p_array``); every - # simulated experiment has all of its events active. - h_mask = jnp.stack( - [jnp.ones(self.model.n_events) for _ in experiments] - ) + def preeq_condition_ids(exp: petabv2.Experiment) -> list[str]: + for period in exp.sorted_periods: + if period.is_preequilibration: + return period.condition_ids + return [] - t_zeros = jnp.stack( - [ - 0.0 - if exp.periods[0].is_preequilibration - else exp.periods[0].time + reinit_condition_ids = [ + preeq_condition_ids(exp) for exp in experiments + ] + 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] + + h_mask = jnp.stack( + [ + jnp.ones(self.model.n_events) + if (exp_id in exp_ids) + else jnp.zeros(self.model.n_events) + for exp_id in all_exp_ids + ] ) if self.parameters.size: @@ -1451,38 +1995,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( - (len(experiments), self._ts_masks.shape[1], 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( - (len(experiments), self._ts_masks.shape[1], 0) - ) + np_array = jnp.zeros((*self._ts_masks.shape, 0)) - reinit_arrays = [self.load_reinitialisation(sc) for sc in conditions] - mask_reinit_array = jnp.stack([m for m, _ in reinit_arrays]) - x_reinit_array = jnp.stack([x for _, x in reinit_arrays]) + 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, @@ -1568,7 +2152,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: @@ -1579,7 +2163,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)), @@ -1648,21 +2232,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. """ - # one entry per experiment, aligned with `experiments` (and thus with - # `p_array` built from it in `_prepare_experiments`), not a - # deduplicated set of condition names. - dynamic_conditions = [ - _get_period_condition_ids(exp, is_preequilibration=False) - for exp in experiments - ] - - # Positions of the requested experiments within the full, __init__-time - # ordering that all `self._*` per-experiment arrays are keyed by. When a - # subset is simulated (``simulation_experiments``), every per-experiment - # array must be indexed by these positions so its leading dimension - # matches ``p_array`` under the vmap. - ei = self._experiment_indices(experiments) - ( p_array, mask_reinit_array, @@ -1673,14 +2242,13 @@ def run_simulations( t_zeros, ) = self._prepare_experiments( experiments, - dynamic_conditions, False, - self._op_numeric[ei], - self._op_mask[ei], - self._op_indices[ei], - self._np_numeric[ei], - self._np_mask[ei], - self._np_indices[ei], + self._op_numeric, + self._op_mask, + self._op_indices, + self._np_numeric, + self._np_mask, + self._np_indices, ) init_override_mask = jnp.stack( @@ -1699,7 +2267,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 @@ -1712,11 +2280,11 @@ def run_simulations( return self.run_simulation( p_array, - self._ts_dyn[ei], - self._ts_posteq[ei], - self._my[ei], - self._iys[ei], - self._iy_trafos[ei], + self._ts_dyn, + self._ts_posteq, + self._my, + self._iys, + self._iy_trafos, op_array, np_array, mask_reinit_array, @@ -1731,7 +2299,7 @@ def run_simulations( max_steps, preeq_array, h_preeqs, - self._ts_masks[ei], + self._ts_masks, t_zeros, jnp.arange(len(experiments)), ret, @@ -1803,18 +2371,8 @@ def run_preequilibrations( ], max_steps: jnp.int_, ): - # one entry per experiment, aligned with `experiments` (and thus with - # `p_array` built from it in `_prepare_experiments`), not a - # deduplicated set of condition names. - preequilibration_conditions = [ - _get_period_condition_ids(exp, is_preequilibration=True) - for exp in 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, @@ -1841,7 +2399,7 @@ def run_simulations( ), steady_state_event: Callable[ ..., diffrax._custom_types.BoolScalarLike - ] = SteadyStateEvent(), + ] = diffrax.steady_state_event(), max_steps: int = 2**13, ret: ReturnValue | str = ReturnValue.llh, ): @@ -1886,24 +2444,19 @@ def run_simulations( if exp.id in simulation_experiments ] - # one entry per experiment, aligned with `experiments` (and thus with the - # rows of `problem._iys`/`problem._ts_masks` built from it), not a - # deduplicated set of condition names. + preeq_condition_ids = _get_preequilibration_condition_ids(experiments) dynamic_conditions = [ - _get_period_condition_ids(exp, is_preequilibration=False) + _dynamic_period_label(exp, period, preeq_condition_ids) for exp in experiments + for period in exp.sorted_periods + if not period.is_preequilibration ] + dynamic_conditions = list(dict.fromkeys(dynamic_conditions)) conditions = { "dynamic_conditions": dynamic_conditions, - # experiment ids aligned with `dynamic_conditions` and the rows of - # `_iys`/`_ts_masks`, so result-building need not reverse-map a - # condition id back to its experiment - "experiment_ids": [exp.id for exp in experiments], } - has_preeq = any(exp.periods[0].is_preequilibration for exp in experiments) - - if has_preeq: + if has_preeq := any(exp.has_preequilibration for exp in experiments): preeqs, preresults, h_preeqs = problem.run_preequilibrations( experiments, solver, @@ -1951,7 +2504,7 @@ def petab_simulate( ), steady_state_event: Callable[ ..., diffrax._custom_types.BoolScalarLike - ] = SteadyStateEvent(), + ] = diffrax.steady_state_event(), max_steps: int = 2**13, ): """ @@ -1980,7 +2533,7 @@ def petab_simulate( ret=ReturnValue.y, ) if isinstance(problem._petab_problem, petabv2.Problem): - return _build_simulation_df_v2(problem, y, r["experiment_ids"]) + return _build_simulation_df_v2(problem, y, r["dynamic_conditions"]) else: dfs = [] for ic, sc in enumerate(r["dynamic_conditions"]): @@ -2052,18 +2605,14 @@ def add_default_experiment_names_to_v2_problem(petab_problem: petabv2.Problem): petab_problem.experiment_df is None or petab_problem.experiment_df.empty ): - # read condition ids from the condition table elements, not - # `condition_df`: a condition with no changes (e.g. the just-added - # default condition, or any other no-op condition) contributes zero - # rows to the long-format `condition_df`, so its id could not be - # recovered from there. + # 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 table in petab_problem.condition_tables - for c in table.elements - ] - condition_ids = [ - c for c in condition_ids if "preequilibration" not in c + for c in petab_problem.conditions + if "preequilibration" not in c.id ] default_experiment = petabv2.core.Experiment( id="__default__", @@ -2089,92 +2638,168 @@ def get_simulation_conditions_v2(petab_problem) -> pd.DataFrame: """Get simulation conditions from PEtab v2 measurement DataFrame. Returns: - A pandas DataFrame mapping experiment_ids to condition ids, one row - per experiment. + A pandas DataFrame mapping experiment_ids to condition ids. """ experiment_df = petab_problem.experiment_df + exps = {} + for exp_id in experiment_df[petabv2.C.EXPERIMENT_ID].unique(): + exps[exp_id] = experiment_df[ + experiment_df[petabv2.C.EXPERIMENT_ID] == exp_id + ][petabv2.C.CONDITION_ID].unique() - # drop preequilibration periods - experiment_df = experiment_df[ - experiment_df[petabv2.C.TIME] != petabv2.C.TIME_PREEQUILIBRATION - ] experiment_df = experiment_df.drop(columns=[petabv2.C.TIME]) - # a dynamic period may reference multiple condition ids (e.g. the - # synthetic preequilibration-indicator condition alongside the actual - # experiment condition); measurements are only ever queried by - # experiment id (see `JAXProblem._get_measurements`), so collapse to - # one row per experiment -- otherwise arrays built per condition row - # here and arrays built per experiment elsewhere (e.g. `p_array` in - # `_prepare_experiments`) end up with mismatched batch sizes. - experiment_df = experiment_df.drop_duplicates( - subset=[petabv2.C.EXPERIMENT_ID] - ) return experiment_df -def _build_simulation_df_v2(problem, y, experiment_ids): - """Build a PEtab simulation DataFrame from PEtab v2 simulation results. - - ``experiment_ids`` is aligned with the rows of ``y`` / - ``problem._iys`` / ``problem._ts_masks`` (one entry per simulated - experiment). +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. """ - # ``y`` rows follow ``experiment_ids`` (the simulated subset), but the - # per-experiment ``problem._*`` arrays are keyed by the full experiment - # ordering, so map each simulated experiment id back to its full position. - full_positions = { - exp.id: i for i, exp in enumerate(problem._petab_problem.experiments) - } + 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 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 + 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 + label = _dynamic_period_label(exp, period, preeq_ids) + positions.setdefault(label, (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, experiment_id in enumerate(experiment_ids): - ie = full_positions[experiment_id] - # the synthetic default experiment id is reported as NaN, but the - # original id is still needed below to query the measurement table - reported_experiment_id = ( - jnp.nan if experiment_id == "__default__" else experiment_id - ) + 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 - # `_get_measurements` pads every experiment's arrays to a common - # length; apply the per-experiment mask consistently to the index - # and every column so experiments with fewer timepoints don't cause - # length mismatches or leak padded/duplicated measurement indices. - mask = problem._ts_masks[ie, :] + mask = problem._ts_masks[exp_idx, period_idx, :] obs = [ - problem.model.observable_ids[io] for io in problem._iys[ie, mask] + problem.model.observable_ids[io] + for io in problem._iys[exp_idx, period_idx, mask] ] - t = jnp.concat((problem._ts_dyn[ie, :], problem._ts_posteq[ie, :]))[ - mask - ] - n = len(t) + t = jnp.concatenate( + ( + 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")] * n, + petabv2.C.MODEL_ID: [float("nan")] * n_real, petabv2.C.OBSERVABLE_ID: obs, - petabv2.C.EXPERIMENT_ID: [reported_experiment_id] * n, - petabv2.C.TIME: t, - petabv2.C.SIMULATION: y[ic, mask], + petabv2.C.EXPERIMENT_ID: [experiment_id] * n_real, + petabv2.C.TIME: t[mask], + petabv2.C.SIMULATION: y_period[mask], }, - index=problem._petab_measurement_indices[ie, mask], + 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` 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 + ] + 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() +def _conditions_to_experiment_map( + experiment_df: pd.DataFrame, +) -> dict[str, str]: + return { + row.conditionId: row.experimentId + for row in experiment_df.itertuples() + } + + +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]]: @@ -2226,3 +2851,54 @@ def _try_float(value): if isinstance(e, ValueError) and "could not convert" in msg: return value raise + + +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__`. + """ + + 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 :class:`_CompiledConditionExpr`. + """ + return _CompiledConditionExpr.compile(sp.sympify(target_value)) 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..47a2a7ae2f --- /dev/null +++ b/python/tests/petab_/test_petab_v2_multiperiod.py @@ -0,0 +1,284 @@ +"""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_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. + + 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 + ) + 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", + ) + 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=measurement_value, + 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) + + 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(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_reevaluated_after_period_reinit( + tmp_path, +): + """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: 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) + 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_reevaluated", 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 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_cross1: + return 5.0 * np.exp(-1.0 * t) + return 2.0 * np.exp(-0.1 * (t - t_cross1)) + + # 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(0.3), + period2(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 + 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/python/tests/test_jax.py b/python/tests/test_jax.py index 230b42e601..8eee17e1cd 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_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`/ + # `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_experiment_unjitted) + + def fun(p, **kw): + return fun_periodic(p[None, :], **kw) for output in ["llh", "x0", "x", "y", "res"]: okwargs = kwargs | { @@ -440,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 @@ -561,19 +653,19 @@ 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( - jnp.array([k_val]), - ts, - jnp.array([]), - my, - iys, - iy_trafos, - jnp.zeros((3, 0)), - jnp.zeros((3, 0)), + model.simulate_experiment( + 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"], @@ -583,7 +675,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" ) @@ -608,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.""" diff --git a/tests/performance/test_jax_regression.py b/tests/performance/test_jax_regression.py index 0ba530e38a..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,23 +89,31 @@ def _sim_kwargs(model, solver_kwargs) -> dict: model_name = type(model).__name__ ts_dyn, my, iys, iy_trafos, ops, nps = _MAP[model_name] + # `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( - 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, ) 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) + # `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,14 +134,15 @@ 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) - 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( @@ -154,11 +163,12 @@ 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): - 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 @@ -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), @@ -189,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/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() diff --git a/tests/sbml/testSBMLSuite.py b/tests/sbml/testSBMLSuite.py index ca74e1dfe4..3b4cf4a170 100755 --- a/tests/sbml/testSBMLSuite.py +++ b/tests/sbml/testSBMLSuite.py @@ -220,15 +220,15 @@ def jax_sensitivity_check( root_finder = optimistix.Newton(**DEFAULT_ROOT_FINDER_SETTINGS) def simulate(pars): - x, _ = jax_model.simulate_condition( - 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)), + x, _ = jax_model.simulate_experiment( + 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, diff --git a/tests/sbml/testSBMLSuiteJax.py b/tests/sbml/testSBMLSuiteJax.py index 4bf9f34e2b..15b00ecce3 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) - 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)), + # `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_experiment( + 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,