Track discontinuities in JAX & fix gradient for models with parameter dependent roots (#2815)

* add tracking of known roots

* respect private APIs

* pin diffrax

* Update test_sbml_semantic_test_suite_jax.yml

* Update python/sdist/amici/jax/ode_export.py

Co-authored-by: Daniel Weindl <[email protected]>

---------

Co-authored-by: Daniel Weindl <[email protected]>

Author: FFroehlich

.github/workflows/test_benchmark_collection_models.yml

@@ -134,6 +134,8 @@ jobs:
         &&  python3 -m pip install -U sympy \
         &&  python3 -m pip install git+https://github.com/ICB-DCM/fiddy.git
 
+    - run: pip uninstall -y diffrax && pip install git+https://github.com/patrick-kidger/diffrax # TODO FIXME https://github.com/patrick-kidger/diffrax/issues/654
+
     - name: Download benchmark collection
       run: |
         pip install git+https://github.com/Benchmarking-Initiative/Benchmark-Models-PEtab.git@master#subdirectory=src/python

.github/workflows/test_sbml_semantic_test_suite_jax.yml

@@ -42,6 +42,7 @@ jobs:
       uses: ./.github/actions/install-apt-dependencies
 
     - run: AMICI_PARALLEL_COMPILE="" ./scripts/installAmiciSource.sh
+    - run: source ./venv/bin/activate && pip uninstall -y diffrax && pip install git+https://github.com/patrick-kidger/diffrax # TODO FIXME https://github.com/patrick-kidger/diffrax/issues/654
     - run: ./scripts/run-SBMLTestsuite.sh --jax ${{ matrix.cases }}
 
     - name: "Upload artifact: SBML semantic test suite results"

python/sdist/amici/de_model.py

@@ -2362,3 +2362,57 @@ def _process_heavisides(
                 dxdt = dxdt.subs(heaviside_sympy, heaviside_amici)
 
         return dxdt
+
+    def get_explicit_roots(self) -> set[sp.Expr]:
+        """
+        Returns explicit formulas for all discontinuities (events)
+        that can be precomputed
+
+        :return:
+            set of symbolic roots
+        """
+        return {root for e in self._events for root in e.get_trigger_times()}
+
+    def get_implicit_roots(self) -> set[sp.Expr]:
+        """
+        Returns implicit equations for all discontinuities (events)
+        that have to be located via rootfinding
+
+        :return:
+            set of symbolic roots
+        """
+        return {
+            e.get_val()
+            for e in self._events
+            if not e.has_explicit_trigger_times()
+        }
+
+    def has_algebraic_states(self) -> bool:
+        """
+        Checks whether the model has algebraic states
+
+        :return:
+            boolean indicating if algebraic states are present
+        """
+        return len(self._algebraic_states) > 0
+
+    def has_event_assignments(self) -> bool:
+        """
+        Checks whether the model has event assignments
+
+        :return:
+            boolean indicating if event assignments are present
+        """
+        return any(event.updates_state for event in self._events)
+
+    def has_parameter_dependent_implicit_roots(self) -> bool:
+        """
+        Checks whether the model has events with parameter-dependent implicit roots
+
+        :return:
+            boolean indicating if parameter-dependent implicit roots are present
+        """
+        return any(
+            self.sym("p").has(root.free_symbols)
+            for root in self.get_implicit_roots()
+        )

python/sdist/amici/de_model_components.py

@@ -829,6 +829,22 @@ def get_trigger_time(self) -> sp.Float:
             )
         return self._t_root[0]
 
+    def has_explicit_trigger_times(self) -> bool:
+        """Check whether the event has explicit trigger times.
+
+        Explicit trigger times do not require root finding to determine
+        the time points at which the event triggers.
+        """
+        return len(self._t_root) > 0
+
+    def get_trigger_times(self) -> set[sp.Expr]:
+        """Get the time points at which the event triggers.
+
+        Returns a set of expressions, which may contain multiple time points
+        for events that trigger at multiple time points.
+        """
+        return set(self._t_root)
+
     @property
     def uses_values_from_trigger_time(self) -> bool:
         """Whether the event assignment is evaluated using the state from

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

@@ -1,5 +1,6 @@
 # ruff: noqa: F401, F821, F841
 import jax.numpy as jnp
+import jaxtyping as jt
 from interpax import interp1d
 from pathlib import Path
 from jax.numpy import inf as oo
@@ -99,6 +100,11 @@ def _nllh(self, t, x, p, tcl, my, iy, op, np):
 
         return TPL_JY_RET.at[iy].get()
 
+    def _known_discs(self, p):
+        TPL_P_SYMS = p
+
+        return TPL_ROOTS
+
     @property
     def observable_ids(self):
         return TPL_Y_IDS

python/sdist/amici/jax/model.py

@@ -238,6 +238,20 @@ def _nllh(
         """
         ...
 
+    @abstractmethod
+    def _known_discs(
+        self, p: jt.Float[jt.Array, "np"]
+    ) -> jt.Float[jt.Array, "ndiscs"]:
+        """
+        Compute the known discontinuity points of the ODE system.
+
+        :param p:
+            parameters
+        :return:
+            known discontinuity points in the ODE system
+        """
+        ...
+
     @property
     @abstractmethod
     def state_ids(self) -> list[str]:
@@ -319,7 +333,9 @@ def _eq(
             t1=jnp.inf,
             dt0=None,
             y0=x0,
-            stepsize_controller=controller,
+            stepsize_controller=self._get_clipped_stepsize_controller(
+                p, controller
+            ),
             max_steps=max_steps,
             adjoint=diffrax.DirectAdjoint(),
             event=diffrax.Event(
@@ -378,7 +394,9 @@ def _solve(
             t1=ts[-1],
             dt0=None,
             y0=x0,
-            stepsize_controller=controller,
+            stepsize_controller=self._get_clipped_stepsize_controller(
+                p, controller
+            ),
             max_steps=max_steps,
             adjoint=adjoint,
             saveat=diffrax.SaveAt(ts=ts),
@@ -747,6 +765,19 @@ def preequilibrate_condition(
 
         return self._x_rdata(current_x, tcl), dict(stats_preeq=stats_preeq)
 
+    def _get_clipped_stepsize_controller(
+        self,
+        p: jt.Float[jt.Array, "np"],
+        controller: diffrax.AbstractStepSizeController,
+    ) -> diffrax.AbstractStepSizeController:
+        if not self._known_discs(p).size:
+            return controller
+
+        return diffrax.ClipStepSizeController(
+            controller,
+            jump_ts=self._known_discs(p),
+        )
+
 
 def safe_log(x: jnp.float_) -> jnp.float_:
     """

python/sdist/amici/jax/ode_export.py

@@ -25,6 +25,7 @@
 from amici.jax.jaxcodeprinter import AmiciJaxCodePrinter, _jnp_array_str
 from amici.jax.model import JAXModel
 from amici.de_model import DEModel
+
 from amici.de_export import is_valid_identifier
 from amici.import_utils import (
     strip_pysb,
@@ -142,14 +143,18 @@ def __init__(
         """
         set_log_level(logger, verbose)
 
-        if any(event.updates_state for event in ode_model._events):
+        if ode_model.has_event_assignments():
             raise NotImplementedError(
                 "The JAX backend does not support models with event assignments."
             )
 
-        if ode_model._algebraic_equations:
+        if ode_model.has_algebraic_states():
+            raise NotImplementedError(
+                "The JAX backend does not support models with algebraic states."
+            )
+        if ode_model.has_parameter_dependent_implicit_roots():
             raise NotImplementedError(
-                "The JAX backend does not support models with algebraic equations."
+                "The JAX backend does not support models with parameter dependent implicit event triggers."
             )
 
         self.verbose: bool = logger.getEffectiveLevel() <= logging.DEBUG
@@ -243,6 +248,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")),
+            "ROOTS": _jnp_array_str(
+                {root for e in self.model._events for root in e.get_trigger_times()}
+            ),
             **{
                 "MODEL_NAME": self.model_name,
                 # keep track of the API version that the model was generated with so we

tests/benchmark_models/test_petab_benchmark_jax.py

@@ -57,7 +57,7 @@ def test_jax_llh(benchmark_problem):
     if problem_id in problems_for_gradient_check:
         point = flat_petab_problem.x_nominal_free_scaled
         for _ in range(20):
-            amici_solver.setSensitivityMethod(amici.SensitivityMethod.adjoint)
+            amici_solver.setSensitivityMethod(amici.SensitivityMethod.forward)
             amici_solver.setSensitivityOrder(amici.SensitivityOrder.first)
             amici_model.setSteadyStateSensitivityMode(
                 cur_settings.ss_sensitivity_mode