From a35ce7dab1e84329793d0690780a1b0d037f074a Mon Sep 17 00:00:00 2001
From: Daniel Weindl <daniel.weindl@uni-bonn.de>
Date: Thu, 23 Oct 2025 22:04:37 +0200
Subject: [PATCH] PySB events

Add option to sneak general events into PySB-derived AMICI models.

At least for now, this is intended only for internal use,
in particular for PEtab v2 handling.
---
 python/sdist/amici/de_model_components.py |  2 +-
 python/sdist/amici/pysb_import.py         | 14 ++++++-
 python/tests/test_pysb.py                 | 46 +++++++++++++++++++++++
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/python/sdist/amici/de_model_components.py b/python/sdist/amici/de_model_components.py
index bf11cca3d0..0922772bc3 100644
--- a/python/sdist/amici/de_model_components.py
+++ b/python/sdist/amici/de_model_components.py
@@ -783,7 +783,7 @@ def get_state_update(
         if len(self._assignments) == 0:
             return None
 
-        x_to_x_old = dict(zip(x, x_old))
+        x_to_x_old = dict(zip(x, x_old, strict=True))
 
         def get_bolus(x_i: sp.Symbol) -> sp.Expr:
             """
diff --git a/python/sdist/amici/pysb_import.py b/python/sdist/amici/pysb_import.py
index 727e5c59c6..9f55cfc60a 100644
--- a/python/sdist/amici/pysb_import.py
+++ b/python/sdist/amici/pysb_import.py
@@ -35,7 +35,7 @@
     SigmaY,
 )
 from .de_model import DEModel
-from .de_model_components import NoiseParameter, ObservableParameter
+from .de_model_components import Event, NoiseParameter, ObservableParameter
 from .import_utils import (
     MeasurementChannel,
     _default_simplify,
@@ -163,6 +163,7 @@ def pysb2amici(
     generate_sensitivity_code: bool = True,
     model_name: str | None = None,
     pysb_model_has_obs_and_noise: bool = False,
+    _events: list[Event] = None,
 ) -> amici.Model | None:
     r"""
     Generate AMICI C++ files for the provided model.
@@ -262,6 +263,7 @@ def pysb2amici(
         cache_simplify=cache_simplify,
         verbose=verbose,
         pysb_model_has_obs_and_noise=pysb_model_has_obs_and_noise,
+        events=_events,
     )
     exporter = DEExporter(
         ode_model,
@@ -303,6 +305,7 @@ def ode_model_from_pysb_importer(
     verbose: int | bool = False,
     jax: bool = False,
     pysb_model_has_obs_and_noise: bool = False,
+    events: list[Event] = None,
 ) -> DEModel:
     """
     Creates an :class:`amici.DEModel` instance from a :class:`pysb.Model`
@@ -360,6 +363,11 @@ def ode_model_from_pysb_importer(
     _process_pysb_species(model, ode)
     _process_pysb_parameters(model, ode, constant_parameters, jax)
     if compute_conservation_laws:
+        if events:
+            raise NotImplementedError(
+                "Conservation law computation is not supported for models "
+                "with events. Use `compute_conservation_laws=False`."
+            )
         _process_pysb_conservation_laws(model, ode)
     _process_pysb_observables(
         model,
@@ -373,6 +381,10 @@ def ode_model_from_pysb_importer(
         observation_model,
         pysb_model_has_obs_and_noise,
     )
+
+    for event in events or []:
+        ode.add_component(event)
+
     ode._has_quadratic_nllh = all(
         channel.noise_distribution
         in ["normal", "lin-normal", "log-normal", "log10-normal"]
diff --git a/python/tests/test_pysb.py b/python/tests/test_pysb.py
index ab8bec9ed8..4803494550 100644
--- a/python/tests/test_pysb.py
+++ b/python/tests/test_pysb.py
@@ -17,7 +17,9 @@
 import pytest
 import sympy as sp
 from amici import ParameterScaling, parameter_scaling_from_int_vector
+from amici.de_model_components import Event
 from amici.gradient_check import check_derivatives
+from amici.import_utils import amici_time_symbol
 from amici.pysb_import import pysb2amici
 from amici.testing import TemporaryDirectoryWinSafe, skip_on_valgrind
 from numpy.testing import assert_allclose
@@ -391,3 +393,47 @@ def test_energy():
         check_derivatives(
             amici_model, solver, epsilon=1e-4, rtol=1e-2, atol=1e-2
         )
+
+
+@skip_on_valgrind
+def test_pysb_event(tempdir):
+    """Test adding events to PySB models."""
+    pysb.SelfExporter.cleanup()  # reset pysb
+    pysb.SelfExporter.do_export = True
+
+    model = pysb.Model("pysb_event_test")
+    a = pysb.Monomer("A")
+    pysb.Initial(a(), pysb.Parameter("a0"))
+    pysb.Rule("deg", a() >> None, pysb.Parameter("kk", 1.0))
+
+    events = [
+        Event(
+            # note that unlike for SBML import, we must omit the real=True here
+            identifier=sp.Symbol("event1"),
+            name="Event1",
+            value=amici_time_symbol - 5,
+            assignments={sp.Symbol("__s0"): sp.Symbol("__s0") + 1000},
+            use_values_from_trigger_time=False,
+        )
+    ]
+
+    outdir = tempdir
+    pysb2amici(
+        model,
+        outdir,
+        verbose=True,
+        observation_model=[amici.MeasurementChannel("a")],
+        compute_conservation_laws=False,
+        _events=events,
+    )
+
+    model_module = amici.import_model_module(
+        module_name=model.name, module_path=outdir
+    )
+    amici_model = model_module.get_model()
+    assert amici_model.ne
+    amici_model.set_timepoints([0, 4, 5])
+
+    np.testing.assert_allclose(
+        amici_model.simulate().x, np.array([[0.0], [0.0], [1000.0]])
+    )