Let froot depend on w

WIP

This should hopefully allow for more efficient computation of `root`-derivatives
by avoiding flattening `w` into `root` which, so far, made computing `drootdt_total`
prohibitively expensive in case of large `w` dependencies in `root`.

Author: dweindl

cmake/AmiciFindBLAS.cmake

@@ -15,6 +15,11 @@ if(DEFINED ENV{AMICI_BLAS_USE_SCIPY_OPENBLAS})
       "Using AMICI_BLAS_USE_SCIPY_OPENBLAS=${AMICI_BLAS_USE_SCIPY_OPENBLAS} from environment variable."
   )
   set(AMICI_BLAS_USE_SCIPY_OPENBLAS $ENV{AMICI_BLAS_USE_SCIPY_OPENBLAS})
+elseif(NOT DEFINED AMICI_BLAS_USE_SCIPY_OPENBLAS
+       AND NOT AMICI_PYTHON_BUILD_EXT_ONLY)
+  # If were are not building the Python extension, it's unlikely that we want to
+  # use scipy-openblas
+  set(AMICI_BLAS_USE_SCIPY_OPENBLAS FALSE)
 endif()
 
 if((${BLAS} STREQUAL "MKL" OR DEFINED ENV{MKLROOT})

include/amici/abstract_model.h

@@ -299,6 +299,7 @@ class AbstractModel {
      * @param p parameter vector
      * @param k constant vector
      * @param h Heaviside vector
+     * @param w vector with helper variables
      * @param dx time derivative of state (DAE only)
      * @param tcl total abundances for conservation laws
      * @param sx current state sensitivity
@@ -307,8 +308,9 @@ class AbstractModel {
      */
     virtual void fstau(
         realtype* stau, realtype t, realtype const* x, realtype const* p,
-        realtype const* k, realtype const* h, realtype const* dx,
-        realtype const* tcl, realtype const* sx, int ip, int ie
+        realtype const* k, realtype const* h, realtype const* w,
+        realtype const* dx, realtype const* tcl, realtype const* sx, int ip,
+        int ie
     );
 
     /**
@@ -542,6 +544,7 @@ class AbstractModel {
      * @param p parameter vector
      * @param k constant vector
      * @param h Heaviside vector
+     * @param w vector with helper variables
      * @param dx time derivative of state (DAE only)
      * @param ie event index
      * @param xdot new model right hand side
@@ -552,9 +555,10 @@ class AbstractModel {
      */
     virtual void fdeltaxB(
         realtype* deltaxB, realtype t, realtype const* x, realtype const* p,
-        realtype const* k, realtype const* h, realtype const* dx, int ie,
-        realtype const* xdot, realtype const* xdot_old, realtype const* x_old,
-        realtype const* xB, realtype const* tcl
+        realtype const* k, realtype const* h, realtype const* w,
+        realtype const* dx, int ie, realtype const* xdot,
+        realtype const* xdot_old, realtype const* x_old, realtype const* xB,
+        realtype const* tcl
     );
 
     /**
@@ -565,6 +569,7 @@ class AbstractModel {
      * @param p parameter vector
      * @param k constant vector
      * @param h Heaviside vector
+     * @param w vector with helper variables
      * @param dx time derivative of state (DAE only)
      * @param ip sensitivity index
      * @param ie event index
@@ -575,9 +580,9 @@ class AbstractModel {
      */
     virtual void fdeltaqB(
         realtype* deltaqB, realtype t, realtype const* x, realtype const* p,
-        realtype const* k, realtype const* h, realtype const* dx, int ip,
-        int ie, realtype const* xdot, realtype const* xdot_old,
-        realtype const* x_old, realtype const* xB
+        realtype const* k, realtype const* h, realtype const* w,
+        realtype const* dx, int ip, int ie, realtype const* xdot,
+        realtype const* xdot_old, realtype const* x_old, realtype const* xB
     );
 
     /**

include/amici/model_dae.h

@@ -359,11 +359,13 @@ class Model_DAE : public Model {
      * @param p parameter vector
      * @param k constants vector
      * @param h Heaviside vector
+     * @param w vector with helper variables
      * @param dx Vector with the derivative states
      **/
     virtual void froot(
         realtype* root, realtype t, realtype const* x, double const* p,
-        double const* k, realtype const* h, realtype const* dx
+        double const* k, realtype const* h, realtype const* w,
+        realtype const* dx
     );
 
     /**

include/amici/model_ode.h

@@ -316,11 +316,13 @@ class Model_ODE : public Model {
      * @param p parameter vector
      * @param k constants vector
      * @param h Heaviside vector
+     * @param w vector with helper variables
      * @param tcl total abundances for conservation laws
      **/
     virtual void froot(
         realtype* root, realtype t, realtype const* x, realtype const* p,
-        realtype const* k, realtype const* h, realtype const* tcl
+        realtype const* k, realtype const* h, realtype const* w,
+        realtype const* tcl
     );
 
     /**

python/sdist/amici/de_model.py

@@ -296,6 +296,8 @@ def __init__(
         ] = {
             "sroot": {
                 "eq": "root",
+                # TODO?
+                # "chainvars": ["x", "w"],
                 "chainvars": ["x"],
                 "var": "p",
                 "dxdz_name": "sx",
@@ -528,17 +530,9 @@ def get_rate(symbol: sp.Symbol):
 
         for component in chain(
             self.observables(),
-            self.events(),
             self._algebraic_equations,
         ):
             if rate_ofs := component.get_val().find(rate_of_func):
-                if isinstance(component, Event):
-                    # TODO froot(...) can currently not depend on `w`, so this substitution fails for non-zero rates
-                    #  see, e.g., sbml test case 01293
-                    raise SBMLException(
-                        "AMICI does currently not support rateOf(.) inside event trigger functions."
-                    )
-
                 if isinstance(component, AlgebraicEquation):
                     # TODO IDACalcIC fails with
                     #   "The linesearch algorithm failed: step too small or too many backtracks."
@@ -1340,10 +1334,10 @@ def parse_events(self) -> None:
             # add roots of heaviside functions
             self.add_component(root)
 
-        # Substitute 'w' expressions into root expressions, to avoid rewriting
-        # 'root.cpp' and 'stau.cpp' headers to include 'w.h'.
-        for event in self.events():
-            event.set_val(event.get_val().subs(w_toposorted))
+        # # Substitute 'w' expressions into root expressions, to avoid rewriting
+        # # 'root.cpp' and 'stau.cpp' headers to include 'w.h'.
+        # for event in self.events():
+        #     event.set_val(event.get_val().subs(w_toposorted))
 
         # re-order events - first those that require root tracking, then the others
         constant_syms = set(self.sym("k")) | set(self.sym("p"))
@@ -1668,18 +1662,23 @@ def _compute_equation(self, name: str) -> None:
             self._eqs[name] = smart_jacobian(self.eq("root"), time_symbol)
 
         elif name == "drootdt_total":
+            # root(t, x(t), w(t, x(t)))
+            # drootdt_total = drootdt + drootdx * dxdt + drootdw * dwdt_total
+            # dwdt_total = dwdt + dwdx * dxdt
             self._eqs[name] = self.eq("drootdt")
-            # backsubstitution of optimized right-hand side terms into RHS
-            # calling subs() is costly. We can skip it if we don't have any
-            # state-dependent roots.
+
+            xdot = self.eq("xdot")
             if self.num_states_solver() and not smart_is_zero_matrix(
                 drootdx := self.eq("drootdx")
             ):
-                w_sorted = toposort_symbols(
-                    dict(zip(self.sym("w"), self.eq("w"), strict=True))
-                )
-                tmp_xdot = smart_subs_dict(self.eq("xdot"), w_sorted)
-                self._eqs[name] += smart_multiply(drootdx, tmp_xdot)
+                self._eqs[name] += smart_multiply(drootdx, xdot)
+
+            drootdw = self.eq("drootdw")
+            dwdt = self.eq("dwdt")
+            dwdx = self.eq("dwdx")
+            dwdt_total = dwdt + smart_multiply(dwdx, xdot)
+
+            self._eqs[name] += smart_multiply(drootdw, dwdt_total)
 
         elif name == "deltax":
             # fill boluses for Heaviside functions, as empty state updates
@@ -1763,16 +1762,26 @@ def _compute_equation(self, name: str) -> None:
             ]
 
         elif name == "dtaudx":
+            # TODO drootdx + drootdw * dwdx
             self._eqs[name] = [
-                self.eq("drootdx")[ie, :] / self.eq("drootdt_total")[ie]
+                (
+                    self.eq("drootdx")[ie, :]
+                    + self.eq("drootdw")[ie, :] * self.eq("dwdx")
+                )
+                / self.eq("drootdt_total")[ie]
                 if not self.eq("drootdt_total")[ie].is_zero
                 else sp.zeros(*self.eq("drootdx")[ie, :].shape)
                 for ie in range(self.num_events())
             ]
 
         elif name == "dtaudp":
+            # TODO drootdp + drootdw * dwdp
             self._eqs[name] = [
-                self.eq("drootdp")[ie, :] / self.eq("drootdt_total")[ie]
+                (
+                    self.eq("drootdp")[ie, :]
+                    + self.eq("drootdw")[ie, :] * self.eq("dwdp")
+                )
+                / self.eq("drootdt_total")[ie]
                 if not self.eq("drootdt_total")[ie].is_zero
                 else sp.zeros(*self.eq("drootdp")[ie, :].shape)
                 for ie in range(self.num_events())
@@ -1922,6 +1931,9 @@ def _compute_equation(self, name: str) -> None:
                 smart_jacobian(self.eq("w")[self.num_cons_law() :, :], x)
             )
 
+        elif name == "dwdt":
+            self._eqs[name] = smart_jacobian(self.eq("w"), time_symbol)
+
         elif name == "iroot":
             self._eqs[name] = sp.Matrix(
                 [
@@ -2110,7 +2122,9 @@ def _derivative(self, eq: str, var: str, name: str = None) -> None:
                             "attach this model."
                         )
 
-        if name == "dydw" and not smart_is_zero_matrix(derivative):
+        elif name in ("dydw", "drootdw") and not smart_is_zero_matrix(
+            derivative
+        ):
             dwdw = self.eq("dwdw")
             # h(k) = d{eq}dw*dwdw^k* (k=1)
             h = smart_multiply(derivative, dwdw)
@@ -2407,6 +2421,8 @@ def _expr_is_time_dependent(self, expr: sp.Expr) -> bool:
         :returns:
             Whether the expression is time-dependent.
         """
+        # TODO: handle w-dependency
+
         # `expr.free_symbols` will be different to `self._states.keys()`, so
         # it's easier to compare as `str`.
         expr_syms = {str(sym) for sym in expr.free_symbols}
@@ -2515,13 +2531,9 @@ def _process_heavisides(
         tmp_roots_old = self._collect_heaviside_roots((dxdt,))
         # substitute 'w' symbols in the root expression by their equations,
         #  because currently,
-        #  1) root functions must not depend on 'w'
-        #  2) the check for time-dependence currently assumes only state
+        # #  1) root functions must not depend on 'w'
+        # FIXME 2) the check for time-dependence currently assumes only state
         #     variables are implicitly time-dependent
-        tmp_roots_old = [
-            (a.subs(w_toposorted), b.subs(w_toposorted))
-            for a, b in tmp_roots_old
-        ]
         for tmp_root_old, tmp_x0_old in unique_preserve_order(tmp_roots_old):
             # we want unique identifiers for the roots
             tmp_root_new = self._get_unique_root(tmp_root_old, roots)

python/sdist/amici/exporters/sundials/cxx_functions.py

@@ -144,7 +144,7 @@ def var_in_signature(self, varname: str, ode: bool = True) -> bool:
     "root": _FunctionInfo(
         "realtype *root, const realtype t, const realtype *x, "
         "const realtype *p, const realtype *k, const realtype *h, "
-        "const realtype *tcl"
+        "const realtype *w, const realtype *tcl"
     ),
     "dwdp": _FunctionInfo(
         "realtype *dwdp, const realtype t, const realtype *x, "
@@ -288,7 +288,7 @@ def var_in_signature(self, varname: str, ode: bool = True) -> bool:
     "stau": _FunctionInfo(
         "realtype *stau, const realtype t, const realtype *x, "
         "const realtype *p, const realtype *k, const realtype *h, "
-        "const realtype *dx, "
+        "const realtype *w, const realtype *dx, "
         "const realtype *tcl, const realtype *sx, const int ip, "
         "const int ie"
     ),
@@ -313,15 +313,15 @@ def var_in_signature(self, varname: str, ode: bool = True) -> bool:
     "deltaxB": _FunctionInfo(
         "realtype *deltaxB, const realtype t, const realtype *x, "
         "const realtype *p, const realtype *k, const realtype *h, "
-        "const realtype *dx, "
+        "const realtype *w, const realtype *dx, "
         "const int ie, const realtype *xdot, const realtype *xdot_old, "
         "const realtype *x_old, "
         "const realtype *xB, const realtype *tcl"
     ),
     "deltaqB": _FunctionInfo(
         "realtype *deltaqB, const realtype t, const realtype *x, "
         "const realtype *p, const realtype *k, const realtype *h, "
-        "const realtype *dx, "
+        "const realtype *w, const realtype *dx, "
         "const int ip, const int ie, const realtype *xdot, "
         "const realtype *xdot_old, const realtype *x_old, const realtype *xB"
     ),

python/sdist/amici/gradient_check.py

@@ -133,17 +133,22 @@ def check_finite_difference(
         else:
             raise NotImplementedError()
 
-        _check_close(
-            sensi,
-            fd,
-            atol=atol,
-            rtol=rtol,
-            field=field,
-            ip=ip,
-            parameter_id=model.get_parameter_ids()[ip]
-            if model.has_parameter_ids()
-            else None,
-        )
+        try:
+            _check_close(
+                sensi,
+                fd,
+                atol=atol,
+                rtol=rtol,
+                field=field,
+                ip=ip,
+                parameter_id=model.get_parameter_ids()[ip]
+                if model.has_parameter_ids()
+                else None,
+            )
+        except AssertionError as e:
+            sm = SensitivityMethod(solver.get_sensitivity_method())
+            e.add_note(f"Sensitivity method was {sm!r}")
+            raise e
 
     solver.set_sensitivity_order(og_sensitivity_order)
     model.set_parameters(og_parameters)

python/tests/test_events.py

@@ -1294,3 +1294,58 @@ def test_gh2926(tempdir):
     rdata = amici.run_simulation(model, solver)
     assert rdata.status == amici.AMICI_SUCCESS
     assert rdata.by_id("x1").tolist() == [1.0, 1.0, 2.0, 2.0]
+
+
+@skip_on_valgrind
+def test_event_with_w_dependent_trigger(tempdir):
+    """Test sensitivities for events with trigger depending on
+    cascading expressions in `w`."""
+
+    model_name = "test_event_with_w_dependent_trigger"
+    model = antimony2amici(
+        r"""
+        one = 1
+        two = 2
+        a := two^2 - 2 # 2
+        b := a^2 - 1  # 3
+        c := b * 2 + x / 10 # 6
+        d := c + a + (a - 1) * time / 10 # -> d = 8 + time / 5
+
+        x = 0
+        x' = 1
+        target = 0
+
+        E1: at x >= d:  # triggers at time = 8 + time / 5 <=> time = 10
+            target = x + one;
+        """,
+        model_name=model_name,
+        output_dir=tempdir,
+    )
+
+    model.set_timepoints([0, 5, 9, 11])
+    solver = model.create_solver()
+    solver.set_sensitivity_order(SensitivityOrder.first)
+    solver.set_sensitivity_method(SensitivityMethod.forward)
+
+    # generate synthetic measurements
+    rdata = amici.run_simulation(model, solver)
+    assert rdata.status == amici.AMICI_SUCCESS
+    # check that event triggered correctly
+    assert np.isclose(rdata.by_id("target")[-1], 11.0)
+    edata = amici.ExpData(rdata, 1, 0)
+
+    # check sensitivities against finite differences
+
+    for sens_method in (
+        SensitivityMethod.forward,
+        SensitivityMethod.adjoint,
+    ):
+        solver.set_sensitivity_method(sens_method)
+        check_derivatives(
+            model,
+            solver,
+            edata=edata,
+            atol=1e-6,
+            rtol=1e-6,
+            epsilon=1e-8,
+        )