tidying - add docstrings - rm outputs in notebook

Author: BSnelling

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

@@ -2720,6 +2720,7 @@ def _args_containing_t(self, expr):
     def has_implicit_event_assignments(self) -> bool:
         """
         Checks whether the model has event assignments with implicit triggers
+        (i.e. triggers that are not time based).
 
         :return:
             boolean indicating if event assignments with implicit triggers are present

python/sdist/amici/_symbolic/de_model.py

@@ -877,6 +877,9 @@ def get_trigger_times(self) -> set[sp.Expr]:
         return set(self._t_root)
 
     def _implicit_symbols(self):
+        """Get implicit symbols in the event trigger function.
+        That is, all symbols except time and petab indicator variables.
+        """
         symbols = [str(s) for s in list(self.get_val().free_symbols)]
         implicit_symbols = []
         for s in symbols:

python/sdist/amici/_symbolic/de_model_components.py

@@ -474,11 +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.
 
-    states_in_condition_table = get_states_in_condition_table(
+    if states_in_condition_table := get_states_in_condition_table(
         petab_problem, condition
-    )
-
-    if states_in_condition_table:
+    ):
         # set indicator fixed parameter for preeq
         # (we expect here, that this parameter was added during import and
         # that it was not added by the user with a different meaning...)
@@ -527,7 +525,7 @@ def create_parameter_mapping_for_condition(
                 value,
                 fill_fixed_parameters=fill_fixed_parameters,
             )
-                # set dummy value as above
+            # set dummy value as above
             if condition_map_preeq:
                 condition_map_preeq[init_par_id] = 0.0
                 condition_scale_map_preeq[init_par_id] = LIN

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

@@ -615,15 +615,6 @@ def simulate_condition_unjitted(
         if not h_mask.shape[0]:
             h_mask = jnp.ones(self.n_events, dtype=jnp.bool_)
 
-        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]:
@@ -913,6 +904,9 @@ def preequilibrate_condition(
             re-initialized state vector. If not provided, the state vector is not re-initialized.
         :param mask_reinit:
             mask for re-initialization. If `True`, the corresponding state variable is re-initialized.
+        :param h_mask:
+            mask for heaviside variables. If `True`, the corresponding heaviside variable is updated during simulation, otherwise it 
+            it marked as 1.0.
         :param solver:
             ODE solver
         :param controller:
@@ -1025,15 +1019,13 @@ def _handle_t0_event(
 
         if os.getenv("JAX_DEBUG") == "1":
             jax.debug.print(
-                "handle_t0_event h: {}, rf0: {}, rfx: {}, roots_found: {}, roots_dir: {}, h_next: {}, y0_next: {}, y0: {}",
+                "h: {}, rf0: {}, rfx: {}, roots_found: {}, roots_dir: {}, h_next: {}",
                 h,
                 rf0,
                 rfx,
                 roots_found,
                 roots_dir,
                 h_next,
-                y0_next,
-                y0,
             )
 
         return y0_next, t0_next, h_next, stats

python/sdist/amici/jax/model.py

@@ -1491,6 +1491,8 @@ def run_simulation(
             Mask for states that need reinitialisation
         :param x_reinit:
             Reinitialisation values for states
+        :param h_mask:
+            Mask for the events that are part of the current experiment
         :param solver:
             ODE solver to use for simulation
         :param controller:
@@ -1503,6 +1505,8 @@ def run_simulation(
         :param x_preeq:
             Pre-equilibration state. Can be empty if no pre-equilibration is available, in which case the states will
             be initialised to the model default values.
+        :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.
         :param t_zeros:
@@ -1563,6 +1567,9 @@ def run_simulations(
         :param preeq_array:
             Matrix of pre-equilibrated states for the simulation conditions. Ordering must match the simulation
             conditions. If no pre-equilibration is available for a condition, the corresponding row must be empty.
+        :param h_preeqs:
+            Matrix of pre-equilibration event heaviside variables indicating whether an event condition is false or
+            true after preequilibration.
         :param solver:
             ODE solver to use for simulation.
         :param controller:
@@ -1678,6 +1685,8 @@ def run_preequilibration(
             Mask for states that need reinitialisation
         :param x_reinit:
             Reinitialisation values for states
+        :param h_mask:
+            Mask for the events that are part of the current experiment
         :param solver:
             ODE solver to use for simulation
         :param controller:

python/sdist/amici/jax/petab.py

@@ -117,27 +117,6 @@ def _conditions_to_experiment_map(experiment_df: pd.DataFrame) -> dict[str, str]
     }
     return condition_to_experiment
 
-# def get_states_in_condition_table_v2(
-#     petab_problem,
-#     condition: dict | pd.Series = None,
-# ) -> dict[str, tuple[float | str | None, float | str | None]]:
-#     """Get states and their initial condition as specified in the condition table.
-
-#     Returns: Dictionary: ``stateId -> (initial condition simulation)``
-#     """
-#     states = {
-#         target_id: (target_value, None)
-#         if condition_id == condition[petabv1.SIMULATION_CONDITION_ID]
-#         else (None, None)
-#         for condition_id, target_id, target_value in zip(
-#             petab_problem.condition_df[petabv2.C.CONDITION_ID],
-#             petab_problem.condition_df[petabv2.C.TARGET_ID],
-#             petab_problem.condition_df[petabv2.C.TARGET_VALUE],
-#         )
-#     }
-
-#     return states
-
 def _try_float(value):
     try:
         return float(value)

python/sdist/amici/sim/jax/__init__.py

@@ -245,8 +245,7 @@ def run():
     n_total = 0
     version = "v2.0.0"
 
-    # for jax in (False, True):
-    for jax in (True):
+    for jax in (False, True):
         cases = list(petabtests.get_cases("sbml", version=version))
         n_total += len(cases)
         for case in cases: