diff --git a/python/sdist/amici/de_export.py b/python/sdist/amici/de_export.py
index ffcd0f3993..3770068a5e 100644
--- a/python/sdist/amici/de_export.py
+++ b/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
 
@@ -2754,7 +2754,7 @@ def _process_heavisides(
         heavisides = []
         # run through the expression tree and get the roots
         tmp_roots_old = self._collect_heaviside_roots(dt_expanded.args)
-        for tmp_old in set(tmp_roots_old):
+        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.
diff --git a/python/sdist/amici/import_utils.py b/python/sdist/amici/import_utils.py
index 77a2add60b..2dfc60d0c5 100644
--- a/python/sdist/amici/import_utils.py
+++ b/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")
diff --git a/python/sdist/amici/sbml_import.py b/python/sdist/amici/sbml_import.py
index dd24b98cf8..b6a7b37ed4 100644
--- a/python/sdist/amici/sbml_import.py
+++ b/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