Refactor SBML event processing, defer computing state updates (#2742)

To support `use_values_from_trigger_time=true`, `deltax` may have to be
computed from `x_old` instead of `x`.
At the stage of processing the SBML model, there is no concept of `x_old`.
Thus, the state updates should only be computed inside the `DEModel`.
Also independently of `use_values_from_trigger_time`, this seems to be the
more appropriate place to compute the bolus.
Therefore, store only the event assignment in `Event`, and compute the bolus
on demand.

Author: dweindl

python/sdist/amici/de_model.py

@@ -534,11 +534,14 @@ def get_rate(symbol: sp.Symbol):
                 )
 
         for event in self.events():
-            if event._state_update is None:
+            state_update = event.get_state_update(
+                x=self.sym("x"), x_old=self.sym("x")
+            )
+            if state_update is None:
                 continue
 
-            for i_state in range(len(event._state_update)):
-                if rate_ofs := event._state_update[i_state].find(rate_of_func):
+            for i_state in range(len(state_update)):
+                if rate_ofs := state_update[i_state].find(rate_of_func):
                     raise SBMLException(
                         "AMICI does currently not support rateOf(.) inside event state updates."
                     )
@@ -1612,10 +1615,15 @@ def _compute_equation(self, name: str) -> None:
             # would cause problems when writing the function file later
             event_eqs = []
             for event in self._events:
-                if event._state_update is None:
+                # TODO https://github.com/AMICI-dev/AMICI/issues/2719
+                #   with use_values_from_trigger_time=True: x_old != x
+                state_update = event.get_state_update(
+                    x=self.sym("x"), x_old=self.sym("x")
+                )
+                if state_update is None:
                     event_eqs.append(sp.zeros(self.num_states_solver(), 1))
                 else:
-                    event_eqs.append(event._state_update)
+                    event_eqs.append(state_update)
 
             self._eqs[name] = event_eqs
 
@@ -1718,8 +1726,8 @@ def _compute_equation(self, name: str) -> None:
                         self.sym("stau").T,
                     )
 
-                # only add deltax part if there is state update
-                if event._state_update is not None:
+                # only add deltax part if there is a state update
+                if event._assignments is not None:
                     # partial derivative for the parameters
                     tmp_eq += self.eq("ddeltaxdp")[ie]
 
@@ -2259,7 +2267,7 @@ def _get_unique_root(
                 identifier=sp.Symbol(root_symstr),
                 name=root_symstr,
                 value=root_found,
-                state_update=None,
+                assignments=None,
             )
         )
         return roots[-1].get_id()

python/sdist/amici/de_model_components.py

@@ -708,7 +708,7 @@ def __init__(
         identifier: sp.Symbol,
         name: str,
         value: sp.Expr,
-        state_update: sp.Expr | None,
+        assignments: dict[sp.Symbol, sp.Expr] | None = None,
         initial_value: bool | None = True,
         priority: sp.Basic | None = None,
     ):
