implement implicit triggers using fixed parameters check

Author: BSnelling

python/sdist/amici/_symbolic/de_model.py

@@ -2684,7 +2684,9 @@ def has_implicit_event_assignments(self) -> bool:
         :return:
             boolean indicating if event assignments with implicit triggers are present
         """
-        return any(event.updates_state and event._implicit_symbols() for event in self._events)
+        fixed_symbols = set([k._symbol for k in self._fixed_parameters])
+        allowed_symbols = fixed_symbols | {amici_time_symbol}
+        return any(event.updates_state and not event.has_explicit_trigger_times(allowed_symbols) for event in self._events)
 
     def toposort_expressions(
         self, reorder: bool = True

python/sdist/amici/_symbolic/de_model_components.py

@@ -863,8 +863,12 @@ def has_explicit_trigger_times(
         """
         if allowed_symbols is None:
             return len(self._t_root) > 0
+        
+        if len(self._t_root) == 0:
+            t = self.get_val()
+            return t.is_Number or t.free_symbols.issubset(allowed_symbols)
 
-        return len(self._t_root) > 0 and all(
+        return all(
             t.is_Number or t.free_symbols.issubset(allowed_symbols)
             for t in self._t_root
         )
@@ -883,18 +887,6 @@ def get_trigger_times(self) -> set[sp.Expr]:
         time points at which the event triggers.
         """
         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:
-            if (s.startswith("_petab_") and "indicator" in s) or s == "t":
-                continue
-            implicit_symbols.append(s)
-        return len(implicit_symbols) > 0
 
     @property
     def uses_values_from_trigger_time(self) -> bool:

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

@@ -353,6 +353,7 @@ def _do_import_sbml(self):
                 model_name=self._module_name,
                 output_dir=self.output_dir,
                 observation_model=observation_model,
+                fixed_parameters=fixed_parameters,
                 verbose=self._verbose,
                 # **kwargs,
             )

python/sdist/amici/importers/sbml/__init__.py

@@ -438,6 +438,7 @@ def sbml2jax(
         self,
         model_name: str,
         output_dir: str | Path = None,
+        fixed_parameters: Iterable[str] = None,
         observation_model: list[MeasurementChannel] = None,
         verbose: int | bool = logging.ERROR,
         compute_conservation_laws: bool = True,
@@ -465,6 +466,9 @@ def sbml2jax(
         :param output_dir:
             Directory where the generated model package will be stored.
 
+        :param fixed_parameters:
+            SBML Ids to be excluded from sensitivity analysis
+
         :param observation_model:
             The different measurement channels that make up the observation
             model, see :class:`amici.importers.utils.MeasurementChannel`.
@@ -513,6 +517,7 @@ def sbml2jax(
         set_log_level(logger, verbose)
 
         ode_model = self._build_ode_model(
+            fixed_parameters=fixed_parameters,
             observation_model=observation_model,
             verbose=verbose,
             compute_conservation_laws=compute_conservation_laws,

python/sdist/amici/jax/jax.template.py

@@ -21,14 +21,14 @@ class JAXModel_TPL_MODEL_NAME(JAXModel):
     def __init__(self):
         self.jax_py_file = Path(__file__).resolve()
         self.nns = {TPL_NETS}
-        self.parameters = TPL_P_VALUES
+        self.parameters = TPL_ALL_P_VALUES
         super().__init__()
 
     def _xdot(self, t, x, args):
         p, tcl, h = args
 
         TPL_X_SYMS = x
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
         TPL_TCL_SYMS = tcl
         TPL_H_SYMS = h
         TPL_W_SYMS = self._w(t, x, p, tcl, h)
@@ -39,7 +39,7 @@ def _xdot(self, t, x, args):
 
     def _w(self, t, x, p, tcl, h):
         TPL_X_SYMS = x
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
         TPL_TCL_SYMS = tcl
         TPL_H_SYMS = h
 
@@ -48,7 +48,7 @@ def _w(self, t, x, p, tcl, h):
         return TPL_W_RET
 
     def _x0(self, t, p):
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
 
         TPL_X0_EQ
 
@@ -71,15 +71,15 @@ def _x_rdata(self, x, tcl):
 
     def _tcl(self, x, p):
         TPL_X_RDATA_SYMS = x
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
 
         TPL_TOTAL_CL_EQ
 
         return TPL_TOTAL_CL_RET
 
     def _y(self, t, x, p, tcl, h, op):
         TPL_X_SYMS = x
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
         TPL_W_SYMS = self._w(t, x, p, tcl, h)
         TPL_OP_SYMS = op
 
@@ -88,7 +88,7 @@ def _y(self, t, x, p, tcl, h, op):
         return TPL_Y_RET
 
     def _sigmay(self, y, p, np):
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
 
         TPL_Y_SYMS = y
         TPL_NP_SYMS = np
@@ -110,15 +110,15 @@ def _nllh(self, t, x, p, tcl, h, my, iy, op, np):
         return TPL_JY_RET.at[iy].get()
 
     def _known_discs(self, p):
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
 
         return TPL_ROOTS
 
     def _root_cond_fn(self, t, y, args, **_):
         p, tcl, h = args
 
         TPL_X_SYMS = y
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
         TPL_TCL_SYMS = tcl
         TPL_H_SYMS = h
         TPL_W_SYMS = self._w(t, y, p, tcl, h)
@@ -130,7 +130,7 @@ def _root_cond_fn(self, t, y, args, **_):
 
     def _delta_x(self, y, p, tcl):
         TPL_X_SYMS = y
-        TPL_P_SYMS = p
+        TPL_ALL_P_SYMS = p
         TPL_TCL_SYMS = tcl
         # FIXME: workaround until state from event time is properly passed
         TPL_X_OLD_SYMS = y
@@ -157,7 +157,7 @@ def state_ids(self):
 
     @property
     def parameter_ids(self):
-        return TPL_P_IDS
+        return TPL_ALL_P_IDS
 
     @property
     def expression_ids(self):

python/sdist/amici/jax/ode_export.py

@@ -213,6 +213,7 @@ def _generate_jax_code(self) -> None:
         )
         sym_names = (
             "p",
+            "k",
             "np",
             "op",
             "x",
@@ -261,6 +262,9 @@ def _generate_jax_code(self) -> None:
             # tuple of variable names (ids as they are unique)
             **_jax_variable_ids(self.model, ("p", "k", "y", "w", "x_rdata")),
             "P_VALUES": _jnp_array_str(self.model.val("p")),
+            "ALL_P_VALUES": _jnp_array_str(self.model.val("p") + self.model.val("k")),
+            "ALL_P_IDS": "".join(f'"{s.name}", ' for s in self._get_all_p_syms()),
+            "ALL_P_SYMS": "".join(f"{s.name}, " for s in self._get_all_p_syms()),
             "ROOTS": _jnp_array_str(
                 {
                     _print_trigger_root(root)
@@ -297,6 +301,9 @@ def _generate_jax_code(self) -> None:
             tpl_data,
         )
 
+    def _get_all_p_syms(self) -> list[sp.Symbol]:
+        return list(self.model.sym("p")) + list(self.model.sym("k"))
+
     def _generate_nn_code(self) -> None:
         for net_name, net in self.hybridization.items():
             generate_equinox(