Fix wrong branch for piecewise nested in another piecewise's condition

CHANGELOG.md

@@ -31,6 +31,13 @@ See also our [versioning policy](https://amici.readthedocs.io/en/latest/versioni
 
 **Fixes**
 
+* Fixed both simulation backends selecting the wrong branch of a piecewise
+  expression whose condition depends on the value of another piecewise
+  expression. The Heaviside variables were initialized in a single pass, so
+  the outer condition was evaluated using the not-yet-updated value of the
+  inner expression. They are now iterated to a fixed point, which also makes
+  the result independent of the (previously differing) value the Heaviside
+  variables happened to be seeded with in either backend (#3233).
 * Demote the module import mtime check during sundials model import to a
   warning, as it can lead to false positives in some environments
   (e.g., when using network file systems with not synchronized clocks).

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

@@ -990,7 +990,35 @@ def _handle_t0_event(
             rf0 = jnp.where(h > 0.5, 0.5, -0.5)
         else:
             h = jnp.where(h_mask, jnp.heaviside(rf0, 0.0), jnp.ones_like(rf0))
-        args = (p, tcl, h)
+
+        # The root functions may themselves depend on the Heaviside variables,
+        # e.g., for a piecewise expression occurring inside the condition of
+        # another piecewise expression. Evaluating them once at the seeded `h`
+        # would use a stale value for the inner expression and select the wrong
+        # branch of the outer one, so refine `h` to a fixed point before
+        # deciding which events triggered (see #3233).
+        # `n_events` refinement steps plus the final evaluation below resolve
+        # any acyclic dependency hierarchy. The trip count is static so that
+        # this stays reverse-mode differentiable, unlike `jax.lax.while_loop`.
+        # Note that `h` itself keeps its seeded (pre-event) value -- it is the
+        # reference that `roots_found` flips below -- only the point at which
+        # the root functions are evaluated changes.
+        h_eval = h
+        if self.n_events:
+            ones = jnp.ones_like(rf0)
+
+            def refine_h(h_cur, _):
+                rfx_cur = root_cond_fn(t0_next, y0_next, (p, tcl, h_cur))
+                h_new = jnp.where(
+                    h_mask, jnp.heaviside(rfx_cur, 1.0), ones
+                ).astype(h_cur.dtype)
+                return h_new, None
+
+            h_eval, _ = jax.lax.scan(
+                refine_h, h_eval, None, length=self.n_events
+            )
+
+        args = (p, tcl, h_eval)
         rfx = root_cond_fn(t0_next, y0_next, args)
         roots_dir = jnp.sign(rfx - rf0)
         roots_found = jnp.sign(rfx) != jnp.sign(rf0)

python/tests/test_heavisides.py

@@ -2,7 +2,12 @@
 
 import numpy as np
 import pytest
+from amici import import_model_module
+from amici.importers.antimony import antimony2amici
+from amici.sim.sundials import AMICI_SUCCESS, run_simulation
+from amici.testing import skip_on_valgrind
 from amici.testing.models import create_amici_model, create_sbml_model
+from numpy.testing import assert_allclose
 from util import (
     check_trajectories_with_adjoint_sensitivities,
     check_trajectories_with_forward_sensitivities,
@@ -355,3 +360,67 @@ def sx_expected(t, x_1_0):
         x_expected,
         sx_expected,
     )
+
+
+@skip_on_valgrind
+def test_nested_piecewise_in_condition(tempdir):
+    """Test a piecewise expression inside the condition of another one.
+
+    The condition of the outer piecewise expression depends (via `p3`) on the
+    value of the inner one. This requires the Heaviside variable of the outer
+    expression to be initialized based on the *updated* value of the inner one.
+
+    See https://github.com/AMICI-dev/AMICI/issues/3233.
+    """
+    ant_model = """
+    model nested_piecewise
+      p1 := 0.001;
+      p2 := piecewise(0, p1 < 0, p1);
+      p3 := p2 - 0.0001;
+      p4 := piecewise(0, p3 < 0, p3);
+      # three levels of nesting
+      p5 := p4 - 0.0001;
+      p6 := piecewise(0, p5 < 0, p5);
+      x = 0;
+      x' = p6;
+    end
+    """
+    module_name = "test_nested_piecewise_in_condition"
+    antimony2amici(
+        ant_model,
+        model_name=module_name,
+        output_dir=tempdir,
+    )
+    amici_model = import_model_module(
+        module_name=module_name, module_path=tempdir
+    ).get_model()
+    amici_model.set_timepoints([0.0, 1.0])
+    rdata = run_simulation(amici_model, amici_model.create_solver())
+    assert rdata.status == AMICI_SUCCESS
+
+    expr_ids = list(amici_model.get_expression_ids())
+    w = np.asarray(rdata.w).reshape(len(rdata.ts), -1)
+    expected = {
+        "p1": 0.001,
+        "p2": 0.001,
+        "p3": 0.0009,
+        "p4": 0.0009,
+        "p5": 0.0008,
+        "p6": 0.0008,
+    }
+    for expr_id, expected_val in expected.items():
+        assert_allclose(
+            w[:, expr_ids.index(expr_id)],
+            expected_val,
+            rtol=1e-10,
+            atol=1e-14,
+            err_msg=f"Unexpected value for {expr_id}",
+        )
+
+    # dx/dt == p6 == const
+    assert_allclose(
+        np.asarray(rdata.x).flatten(),
+        expected["p6"] * np.asarray(rdata.ts),
+        rtol=1e-8,
+        atol=1e-14,
+    )

python/tests/test_jax.py

@@ -1175,3 +1175,72 @@ def test_resolve_net_id():
     assert JAXProblem._resolve_net_id("", nns) is None
     assert JAXProblem._resolve_net_id(float("nan"), nns) is None
     assert JAXProblem._resolve_net_id(None, nns) is None
+
+
+@skip_on_valgrind
+def test_nested_piecewise_in_condition(tmp_path):
+    """A piecewise expression inside the condition of another one.
+
+    The condition of the outer piecewise expression depends (via ``p3``) on the
+    value of the inner one, so the Heaviside variables have to be consistent
+    with the root functions before the branches are selected.
+
+    ``p2`` encodes as ``p1 * H(p1)`` with a consistent ``H(p1) = 0``. Seeding
+    the Heaviside variables from the trigger initial values (``h = 1``) would
+    give a stale ``p2 = -1e-3``, hence ``p3 = -9e-4 < 0``, and collapse ``p4``
+    to zero.
+
+    See https://github.com/AMICI-dev/AMICI/issues/3233.
+    """
+    from amici.importers.antimony import antimony2sbml
+    from amici.importers.sbml import SbmlImporter
+    from amici.sim.jax.petab import (
+        DEFAULT_CONTROLLER_SETTINGS,
+        DEFAULT_ROOT_FINDER_SETTINGS,
+        SteadyStateEvent,
+    )
+
+    ant_model = """
+    model nested_piecewise
+        p1 := -0.001;
+        p2 := piecewise(0, p1 < 0, p1);
+        p3 := p2 + 0.0001;
+        p4 := piecewise(0, p3 < 0, p3);
+        x = 0.0;
+        x' = p4;
+    end
+    """
+    sbml = antimony2sbml(ant_model)
+    SbmlImporter(sbml, from_file=False).sbml2jax(
+        "nested_piecewise", output_dir=tmp_path
+    )
+    model = amici._module_from_path(
+        "nested_piecewise", tmp_path / "__init__.py"
+    ).Model()
+
+    ts = jnp.array([0.0, 1.0])
+    x, _ = model.simulate_condition(
+        jnp.array(model.parameters),
+        ts,
+        jnp.array([]),
+        jnp.zeros_like(ts),
+        jnp.zeros_like(ts, dtype=int),
+        jnp.zeros_like(ts, dtype=int),
+        jnp.zeros((ts.shape[0], 0)),
+        jnp.zeros((ts.shape[0], 0)),
+        diffrax.Kvaerno5(),
+        diffrax.PIDController(**DEFAULT_CONTROLLER_SETTINGS),
+        optimistix.Newton(**DEFAULT_ROOT_FINDER_SETTINGS),
+        diffrax.DirectAdjoint(),
+        SteadyStateEvent(),
+        1000,
+        ret=ReturnValue.x,
+    )
+
+    # dx/dt == p4 == 1e-4, constant
+    assert_allclose(
+        np.asarray(x).flatten(),
+        1e-4 * np.asarray(ts),
+        rtol=1e-5,
+        atol=1e-12,
+    )

src/model.cpp

@@ -422,18 +422,36 @@ void Model::reinit_events(
     realtype t, AmiVector const& x, AmiVector const& dx,
     std::vector<realtype> const& h_old, std::vector<int>& roots_found
 ) {
+    // Root functions may themselves depend on Heaviside variables, e.g., for
+    // a piecewise expression occurring inside the condition of another
+    // piecewise expression. In that case, a single pass would evaluate the
+    // outer root function using stale Heaviside variables. Therefore, iterate
+    // until `h` is consistent with the root function values.
+    // Each iteration resolves at least one more level of the dependency
+    // hierarchy, so `ne` iterations always suffice for acyclic dependencies.
+    // The iteration count is bounded to avoid infinite loops in case of
+    // (pathological) cyclic dependencies.
     std::vector<realtype> rootvals(ne, 0.0);
-    froot(t, x, dx, rootvals);
+    for (int iter = 0; iter <= ne; ++iter) {
+        froot(t, x, dx, rootvals);
+        bool h_changed = false;
+        for (int ie = 0; ie < ne; ie++) {
+            realtype const h_new = rootvals.at(ie) < 0.0 ? 0.0 : 1.0;
+            if (h_new != state_.h.at(ie)) {
+                state_.h.at(ie) = h_new;
+                h_changed = true;
+            }
+        }
+        if (!h_changed) {
+            break;
+        }
+    }
+
     std::ranges::fill(roots_found, 0);
     for (int ie = 0; ie < ne; ie++) {
-        if (rootvals.at(ie) < 0.0) {
-            state_.h.at(ie) = 0.0;
-        } else {
-            state_.h.at(ie) = 1.0;
-            if (h_old.at(ie) <= 0.0) {
-                // only false->true triggers event
-                roots_found.at(ie) = 1;
-            }
+        if (state_.h.at(ie) > 0.0 && h_old.at(ie) <= 0.0) {
+            // only false->true triggers event
+            roots_found.at(ie) = 1;
         }
     }