diff --git a/CHANGELOG.md b/CHANGELOG.md index 77a028f8b7..03afe7dfd2 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/python/sdist/amici/sim/jax/model.py b/python/sdist/amici/sim/jax/model.py index 7ba1ec76ee..39f5540873 100644 --- a/python/sdist/amici/sim/jax/model.py +++ b/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) diff --git a/python/tests/test_heavisides.py b/python/tests/test_heavisides.py index a6e3b1953d..11e9127df2 100644 --- a/python/tests/test_heavisides.py +++ b/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, + ) diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 07c7406fa6..fd8f431383 100644 --- a/python/tests/test_jax.py +++ b/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, + ) diff --git a/src/model.cpp b/src/model.cpp index 375556fec8..ce8fd0d025 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -422,18 +422,36 @@ void Model::reinit_events( realtype t, AmiVector const& x, AmiVector const& dx, std::vector const& h_old, std::vector& 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 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; } }