From 594b07e599334f5235c1f70bbddf08446d7a9673 Mon Sep 17 00:00:00 2001
From: Daniel Weindl <dweindl@users.noreply.github.com>
Date: Thu, 14 Dec 2023 10:45:35 +0100
Subject: [PATCH] Fix piecewise/Heaviside handling (#2234)

* Fixes a bug where Heaviside functions weren't correctly substituted, potentially resulting in DiracDelta's in the model equations. The problem was, the substitution targets were collected from an expanded expression, but the actual substitution was attempted on the non-expanded expression, which did not always succeed.

  Fixed by *not* expanding the expressions beforehand. Closes #2231

* Fixes redundant trigger-function processing

https://github.com/AMICI-dev/AMICI/pull/2234
---
 python/sdist/amici/de_export.py                | 10 +++-------
 python/sdist/amici/import_utils.py             | 15 +++++++++++++++
 python/sdist/amici/sbml_import.py              |  2 +-
 tests/benchmark-models/test_petab_benchmark.py |  0
 4 files changed, 19 insertions(+), 8 deletions(-)
 mode change 100755 => 100644 tests/benchmark-models/test_petab_benchmark.py

diff --git a/python/sdist/amici/de_export.py b/python/sdist/amici/de_export.py
index ac857b366c..2694f753ad 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
 
@@ -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.
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
diff --git a/tests/benchmark-models/test_petab_benchmark.py b/tests/benchmark-models/test_petab_benchmark.py
old mode 100755
new mode 100644