Fix piecewise/Heavisides handling

* Substitute non-time-dependent Heavisides
* Substitute inside the expanded expression, otherwise not all substitution targets are found

Closes #2231

python/sdist/amici/de_export.py

@@ -2757,19 +2757,22 @@ def _process_heavisides(
         for tmp_old in tmp_roots_old:
             # we want unique identifiers for the roots
             tmp_new = self._get_unique_root(tmp_old, roots)
+            heavisides.append((sp.Heaviside(tmp_old), tmp_new))
+
             # `tmp_new` is None if the root is not time-dependent.
             if tmp_new is None:
                 continue
             # For Heavisides, we need to add the negative function as well
             self._get_unique_root(sp.sympify(-tmp_old), roots)
-            heavisides.append((sp.Heaviside(tmp_old), tmp_new))
 
         if heavisides:
             # only apply subs if necessary
             for heaviside_sympy, heaviside_amici in heavisides:
-                dxdt = dxdt.subs(heaviside_sympy, heaviside_amici)
+                dt_expanded = dt_expanded.subs(
+                    heaviside_sympy, heaviside_amici
+                )
 
-        return dxdt
+        return dt_expanded
 
 
 class DEExporter: