From c3d541295c03c4c810ac098da6948cee8a1456b0 Mon Sep 17 00:00:00 2001
From: Daniel Weindl <daniel.weindl@helmholtz-munich.de>
Date: Sat, 18 Nov 2023 11:47:45 +0100
Subject: [PATCH] missed

---
 python/sdist/amici/petab/util.py | 103 +++++++++++++++++++++++++++++++
 1 file changed, 103 insertions(+)
 create mode 100644 python/sdist/amici/petab/util.py

diff --git a/python/sdist/amici/petab/util.py b/python/sdist/amici/petab/util.py
new file mode 100644
index 0000000000..d30c1a6e9b
--- /dev/null
+++ b/python/sdist/amici/petab/util.py
@@ -0,0 +1,103 @@
+"""Various helper functions for working with PEtab problems."""
+import re
+from typing import Dict, Tuple, Union
+
+import libsbml
+import pandas as pd
+import petab
+from petab.C import PREEQUILIBRATION_CONDITION_ID, SIMULATION_CONDITION_ID
+from petab.mapping import resolve_mapping
+from petab.models import MODEL_TYPE_PYSB, MODEL_TYPE_SBML
+
+
+def get_states_in_condition_table(
+    petab_problem: petab.Problem,
+    condition: Union[Dict, pd.Series] = None,
+    return_patterns: bool = False,
+) -> Dict[str, Tuple[Union[float, str, None], Union[float, str, None]]]:
+    """Get states and their initial condition as specified in the condition table.
+
+    Returns: Dictionary: ``stateId -> (initial condition simulation, initial condition preequilibration)``
+    """
+    if petab_problem.model.type_id not in (MODEL_TYPE_SBML, MODEL_TYPE_PYSB):
+        raise NotImplementedError()
+
+    species_check_funs = {
+        MODEL_TYPE_SBML: lambda x: _element_is_sbml_state(
+            petab_problem.sbml_model, x
+        ),
+        MODEL_TYPE_PYSB: lambda x: _element_is_pysb_pattern(
+            petab_problem.model.model, x
+        ),
+    }
+    states = {
+        resolve_mapping(petab_problem.mapping_df, col): (None, None)
+        if condition is None
+        else (
+            petab_problem.condition_df.loc[
+                condition[SIMULATION_CONDITION_ID], col
+            ],
+            petab_problem.condition_df.loc[
+                condition[PREEQUILIBRATION_CONDITION_ID], col
+            ]
+            if PREEQUILIBRATION_CONDITION_ID in condition
+            else None,
+        )
+        for col in petab_problem.condition_df.columns
+        if species_check_funs[petab_problem.model.type_id](
+            resolve_mapping(petab_problem.mapping_df, col)
+        )
+    }
+
+    if petab_problem.model.type_id == MODEL_TYPE_PYSB:
+        if return_patterns:
+            return states
+        import pysb.pattern
+
+        if not petab_problem.model.model.species:
+            import pysb.bng
+
+            pysb.bng.generate_equations(petab_problem.model.model)
+
+        try:
+            spm = pysb.pattern.SpeciesPatternMatcher(
+                model=petab_problem.model.model
+            )
+        except NotImplementedError as e:
+            raise NotImplementedError(
+                "Requires https://github.com/pysb/pysb/pull/570. "
+                "To use this functionality, update pysb via "
+                "`pip install git+https://github.com/FFroehlich/pysb@fix_pattern_matching`"
+            )
+
+        # expose model components as variables so we can evaluate patterns
+        for c in petab_problem.model.model.components:
+            globals()[c.name] = c
+
+        states = {
+            f"__s{ix}": value
+            for pattern, value in states.items()
+            for ix in spm.match(eval(pattern), index=True, exact=True)
+        }
+    return states
+
+
+def _element_is_pysb_pattern(model: "pysb.Model", element: str) -> bool:
+    """Check if element is a pysb pattern"""
+    if match := re.match(r"[a-zA-Z_][\w_]*\(", element):
+        return match[0][:-1] in [m.name for m in model.monomers]
+    return False
+
+
+def _element_is_sbml_state(sbml_model: libsbml.Model, sbml_id: str) -> bool:
+    """Does the element with ID `sbml_id` correspond to a state variable?"""
+    if sbml_model.getCompartment(sbml_id) is not None:
+        return True
+    if sbml_model.getSpecies(sbml_id) is not None:
+        return True
+    if (
+        rule := sbml_model.getRuleByVariable(sbml_id)
+    ) is not None and rule.getTypeCode() == libsbml.SBML_RATE_RULE:
+        return True
+
+    return False