From e9536fd4ac4b62319f130c5a73c49d4b4b0405af Mon Sep 17 00:00:00 2001
From: FFroehlich <fabian@schaluck.com>
Date: Mon, 18 Jul 2022 17:46:02 +0200
Subject: [PATCH 1/2] fix_cls

---
 python/amici/pysb_import.py | 27 +++++++++++++++++++++++----
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/python/amici/pysb_import.py b/python/amici/pysb_import.py
index bc7ac80bd3..8049ad9112 100644
--- a/python/amici/pysb_import.py
+++ b/python/amici/pysb_import.py
@@ -1195,8 +1195,7 @@ def _apply_conseration_law_sub(cl: ConservationLaw,
 
     :return: boolean flag indicating whether the substitution was applied
     """
-    coeff = cl['coefficients'].get(sub[0], 0.0)
-    if coeff == 0.0 or cl['state'] == sub[0]:
+    if not _state_in_cl_formula(sub[0], cl):
         return False
 
     del cl['coefficients'][sub[0]]
@@ -1216,6 +1215,27 @@ def _apply_conseration_law_sub(cl: ConservationLaw,
     return True
 
 
+def _state_in_cl_formula(
+        state: sp.Symbol, cl: ConservationLaw
+) -> bool:
+    """
+    Checks whether state appears in the formula the provided cl
+
+    :param state:
+        state
+
+    :param cl:
+        conservation law
+
+    :return:
+        boolean indicator
+    """
+    if cl['state'] == state:
+        return False
+
+    return cl['coefficients'].get(state, 0.0) != 0.0
+
+
 def _get_conservation_law_subs(
         conservation_laws: List[ConservationLaw]
 ) -> List[Tuple[sp.Symbol, Dict[sp.Symbol, sp.Expr]]]:
@@ -1233,9 +1253,8 @@ def _get_conservation_law_subs(
     return [
         (cl['state'], cl['coefficients']) for cl in conservation_laws
         if any(
-            cl['state'] in other_cl['coefficients']
+            _state_in_cl_formula(cl['state'], other_cl)
             for other_cl in conservation_laws
-            if other_cl != cl
         )
     ]
 

From ef1a6ee4433686e286a58d6d39c0dacb90af4f45 Mon Sep 17 00:00:00 2001
From: FFroehlich <fabian@schaluck.com>
Date: Tue, 19 Jul 2022 09:59:29 +0200
Subject: [PATCH 2/2] fixup

---
 python/amici/pysb_import.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/amici/pysb_import.py b/python/amici/pysb_import.py
index 8049ad9112..d1b98ea1ac 100644
--- a/python/amici/pysb_import.py
+++ b/python/amici/pysb_import.py
@@ -29,7 +29,7 @@
                          ODEModel, Observable, Parameter, SigmaY, State)
 
 CL_Prototype = Dict[str, Dict[str, Any]]
-ConservationLaw = Dict[str, Union[str, sp.Basic]]
+ConservationLaw = Dict[str, Union[Dict, str, sp.Basic]]
 
 logger = get_logger(__name__, logging.ERROR)
 
@@ -1198,7 +1198,7 @@ def _apply_conseration_law_sub(cl: ConservationLaw,
     if not _state_in_cl_formula(sub[0], cl):
         return False
 
-    del cl['coefficients'][sub[0]]
+    coeff = cl['coefficients'].pop(sub[0], 0.0)
     # x_j = T/b_j - sum_{i≠j}(x_i * b_i) / b_j
     # don't need to account for totals here as we can simply
     # absorb that into the new total