rework petabv2 jax test cases with ExperimentsToSbmlEvents and no v1 parameter_mapping

Author: BSnelling

python/sdist/amici/importers/petab/v1/parameter_mapping.py

@@ -26,7 +26,6 @@
 import numpy as np
 import pandas as pd
 import petab.v1 as petab
-import petab.v2 as petabv2
 import sympy as sp
 from petab.v1.C import *  # noqa: F403
 from petab.v1.C import (
@@ -42,9 +41,6 @@
 
 from amici.importers.sbml import get_species_initial
 from amici.sim.sundials import AmiciModel, ParameterScaling
-from amici.sim.jax import (
-    _set_initial_state_v2, get_states_in_condition_table_v2, reformat_for_v2, fixup_v2_parameter_mapping
-)
 
 from . import PREEQ_INDICATOR_ID
 from .util import get_states_in_condition_table
@@ -384,15 +380,10 @@ def create_parameter_mapping(
         else petab_problem.mapping_df
     )
 
-    # Do some reformatting if V2 problem
-    measurement_df, condition_df = reformat_for_v2(petab_problem) if isinstance(
-        petab_problem, petabv2.Problem
-    ) else (petab_problem.measurement_df, petab_problem.condition_df)
-
     prelim_parameter_mapping = (
         petab.get_optimization_to_simulation_parameter_mapping(
-            condition_df=condition_df,
-            measurement_df=measurement_df,
+            condition_df=petab_problem.condition_df,
+            measurement_df=petab_problem.measurement_df,
             parameter_df=petab_problem.parameter_df,
             observable_df=petab_problem.observable_df,
             mapping_df=mapping,
@@ -409,11 +400,6 @@ def create_parameter_mapping(
     for (_, condition), prelim_mapping_for_condition in zip(
         simulation_conditions.iterrows(), prelim_parameter_mapping, strict=True
     ):
-        if isinstance(petab_problem, petabv2.Problem):
-            prelim_mapping_for_condition = fixup_v2_parameter_mapping(
-                prelim_mapping_for_condition, petab_problem
-            )
-
         mapping_for_condition = create_parameter_mapping_for_condition(
             prelim_mapping_for_condition,
             condition,
@@ -488,14 +474,9 @@ def create_parameter_mapping_for_condition(
     # ExpData.x0, but in the case of pre-equilibration this would not allow for
     # resetting initial states.
 
-    if isinstance(petab_problem, petabv2.Problem):
-        states_in_condition_table = get_states_in_condition_table_v2(
-            petab_problem, condition
-        )
-    else:
-        states_in_condition_table = get_states_in_condition_table(
-            petab_problem, condition
-        )
+    states_in_condition_table = get_states_in_condition_table(
+        petab_problem, condition
+    )
 
     if states_in_condition_table:
         # set indicator fixed parameter for preeq
@@ -536,26 +517,16 @@ def create_parameter_mapping_for_condition(
             # for simulation
             condition_id = condition[SIMULATION_CONDITION_ID]
             init_par_id = f"initial_{element_id}_sim"
-            if isinstance(petab_problem, petabv2.Problem):
-                _set_initial_state_v2(
-                    petab_problem,
-                    init_par_id,
-                    condition_map_sim,
-                    condition_scale_map_sim,
-                    value,
-                    fill_fixed_parameters=fill_fixed_parameters,
-                )
-            else:
-                _set_initial_state(
-                    petab_problem,
-                    condition_id,
-                    element_id,
-                    init_par_id,
-                    condition_map_sim,
-                    condition_scale_map_sim,
-                    value,
-                    fill_fixed_parameters=fill_fixed_parameters,
-                )
+            _set_initial_state(
+                petab_problem,
+                condition_id,
+                element_id,
+                init_par_id,
+                condition_map_sim,
+                condition_scale_map_sim,
+                value,
+                fill_fixed_parameters=fill_fixed_parameters,
+            )
                 # set dummy value as above
             if condition_map_preeq:
                 condition_map_preeq[init_par_id] = 0.0
@@ -609,17 +580,14 @@ def create_parameter_mapping_for_condition(
     )
     logger.debug(f"Variable parameters simulation: {condition_map_sim_var}")
 
-    if isinstance(petab_problem, petabv2.Problem):
-        pass
-    else:
-        petab.merge_preeq_and_sim_pars_condition(
-            condition_map_preeq_var,
-            condition_map_sim_var,
-            condition_scale_map_preeq_var,
-            condition_scale_map_sim_var,
-            condition,
-        )
-        logger.debug(f"Merged: {condition_map_sim_var}")
+    petab.merge_preeq_and_sim_pars_condition(
+        condition_map_preeq_var,
+        condition_map_sim_var,
+        condition_scale_map_preeq_var,
+        condition_scale_map_sim_var,
+        condition,
+    )
+    logger.debug(f"Merged: {condition_map_sim_var}")
 
     if "sciml" in petab_problem.extensions_config:
         hybridizations = [

python/sdist/amici/importers/petab/v1/petab_import.py

@@ -13,12 +13,10 @@
 
 import pandas as pd
 import petab.v1 as petab
-import petab.v2 as petabv2
 from petab.v1.models import MODEL_TYPE_PYSB, MODEL_TYPE_SBML
 
 import amici
 from amici.logging import get_logger
-from amici.sim.jax import reformat_petab_v2_to_v1, add_events_to_sbml
 
 from .import_helpers import (
     _can_import_model,
@@ -88,10 +86,6 @@ def import_petab_problem(
     :return:
         The imported model (if ``jax=False``) or JAX problem (if ``jax=True``).
     """
-    if isinstance(petab_problem, petabv2.Problem):
-        petab_problem = add_events_to_sbml(petab_problem)
-        petab_problem = reformat_petab_v2_to_v1(petab_problem)
-
     if petab_problem.model.type_id not in (MODEL_TYPE_SBML, MODEL_TYPE_PYSB):
         raise NotImplementedError(
             "Unsupported model type " + petab_problem.model.type_id

python/sdist/amici/jax/_simulation.py

@@ -174,6 +174,7 @@ def body_fn(carry):
 
 def solve(
     p: jt.Float[jt.Array, "np"],
+    t0: jnp.float_,
     ts: jt.Float[jt.Array, "nt_dyn"],
     tcl: jt.Float[jt.Array, "ncl"],
     h: jt.Float[jt.Array, "ne"],
@@ -195,6 +196,8 @@ def solve(
 
     :param p:
         parameters
+    :param t0:
+        initial time point
     :param ts:
         time points at which solutions are evaluated
     :param tcl:
@@ -226,7 +229,7 @@ def solve(
     if not root_cond_fns:
         # no events, we can just run a single segment
         sol, _, stats = _run_segment(
-            0.0,
+            t0,
             ts[-1],
             x0,
             p,
@@ -319,7 +322,7 @@ def body_fn(carry):
         body_fn,
         (
             jnp.zeros((ts.shape[0], x0.shape[0]), dtype=x0.dtype) + x0,
-            0.0,
+            t0,
             x0,
             jnp.zeros((ts.shape[0], h.shape[0]), dtype=h.dtype),
             h,

python/sdist/amici/jax/model.py

@@ -599,17 +599,27 @@ def simulate_condition_unjitted(
         init_override_mask: jt.Bool[jt.Array, "*nx"] = jnp.array([]),
         ts_mask: jt.Bool[jt.Array, "nt"] = jnp.array([]),
         h_mask: jt.Bool[jt.Array, "ne"] = jnp.array([]),
+        t_zero: jnp.float_ = 0.0,
         ret: ReturnValue = ReturnValue.llh,
     ) -> tuple[jt.Float[jt.Array, "*nt"], dict]:
         """
         Unjitted version of simulate_condition.
 
         See :meth:`simulate_condition` for full documentation.
         """
-        t0 = 0.0
+        t0 = t_zero
         if p is None:
             p = self.parameters
 
+        if os.getenv("JAX_DEBUG") == "1":
+            jax.debug.print(
+                "x_reinit: {}, x_preeq: {}, x_def: {}. p: {}",
+                x_reinit,
+                x_preeq,
+                self._x0(t0, p),
+                p,
+            )
+
         if x_preeq.shape[0]:
             x = x_preeq
         elif init_override.shape[0]:
@@ -626,6 +636,7 @@ def simulate_condition_unjitted(
         # 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)
 
@@ -645,6 +656,7 @@ def simulate_condition_unjitted(
         if ts_dyn.shape[0]:
             x_dyn, h_dyn, stats_dyn = solve(
                 p,
+                t0,
                 ts_dyn,
                 tcl,
                 h,
@@ -781,6 +793,7 @@ def simulate_condition(
         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]:
         r"""
@@ -863,6 +876,7 @@ def simulate_condition(
             init_override_mask,
             ts_mask,
             h_mask,
+            t_zero,
             ret,
         )
 
@@ -929,6 +943,7 @@ def preequilibrate_condition(
             tcl,
             h,
             current_x,
+            h_mask,
             solver,
             controller,
             root_finder,
@@ -992,7 +1007,7 @@ def _handle_t0_event(
 
         if os.getenv("JAX_DEBUG") == "1":
             jax.debug.print(
-                "h: {}, rf0: {}, rfx: {}, roots_found: {}, roots_dir: {}, h_next: {}, y0_next: {}, y0: {}",
+                "handle_t0_event h: {}, rf0: {}, rfx: {}, roots_found: {}, roots_dir: {}, h_next: {}, y0_next: {}, y0: {}",
                 h,
                 rf0,
                 rfx,

python/sdist/amici/jax/ode_export.py

@@ -150,20 +150,20 @@ def __init__(
                 "The JAX backend does not support models with algebraic states."
             )
 
-        if not ode_model.has_only_time_dependent_event_assignments():
-            raise NotImplementedError(
-                "The JAX backend does not support event assignments with explicit non-time dependent triggers."
-            )
+        # if not ode_model.has_only_time_dependent_event_assignments():
+        #     raise NotImplementedError(
+        #         "The JAX backend does not support event assignments with explicit non-time dependent triggers."
+        #     )
 
         if ode_model.has_priority_events():
             raise NotImplementedError(
                 "The JAX backend does not support event priorities."
             )
 
-        if ode_model.has_implicit_event_assignments():
-            raise NotImplementedError(
-                "The JAX backend does not support event assignments with implicit triggers."
-            )
+        # if ode_model.has_implicit_event_assignments():
+        #     raise NotImplementedError(
+        #         "The JAX backend does not support event assignments with implicit triggers."
+        #     )
 
         self.verbose: bool = logger.getEffectiveLevel() <= logging.DEBUG