Refactor rateOf handling

dweindl · dweindl · commit 5ccb05a58976 · 2025-11-14T22:38:33.000+01:00
Pull rateOf-handling out of DEModel and keep it along other SBML processing where it belongs.

This is easier to follow and prevents some lingering issues with the old approach due to the xdot / w interdependencies.

This also handles rateOf expressions in some additional, previously unsupported places like event assignments.
diff --git a/python/sdist/amici/de_model.py b/python/sdist/amici/de_model.py
@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import contextlib
 import copy
 import itertools
 import logging
@@ -47,7 +46,6 @@
 from .exporters.sundials.cxxcodeprinter import csc_matrix
 from .importers.utils import (
     ObservableTransformation,
-    SBMLException,
     _default_simplify,
     amici_time_symbol,
     smart_subs_dict,
@@ -404,165 +402,6 @@ def states(self) -> list[State]:
         """Get all states."""
         return self._differential_states + self._algebraic_states
 
-    def _process_sbml_rate_of(self) -> None:
-        """Substitute any SBML-rateOf constructs in the model equations"""
-        from sbmlmath import rate_of as rate_of_func
-
-        species_sym_to_xdot = dict(
-            zip(self.sym("x"), self.sym("xdot"), strict=True)
-        )
-        species_sym_to_idx = {x: i for i, x in enumerate(self.sym("x"))}
-
-        def get_rate(symbol: sp.Symbol):
-            """Get rate of change of the given symbol"""
-            if symbol.find(rate_of_func):
-                raise SBMLException("Nesting rateOf() is not allowed.")
-
-            # Replace all rateOf(some_species) by their respective xdot equation
-            with contextlib.suppress(KeyError):
-                return self._eqs["xdot"][species_sym_to_idx[symbol]]
-
-            # For anything other than a state, rateOf(.) is 0 or invalid
-            return 0
-
-        # replace rateOf-instances in xdot by xdot symbols
-        made_substitutions = False
-        for i_state in range(len(self.eq("xdot"))):
-            if rate_ofs := self._eqs["xdot"][i_state].find(rate_of_func):
-                self._eqs["xdot"][i_state] = self._eqs["xdot"][i_state].subs(
-                    {
-                        # either the rateOf argument is a state, or it's 0
-                        rate_of: species_sym_to_xdot.get(rate_of.args[0], 0)
-                        for rate_of in rate_ofs
-                    }
-                )
-                made_substitutions = True
-
-        if made_substitutions:
-            # substitute in topological order
-            subs = toposort_symbols(
-                dict(zip(self.sym("xdot"), self.eq("xdot"), strict=True))
-            )
-            self._eqs["xdot"] = smart_subs_dict(self.eq("xdot"), subs)
-
-        # replace rateOf-instances in w by xdot equation
-        #  here we may need toposort, as xdot may depend on w
-        made_substitutions = False
-        for i_expr in range(len(self.eq("w"))):
-            new, replacement = self._eqs["w"][i_expr].replace(
-                rate_of_func, get_rate, map=True
-            )
-            if replacement:
-                self._eqs["w"][i_expr] = new
-                made_substitutions = True
-
-        if made_substitutions:
-            # Sort expressions in self._expressions, w symbols, and w equations
-            #  in topological order. Ideally, this would already happen before
-            #  adding the expressions to the model, but at that point, we don't
-            #  have access to xdot yet.
-            # NOTE: elsewhere, conservations law expressions are expected to
-            #  occur before any other w expressions, so we must maintain their
-            #  position
-            # toposort everything but conservation law expressions,
-            #  then prepend conservation laws
-            w_sorted = toposort_symbols(
-                dict(
-                    zip(
-                        self.sym("w")[self.num_cons_law() :, :],
-                        self.eq("w")[self.num_cons_law() :, :],
-                        strict=True,
-                    )
-                )
-            )
-            w_sorted = (
-                dict(
-                    zip(
-                        self.sym("w")[: self.num_cons_law(), :],
-                        self.eq("w")[: self.num_cons_law(), :],
-                        strict=True,
-                    )
-                )
-                | w_sorted
-            )
-            old_syms = tuple(self._syms["w"])
-            topo_expr_syms = tuple(w_sorted.keys())
-            new_order = [old_syms.index(s) for s in topo_expr_syms]
-            self._expressions = [self._expressions[i] for i in new_order]
-            self._syms["w"] = sp.Matrix(topo_expr_syms)
-            self._eqs["w"] = sp.Matrix(list(w_sorted.values()))
-
-        # replace rateOf-instances in x0 by xdot equation
-        # indices of state variables whose x0 was modified
-        changed_indices = []
-        for i_state in range(len(self.eq("x0"))):
-            new, replacement = self._eqs["x0"][i_state].replace(
-                rate_of_func, get_rate, map=True
-            )
-            if replacement:
-                self._eqs["x0"][i_state] = new
-                changed_indices.append(i_state)
-        if changed_indices:
-            # Replace any newly introduced state variables
-            #  by their x0 expressions.
-            # Also replace any newly introduced `w` symbols by their
-            #  expressions (after `w` was toposorted above).
-            subs = toposort_symbols(
-                dict(zip(self.sym("x_rdata"), self.eq("x0"), strict=True))
-            )
-            subs = dict(zip(self._syms["w"], self.eq("w"), strict=True)) | subs
-            for i_state in changed_indices:
-                self._eqs["x0"][i_state] = smart_subs_dict(
-                    self._eqs["x0"][i_state], subs
-                )
-
-        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."
-                    #  see, e.g., sbml test case 01482
-                    raise SBMLException(
-                        "AMICI does currently not support rateOf(.) inside AlgebraicRules."
-                    )
-
-                component.set_val(
-                    component.get_val().subs(
-                        {
-                            rate_of: get_rate(rate_of.args[0])
-                            for rate_of in rate_ofs
-                        }
-                    )
-                )
-
-        for event in self.events():
-            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(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."
-                    )
-                    # TODO here we need xdot sym, not eqs
-                    # event._state_update[i_state] = event._state_update[i_state].subs(
-                    #     {rate_of: get_rate(rate_of.args[0]) for rate_of in rate_ofs}
-                    # )
-
     def add_component(
         self, component: ModelQuantity, insert_first: bool | None = False
     ) -> None:
@@ -2752,3 +2591,37 @@ def has_event_assignments(self) -> bool:
             boolean indicating if event assignments are present
         """
         return any(event.updates_state for event in self._events)
+
+    def toposort_expressions(self) -> dict[sp.Symbol, sp.Expr]:
+        """
+        Sort expressions in topological order.
+
+        :return:
+            dict of expression symbols to expressions in topological order
+        """
+        # NOTE: elsewhere, conservations law expressions are expected to
+        #  occur before any other w expressions, so we must maintain their
+        #  position.
+        # toposort everything but conservation law expressions,
+        #  then prepend conservation laws
+        if self._syms or self._eqs:
+            raise AssertionError(
+                "This function must be called before generating any symbols "
+                "or equations."
+            )
+        w_toposorted = toposort_symbols(
+            {
+                e.get_sym(): e.get_val()
+                for e in self.expressions()[self.num_cons_law() :]
+            }
+        )
+
+        w_toposorted = {
+            e.get_sym(): e.get_val()
+            for e in self.expressions()[: self.num_cons_law()]
+        } | w_toposorted
+        old_syms = [e.get_sym() for e in self.expressions()]
+        self._expressions = [
+            self._expressions[old_syms.index(s)] for s in w_toposorted
+        ]
+        return w_toposorted
diff --git a/python/sdist/amici/importers/sbml/__init__.py b/python/sdist/amici/importers/sbml/__init__.py
@@ -5,6 +5,7 @@
 in the `Systems Biology Markup Language (SBML) <https://sbml.org/>`_.
 """
 
+import contextlib
 import copy
 import itertools as itt
 import logging
@@ -14,6 +15,7 @@
 import warnings
 import xml.etree.ElementTree as ET
 from collections.abc import Callable, Iterable, Sequence
+from itertools import chain
 from pathlib import Path
 from typing import (
     Any,
@@ -30,7 +32,12 @@
 from amici import get_model_dir, has_clibs
 from amici.constants import SymbolId
 from amici.de_model import DEModel
-from amici.de_model_components import Expression, symbol_to_type
+from amici.de_model_components import (
+    DifferentialState,
+    Event,
+    Expression,
+    symbol_to_type,
+)
 from amici.importers.sbml.splines import AbstractSpline
 from amici.importers.sbml.utils import SBMLException
 from amici.importers.utils import (
@@ -682,12 +689,12 @@ def _build_ode_model(
         if hybridization:
             ode_model._process_hybridization(hybridization)
 
+        # substitute SBML-rateOf constructs
+        self._process_sbml_rate_of(ode_model)
+
         # fill in 'self._sym' based on prototypes and components in ode_model
         ode_model.generate_basic_variables()
 
-        # substitute SBML-rateOf constructs
-        ode_model._process_sbml_rate_of()
-
         return ode_model
 
     @log_execution_time("importing SBML", logger)
@@ -3014,6 +3021,115 @@ def _transform_dxdt_to_concentration(
 
             return dxdt / v
 
+    def _process_sbml_rate_of(self, de_model: DEModel) -> None:
+        """Substitute any SBML-rateOf constructs in the model equations"""
+        from sbmlmath import rate_of as rate_of_func
+
+        sym_to_state_var: dict[sp.Symbol, DifferentialState] = {
+            state.get_sym(): state for state in de_model.differential_states()
+        }
+
+        is_ode = de_model.is_ode()
+
+        def get_rate(symbol: sp.Symbol):
+            """Get rate of change of the given symbol
+            (symbol is argument of rateOf)"""
+            if not is_ode:
+                # Not implemented for DAE models
+                #  sbml test case 01482
+                raise SBMLException(
+                    "rateOf() is only supported in ODE models."
+                )
+
+            if symbol.find(rate_of_func):
+                raise SBMLException("Nesting rateOf() is not allowed.")
+
+            # Replace all rateOf(some_state_var) by their respective
+            #  xdot equation
+            with contextlib.suppress(KeyError):
+                return sym_to_state_var[symbol].get_dt()
+
+            # For anything other than a state variable,
+            #  rateOf(.) is 0 or invalid
+            return sp.Integer(0)
+
+        def do_subs(expr, rate_ofs) -> sp.Expr:
+            """Substitute rateOf(...) in expr by appropriate expressions"""
+            expr = expr.subs(
+                {rate_of: get_rate(rate_of.args[0]) for rate_of in rate_ofs}
+            )
+            if rate_ofs := expr.find(rate_of_func):
+                # recursively substitute until no rateOf remains
+                expr = do_subs(expr, rate_ofs)
+            return expr
+
+        # replace rateOf-instances in xdot
+        for state in de_model.differential_states():
+            if rate_ofs := state.get_dt().find(rate_of_func):
+                state.set_dt(do_subs(state.get_dt(), rate_ofs))
+
+        # replace rateOf-instances in expressions which we will need for
+        #  substitutions later
+        for expr in de_model.expressions():
+            if rate_ofs := expr.get_val().find(rate_of_func):
+                expr.set_val(do_subs(expr.get_val(), rate_ofs))
+
+        w_toposorted = de_model.toposort_expressions()
+
+        # replace rateOf-instances in x0
+        # indices of state variables whose x0 was modified
+        changed_indices = []
+        for i_state, state in enumerate(de_model.differential_states()):
+            new, replacement = state.get_val().replace(
+                rate_of_func, get_rate, map=True
+            )
+            if replacement:
+                state.set_val(new)
+                changed_indices.append(i_state)
+        if changed_indices:
+            # Replace any newly introduced state variables
+            #  by their x0 expressions.
+            subs = w_toposorted | toposort_symbols(
+                {
+                    state.get_x_rdata(): state.get_val()
+                    for state in de_model.differential_states()
+                }
+            )
+            states = de_model.differential_states()
+            for i_state in changed_indices:
+                states[i_state].set_val(
+                    smart_subs_dict(states[i_state].get_val(), subs)
+                )
+
+        for component in chain(
+            de_model.observables(),
+            de_model.events(),
+        ):
+            if rate_ofs := component.get_val().find(rate_of_func):
+                component.set_val(do_subs(component.get_val(), rate_ofs))
+
+            if isinstance(component, Event):
+                if rate_ofs:
+                    # currently, `root` cannot depend on `w`.
+                    #  this could be changed, but for now,
+                    #  we just flatten out w expressions
+                    component.set_val(
+                        smart_subs_dict(component.get_val(), w_toposorted)
+                    )
+
+                # for events, also substitute in state updates
+                for target, assignment in component._assignments.items():
+                    if rate_ofs := assignment.find(rate_of_func):
+                        new_assignment = do_subs(assignment, rate_ofs)
+                        # currently, deltax cannot depend on `w`.
+                        #  this could be changed, or we could use xdot
+                        #  symbols, but for now, we just flatten out w
+                        #  expressions
+                        new_assignment = smart_subs_dict(
+                            new_assignment, w_toposorted
+                        )
+                        component._assignments[target] = new_assignment
+
 
 def _check_lib_sbml_errors(
     sbml_doc: libsbml.SBMLDocument, show_warnings: bool = False