@@ -724,17 +724,16 @@ def __init__(
         :param value:
             formula for the root / trigger function
 
-        :param state_update:
-            formula for the bolus function (None for Heaviside functions,
-            zero vector for events without bolus)
+        :param assignments:
+            Dictionary of event assignments: state symbol -> new value.
 
         :param initial_value:
             initial boolean value of the trigger function at t0. If set to
             `False`, events may trigger at ``t==t0``, otherwise not.
         """
         super().__init__(identifier, name, value)
         # add the Event specific components
-        self._state_update = state_update
+        self._assignments = assignments if assignments is not None else {}
         self._initial_value = initial_value
 
         if priority is not None and not priority.is_Number:
@@ -751,6 +750,35 @@ def __init__(
             # the trigger can't be solved for `t`
             self._t_root = []
 
+    def get_state_update(
+        self, x: sp.Matrix, x_old: sp.Matrix
+    ) -> sp.Matrix | None:
+        """
+        Get the state update (bolus) expression for the event assignment.
+
+        :param x: The current state vector.
+        :param x_old: The previous state vector.
+        :return: State-update matrix or ``None`` if no state update is defined.
+        """
+        if len(self._assignments) == 0:
+            return None
+
+        x_to_x_old = dict(zip(x, x_old))
+
+        def get_bolus(x_i: sp.Symbol) -> sp.Expr:
+            """
+            Get the bolus expression for a state variable.
+
+            :param x_i: state variable symbol
+            :return: bolus expression
+            """
+            if (assignment := self._assignments.get(x_i)) is not None:
+                return assignment.subs(x_to_x_old) - x_i
+            else:
+                return sp.Float(0.0)
+
+        return sp.Matrix([get_bolus(x_i) for x_i in x])
+
     def get_initial_value(self) -> bool:
         """
         Return the initial value for the root function.

python/sdist/amici/sbml_import.py

@@ -761,7 +761,7 @@ def _build_ode_model(
                 args += ["value"]
 
             if symbol_name == SymbolId.EVENT:
-                args += ["state_update", "initial_value", "priority"]
+                args += ["assignments", "initial_value", "priority"]
             elif symbol_name == SymbolId.OBSERVABLE:
                 args += ["transformation"]
             elif symbol_name == SymbolId.EVENT_OBSERVABLE:
@@ -1769,13 +1769,6 @@ def _process_events(self) -> None:
         """Process SBML events."""
         events = self.sbml.getListOfEvents()
 
-        def get_empty_bolus_value() -> sp.Float:
-            """
-            Used in the event update vector for species that are not affected
-            by the event.
-            """
-            return sp.Symbol("AMICI_EMTPY_BOLUS")
-
         # Used to update species concentrations when an event affects a
         # compartment.
         concentration_species_by_compartment = {
@@ -1811,7 +1804,7 @@ def get_empty_bolus_value() -> sp.Float:
             trigger = _parse_event_trigger(trigger_sym)
 
             # parse the boluses / event assignments
-            bolus = [get_empty_bolus_value() for _ in state_vector]
+            assignment_exprs = {}
             event_assignments = event.getListOfEventAssignments()
             compartment_event_assignments: set[tuple[sp.Symbol, sp.Expr]] = (
                 set()
@@ -1826,8 +1819,8 @@ def get_empty_bolus_value() -> sp.Float:
                 formula = self._sympify(event_assignment)
                 try:
                     # Try to find the species in the state vector.
-                    index = state_vector.index(variable_sym)
-                    bolus[index] = formula
+                    _ = state_vector.index(variable_sym)
+                    assignment_exprs[variable_sym] = formula
                 except ValueError:
                     raise SBMLException(
                         "Could not process event assignment for "
@@ -1864,30 +1857,17 @@ def get_empty_bolus_value() -> sp.Float:
                 ]:
                     # If the species was not affected by an event assignment,
                     # then the old value should be updated.
-                    if (
-                        bolus[state_vector.index(species_sym)]
-                        == get_empty_bolus_value()
-                    ):
+                    if species_sym not in assignment_exprs:
                         species_value = species_sym
                     # else the species was affected by an event assignment,
                     # hence the updated value should be updated further.
                     else:
-                        species_value = bolus[state_vector.index(species_sym)]
+                        species_value = assignment_exprs[species_sym]
                     # New species value is old amount / new volume.
-                    bolus[state_vector.index(species_sym)] = (
+                    assignment_exprs[species_sym] = (
                         species_value * compartment_sym / formula
                     )
 
-            # Subtract the current species value from each species with an
-            # update, as the bolus will be added on to the current species
-            # value during simulation.
-            for index in range(len(bolus)):
-                if bolus[index] != get_empty_bolus_value():
-                    bolus[index] -= state_vector[index]
-                bolus[index] = bolus[index].subs(
-                    get_empty_bolus_value(), sp.Float(0.0)
-                )
-
             initial_value = (
                 trigger_sbml.getInitialValue()
                 if trigger_sbml is not None
@@ -1917,7 +1897,7 @@ def get_empty_bolus_value() -> sp.Float:
             self.symbols[SymbolId.EVENT][event_sym] = {
                 "name": event_id,
                 "value": trigger,
-                "state_update": sp.MutableDenseMatrix(bolus),
+                "assignments": assignment_exprs,
                 "initial_value": initial_value,
                 "use_values_from_trigger_time": use_trig_val,
                 "priority": self._sympify(event.getPriority()),
@@ -1966,10 +1946,10 @@ def try_solve_t(expr: sp.Expr) -> list:
         # if all assignments are absolute (not referring to other non-constant
         # model entities), we are fine.
         if all(
-            update.is_zero or (update + variable).is_Number
+            assignment.is_Number
             for event in self.symbols[SymbolId.EVENT].values()
-            for variable, update in zip(state_vector, event["state_update"])
-            if not update.is_zero
+            for assignment in event["assignments"].values()
+            if event["assignments"] is not None
         ):
             return
 
@@ -2796,13 +2776,16 @@ def _replace_in_all_expressions(
             for element in self.symbols[symbol].values():
                 element["value"] = smart_subs(element["value"], old, new)
 
-        # replace in event state updates (boluses)
+        # replace in event assignments
         if self.symbols.get(SymbolId.EVENT, False):
             for event in self.symbols[SymbolId.EVENT].values():
-                for index in range(len(event["state_update"])):
-                    event["state_update"][index] = smart_subs(
-                        event["state_update"][index], old, new
-                    )
+                if event["assignments"] is not None:
+                    event["assignments"] = {
+                        smart_subs(target, old, new): smart_subs(
+                            expr, old, new
+                        )
+                        for target, expr in event["assignments"].items()
+                    }
 
         for state in {
             **self.symbols[SymbolId.SPECIES],