Chain PEtab v2 experiment periods natively in the JAX simulator

python/sdist/amici/importers/petab/_petab_importer.py

@@ -212,13 +212,12 @@ def _preprocess_sbml(self):
         if not isinstance(self.petab_problem.model, SbmlModel):
             raise ValueError("The PEtab problem must contain an SBML model.")
 
-        # Convert petab experiments to events, because so far,
-        #  AMICI only supports preequilibration/presimulation/simulation, but
-        #  no arbitrary list of periods.
-        exp_event_conv = ExperimentsToSbmlConverter(self.petab_problem)
-        # This will always create a copy of the problem.
         if self._jax:
-            self._unconverted_problem = exp_event_conv._original_problem
+            # The JAX backend natively chains one ODE integration per
+            #  experiment period (see amici.sim.jax.petab), so there is no
+            #  need to convert experiments with more than two periods into
+            #  SBML events. The condition table is left untouched.
+            self._unconverted_problem = None
             condition_targets = {
                 change.target_id
                 for condition in self.petab_problem.conditions
@@ -232,13 +231,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:")
@@ -265,7 +287,15 @@ def _preprocess_pysb(self):
 
         pysb.bng.generate_equations(self.petab_problem.model.model)
 
-        # Convert PEtab v2 experiments/conditions to events
+        # Convert PEtab v2 experiments/conditions to events. Unlike for SBML
+        #  (see `_preprocess_sbml`), this is not skipped for the JAX backend:
+        #  PySB condition-table targets are frequently pysb.Observable
+        #  names that alias an underlying pysb.Initial/Expression rather
+        #  than a state or free parameter directly, and applying those
+        #  requires the same model-rewriting this converter already does.
+        #  JAXProblem's native per-period parameter/state resolution has no
+        #  equivalent for that, so PySB models keep going through event
+        #  conversion for both backends.
         converter = ExperimentsToPySBConverter(self.petab_problem)
         self.petab_problem, self._events = converter.convert()
 
@@ -333,15 +363,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
@@ -771,8 +813,6 @@ def create_simulator(
             Whether to force re-import even if the model module already exists.
         :return: The created PEtab simulator.
         """
-        from amici.sim.sundials.petab import ExperimentManager, PetabSimulator
-
         if self._jax:
             model_module = self.import_module(force_import=force_import)
             model = model_module.Model()
@@ -787,6 +827,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

@@ -250,7 +250,11 @@ def solve(
             term,
             dict(**STARTING_STATS),
         )
-        return sol.ys, jnp.repeat(h[None, :], sol.ys.shape[0]), stats
+        return (
+            sol.ys,
+            jnp.repeat(h[None, :], sol.ys.shape[0], axis=0),
+            stats,
+        )
 
     def cond_fn(carry):
         _, t_start, y0, _, _, stats = carry