add support for heavyside functions

AMICI-dev · Oct 21, 2024 · 7faae32 · 7faae32
1 parent d6c5bcd
commit 7faae32
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/python/sdist/amici/de_export.py b/python/sdist/amici/de_export.py
@@ -21,6 +21,7 @@
     TYPE_CHECKING,
     Literal,
 )
+
 import sympy as sp
 
 from . import (
@@ -305,7 +306,17 @@ def jnp_stack_str(array) -> str:
                 f"{eq_name.upper()}_EQ": "\n".join(
                     self._code_printer_jax._get_sym_lines(
                         (str(strip_pysb(s)) for s in self.model.sym(eq_name)),
-                        self.model.eq(eq_name),
+                        self.model.eq(eq_name).subs(
+                            dict(
+                                zip(
+                                    self.model.sym("h"),
+                                    (
+                                        sp.Heaviside(x)
+                                        for x in self.model.eq("root")
+                                    ),
+                                )
+                            )
+                        ),
                         indent,
                     )
                 )

diff --git a/tests/benchmark-models/test_petab_model.py b/tests/benchmark-models/test_petab_model.py
@@ -145,13 +145,16 @@ def main():
     rdatas = res[RDATAS]
     llh = res[LLH]
 
-    if args.model_name not in (
-        "Beer_MolBioSystems2014",
-        "Brannmark_JBC2010",
-        "Fujita_SciSignal2010",
-        "Isensee_JCB2018",
-        "Smith_BMCSystBiol2013",
-        "Weber_BMC2015",
+    if (
+        args.model_name
+        not in (
+            # "Beer_MolBioSystems2014",
+            # "Brannmark_JBC2010",
+            # "Fujita_SciSignal2010",
+            # "Isensee_JCB2018",
+            # "Smith_BMCSystBiol2013",
+            # "Weber_BMC2015",
+        )
     ):
         # Beer: Heaviside
         # Brannmark: Heaviside