Fix substitution of expressions into event roots (partially reverts #…

…1599) (#1784)

python/amici/ode_export.py

@@ -1297,7 +1297,7 @@ def parse_events(self) -> None:
             expr.set_val(self._process_heavisides(expr.get_val(), roots))
 
         # remove all possible Heavisides from roots, which may arise from
-        # the substitution of `'w'` in `_get_unique_root`
+        # the substitution of `'w'` in `_collect_heaviside_roots`
         for root in roots:
             root.set_val(self._process_heavisides(root.get_val(), roots))
 
@@ -2066,15 +2066,6 @@ def _get_unique_root(
             unique identifier for root, or ``None`` if the root is not
             time-dependent
         """
-        # substitute 'w' expressions into root expressions now, to avoid
-        # rewriting '{model_name}_root.cpp' and '{model_name}_stau.cpp' headers
-        # to include 'w.h'
-        w_sorted = toposort_symbols(dict(zip(
-            [expr.get_id() for expr in self._expressions],
-            [expr.get_val() for expr in self._expressions],
-        )))
-        root_found = root_found.subs(w_sorted)
-
         if not self._expr_is_time_dependent(root_found):
             return None
 
@@ -2115,6 +2106,18 @@ def _collect_heaviside_roots(
             elif arg.has(sp.Heaviside):
                 root_funs.extend(self._collect_heaviside_roots(arg.args))
 
+        # substitute 'w' expressions into root expressions now, to avoid
+        # rewriting '{model_name}_root.cpp' and '{model_name}_stau.cpp' headers
+        # to include 'w.h'
+        w_sorted = toposort_symbols(dict(zip(
+            [expr.get_id()  for expr in self._expressions],
+            [expr.get_val() for expr in self._expressions],
+        )))
+        root_funs = [
+            r.subs(w_sorted)
+            for r in root_funs
+        ]
+
         return root_funs
 
     def _process_heavisides(