Chain PEtab v2 experiment periods natively in the JAX simulator

.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 <amici_version>/ directory instead of under
+# amici_models/<amici_version>/.
+/[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/*
 

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",

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)

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