Fix piecewise/Heaviside handling

python/sdist/amici/de_export.py

@@ -58,8 +58,8 @@
     generate_flux_symbol,
     smart_subs_dict,
     strip_pysb,
-    symbol_with_assumptions,
     toposort_symbols,
+    unique_preserve_order,
 )
 from .logging import get_logger, log_execution_time, set_log_level
 
@@ -2745,16 +2745,12 @@ def _process_heavisides(
         :returns:
             dxdt with Heaviside functions replaced by amici helper variables
         """
-
-        # expanding the rhs will in general help to collect the same
-        # heaviside function
-        dt_expanded = dxdt.expand()
         # track all the old Heaviside expressions in tmp_roots_old
         # replace them later by the new expressions
         heavisides = []
         # run through the expression tree and get the roots
-        tmp_roots_old = self._collect_heaviside_roots(dt_expanded.args)
-        for tmp_old in tmp_roots_old:
+        tmp_roots_old = self._collect_heaviside_roots(dxdt.args)
+        for tmp_old in unique_preserve_order(tmp_roots_old):
             # we want unique identifiers for the roots
             tmp_new = self._get_unique_root(tmp_old, roots)
             # `tmp_new` is None if the root is not time-dependent.

python/sdist/amici/import_utils.py

@@ -734,5 +734,20 @@ def strip_pysb(symbol: sp.Basic) -> sp.Basic:
         return symbol
 
 
+def unique_preserve_order(seq: Sequence) -> list:
+    """Return a list of unique elements in Sequence, keeping only the first
+    occurrence of each element
+
+    Parameters:
+        seq: Sequence to prune
+
+    Returns:
+        List of unique elements in ``seq``
+    """
+    seen = set()
+    seen_add = seen.add
+    return [x for x in seq if not (x in seen or seen_add(x))]
+
+
 sbml_time_symbol = symbol_with_assumptions("time")
 amici_time_symbol = symbol_with_assumptions("t")

python/sdist/amici/sbml_import.py

@@ -38,7 +38,6 @@
     DEModel,
     _default_simplify,
     smart_is_zero_matrix,
-    symbol_with_assumptions,
 )
 from .import_utils import (
     RESERVED_SYMBOLS,
@@ -54,6 +53,7 @@
     sbml_time_symbol,
     smart_subs,
     smart_subs_dict,
+    symbol_with_assumptions,
     toposort_symbols,
 )
 from .logging import get_logger, log_execution_time, set_log_level