mxlpy 0.22.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

mxlpy/model.py

@@ -10,7 +10,9 @@ from __future__ import annotations
 import copy
 import inspect
 import itertools as it
+import logging
 from dataclasses import dataclass, field
+from queue import Empty, SimpleQueue
 from typing import TYPE_CHECKING, Self, cast
 
 import numpy as np
@@ -25,8 +27,8 @@ from mxlpy.meta.sympy_tools import (
 )
 from mxlpy.types import (
     AbstractSurrogate,
-    Array,
     Derived,
+    InitialAssignment,
     Parameter,
     Reaction,
     Readout,
@@ -37,16 +39,25 @@ if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
     from inspect import FullArgSpec
 
+    from sympy.physics.units.quantities import Quantity
+
     from mxlpy.types import Callable, Param, RateFn, RetType
 
+LOGGER = logging.getLogger(__name__)
+
 __all__ = [
     "ArityMismatchError",
     "CircularDependencyError",
     "Dependency",
+    "Failure",
+    "LOGGER",
+    "MdText",
     "MissingDependenciesError",
     "Model",
     "ModelCache",
     "TableView",
+    "UnitCheck",
+    "unit_of",
 ]
 
 
@@ -56,6 +67,80 @@ def _latex_view(expr: sympy.Expr | None) -> str:
     return f"${sympy.latex(expr)}$"
 
 
+def unit_of(expr: sympy.Expr) -> sympy.Expr:
+    """Get unit of sympy expr."""
+    return expr.as_coeff_Mul()[1]
+
+
+@dataclass
+class Failure:
+    """Unit test failure."""
+
+    expected: sympy.Expr
+    obtained: sympy.Expr
+
+    @property
+    def difference(self) -> sympy.Expr:
+        """Difference between expected and obtained unit."""
+        return self.expected / self.obtained  # type: ignore
+
+
+@dataclass
+class MdText:
+    """Generic markdown text."""
+
+    content: list[str]
+
+    def _repr_markdown_(self) -> str:
+        return "\n".join(self.content)
+
+
+@dataclass
+class UnitCheck:
+    """Container for unit check."""
+
+    per_variable: dict[str, dict[str, bool | Failure | None]]
+
+    @staticmethod
+    def _fmt_success(s: str) -> str:
+        return f"<span style='color: green'>{s}</span>"
+
+    @staticmethod
+    def _fmt_failed(s: str) -> str:
+        return f"<span style='color: red'>{s}</span>"
+
+    def correct_diff_eqs(self) -> dict[str, bool]:
+        """Get all correctly annotated reactions by variable."""
+        return {
+            var: all(isinstance(i, bool) for i in checks.values())
+            for var, checks in self.per_variable.items()
+        }
+
+    def report(self) -> MdText:
+        """Export check as markdown report."""
+        report = ["## Type check"]
+        for diff_eq, res in self.correct_diff_eqs().items():
+            txt = self._fmt_success("Correct") if res else self._fmt_failed("Failed")
+            report.append(f"\n### d{diff_eq}dt: {txt}")
+
+            if res:
+                continue
+            for k, v in self.per_variable[diff_eq].items():
+                match v:
+                    case bool():
+                        continue
+                    case None:
+                        report.append(f"\n- {k}")
+                        report.append("  - Failed to parse")
+                    case Failure(expected, obtained):
+                        report.append(f"\n- {k}")
+                        report.append(f"  - expected: {_latex_view(expected)}")
+                        report.append(f"  - obtained: {_latex_view(obtained)}")
+                        report.append(f"  - difference: {_latex_view(v.difference)}")
+
+        return MdText(report)
+
+
 @dataclass(kw_only=True, slots=True)
 class TableView:
     """Markdown view of pandas Dataframe.
@@ -228,8 +313,6 @@ def _sort_dependencies(
         SortError: If circular dependencies are detected
 
     """
-    from queue import Empty, SimpleQueue
-
     _check_if_is_sortable(available, elements)
 
     order = []
@@ -291,12 +374,13 @@ class ModelCache:
 
     """
 
+    order: list[str]  # mostly for debug purposes
     var_names: list[str]
     dyn_order: list[str]
+    base_parameter_values: dict[str, float]
     all_parameter_values: dict[str, float]
     stoich_by_cpds: dict[str, dict[str, float]]
     dyn_stoich_by_cpds: dict[str, dict[str, Derived]]
-    dxdt: pd.Series
     initial_conditions: dict[str, float]
 
 
@@ -317,7 +401,7 @@ class Model:
 
     """
 
-    _ids: dict[str, str] = field(default_factory=dict)
+    _ids: dict[str, str] = field(default_factory=dict, repr=False)
     _variables: dict[str, Variable] = field(default_factory=dict)
     _parameters: dict[str, Parameter] = field(default_factory=dict)
     _derived: dict[str, Derived] = field(default_factory=dict)
@@ -343,11 +427,32 @@ class Model:
             ModelCache: An instance of ModelCache containing the initialized cache data.
 
         """
-        all_parameter_values: dict[str, float] = self.get_parameter_values()
-        all_parameter_names: set[str] = set(all_parameter_values)
+        parameter_names = set(self._parameters)
+        all_parameter_names = set(parameter_names)  # later include static derived
+
+        base_parameter_values: dict[str, float] = {
+            k: val
+            for k, v in self._parameters.items()
+            if not isinstance(val := v.value, InitialAssignment)
+        }
+        base_variable_values: dict[str, float] = {
+            k: init
+            for k, v in self._variables.items()
+            if not isinstance(init := v.initial_value, InitialAssignment)
+        }
+        initial_assignments: dict[str, InitialAssignment] = {
+            k: init
+            for k, v in self._variables.items()
+            if isinstance(init := v.initial_value, InitialAssignment)
+        } | {
+            k: init
+            for k, v in self._parameters.items()
+            if isinstance(init := v.value, InitialAssignment)
+        }
 
         # Sanity checks
         for name, el in it.chain(
+            initial_assignments.items(),
             self._derived.items(),
             self._reactions.items(),
             self._readouts.items(),
@@ -356,25 +461,17 @@ class Model:
                 raise ArityMismatchError(name, el.fn, el.args)
 
         # Sort derived & reactions
+        available = (
+            set(base_parameter_values)
+            | set(base_variable_values)
+            | set(self._data)
+            | {"time"}
+        )
         to_sort = (
-            self._derived
-            | self._reactions
-            | self._surrogates
-            | {
-                k: init
-                for k, v in self._variables.items()
-                if isinstance(init := v.initial_value, Derived)
-            }
+            initial_assignments | self._derived | self._reactions | self._surrogates
         )
         order = _sort_dependencies(
-            available=all_parameter_names
-            | {
-                k
-                for k, v in self._variables.items()
-                if not isinstance(v.initial_value, Derived)
-            }
-            | set(self._data)
-            | {"time"},
+            available=available,
             elements=[
                 Dependency(name=k, required=set(v.args), provided={k})
                 if not isinstance(v, AbstractSurrogate)
@@ -386,14 +483,7 @@ class Model:
         # Calculate all values once, including dynamic ones
         # That way, we can make initial conditions dependent on e.g. rates
         dependent = (
-            all_parameter_values
-            | self._data
-            | {
-                k: init
-                for k, v in self._variables.items()
-                if not isinstance(init := v.initial_value, Derived)
-            }
-            | {"time": 0.0}
+            base_parameter_values | base_variable_values | self._data | {"time": 0.0}
         )
         for name in order:
             to_sort[name].calculate_inpl(name, dependent)
@@ -404,7 +494,7 @@ class Model:
         for name in order:
             if name in self._reactions or name in self._surrogates:
                 dyn_order.append(name)
-            elif name in self._variables:
+            elif name in self._variables or name in self._parameters:
                 static_order.append(name)
             else:
                 derived = self._derived[name]
@@ -445,29 +535,27 @@ class Model:
                         d_static[rxn_name] = factor
 
         var_names = self.get_variable_names()
-        dxdt = pd.Series(np.zeros(len(var_names), dtype=float), index=var_names)
-
         initial_conditions: dict[str, float] = {
-            k: init
-            for k, v in self._variables.items()
-            if not isinstance(init := v.initial_value, Derived)
+            k: cast(float, dependent[k]) for k in self._variables
         }
+        all_parameter_values = dict(base_parameter_values)
         for name in static_order:
             if name in self._variables:
-                initial_conditions[name] = cast(float, dependent[name])
-            elif name in self._derived:
+                continue  # handled in initial_conditions above
+            if name in self._parameters or name in self._derived:
                 all_parameter_values[name] = cast(float, dependent[name])
             else:
                 msg = "Unknown target for static derived variable."
                 raise KeyError(msg)
 
         self._cache = ModelCache(
+            order=order,
             var_names=var_names,
             dyn_order=dyn_order,
+            base_parameter_values=base_parameter_values,
             all_parameter_values=all_parameter_values,
             stoich_by_cpds=stoich_by_compounds,
             dyn_stoich_by_cpds=dyn_stoich_by_compounds,
-            dxdt=dxdt,
             initial_conditions=initial_conditions,
         )
         return self._cache
@@ -529,14 +617,25 @@ class Model:
     def parameters(self) -> TableView:
         """Return view of parameters."""
         index = list(self._parameters.keys())
-        data = [
-            {
-                "value": el.value,
-                "unit": _latex_view(unit) if (unit := el.unit) is not None else "",
-                # "source": ...,
-            }
-            for el in self._parameters.values()
-        ]
+        data = []
+        for name, el in self._parameters.items():
+            if isinstance(init := el.value, InitialAssignment):
+                value_str = _latex_view(
+                    fn_to_sympy(
+                        init.fn,
+                        origin=name,
+                        model_args=list_of_symbols(init.args),
+                    )
+                )
+            else:
+                value_str = str(init)
+            data.append(
+                {
+                    "value": value_str,
+                    "unit": _latex_view(unit) if (unit := el.unit) is not None else "",
+                    # "source": ...,
+                }
+            )
         return TableView(data=pd.DataFrame(data, index=index))
 
     def get_raw_parameters(self, *, as_copy: bool = True) -> dict[str, Parameter]:
@@ -557,7 +656,9 @@ class Model:
                   and the values are parameter values (as floats).
 
         """
-        return {k: v.value for k, v in self._parameters.items()}
+        if (cache := self._cache) is None:
+            cache = self._create_cache()
+        return cache.base_parameter_values
 
     def get_parameter_names(self) -> list[str]:
         """Retrieve the names of the parameters.
@@ -580,7 +681,7 @@ class Model:
     def add_parameter(
         self,
         name: str,
-        value: float,
+        value: float | InitialAssignment,
         unit: sympy.Expr | None = None,
         source: str | None = None,
     ) -> Self:
@@ -603,7 +704,9 @@ class Model:
         self._parameters[name] = Parameter(value=value, unit=unit, source=source)
         return self
 
-    def add_parameters(self, parameters: Mapping[str, float | Parameter]) -> Self:
+    def add_parameters(
+        self, parameters: Mapping[str, float | Parameter | InitialAssignment]
+    ) -> Self:
         """Adds multiple parameters to the model.
 
         Examples:
@@ -671,7 +774,7 @@ class Model:
     def update_parameter(
         self,
         name: str,
-        value: float | None = None,
+        value: float | InitialAssignment | None = None,
         *,
         unit: sympy.Expr | None = None,
         source: str | None = None,
@@ -695,7 +798,7 @@ class Model:
 
         """
         if name not in self._parameters:
-            msg = f"'{name}' not found in parameters"
+            msg = f"{name!r} not found in parameters"
             raise KeyError(msg)
 
         parameter = self._parameters[name]
@@ -707,7 +810,9 @@ class Model:
             parameter.source = source
         return self
 
-    def update_parameters(self, parameters: Mapping[str, float | Parameter]) -> Self:
+    def update_parameters(
+        self, parameters: Mapping[str, float | Parameter | InitialAssignment]
+    ) -> Self:
         """Update multiple parameters of the model.
 
         Examples:
@@ -741,7 +846,17 @@ class Model:
             Self: The instance of the class with the updated parameter.
 
         """
-        return self.update_parameter(name, self._parameters[name].value * factor)
+        old = self._parameters[name].value
+        if isinstance(old, InitialAssignment):
+            LOGGER.warning("Overwriting initial assignment %s", name)
+            if (cache := self._cache) is None:
+                cache = self._create_cache()
+
+            return self.update_parameter(
+                name, cache.all_parameter_values[name] * factor
+            )
+
+        return self.update_parameter(name, old * factor)
 
     def scale_parameters(self, parameters: dict[str, float]) -> Self:
         """Scales the parameters of the model.
@@ -843,7 +958,7 @@ class Model:
         index = list(self._variables.keys())
         data = []
         for name, el in self._variables.items():
-            if isinstance(init := el.initial_value, Derived):
+            if isinstance(init := el.initial_value, InitialAssignment):
                 value_str = _latex_view(
                     fn_to_sympy(
                         init.fn,
@@ -909,7 +1024,7 @@ class Model:
     def add_variable(
         self,
         name: str,
-        initial_value: float | Derived,
+        initial_value: float | InitialAssignment,
         unit: sympy.Expr | None = None,
         source: str | None = None,
     ) -> Self:
@@ -935,7 +1050,7 @@ class Model:
         return self
 
     def add_variables(
-        self, variables: Mapping[str, float | Variable | Derived]
+        self, variables: Mapping[str, float | Variable | InitialAssignment]
     ) -> Self:
         """Adds multiple variables to the model with their initial conditions.
 
@@ -1001,7 +1116,7 @@ class Model:
     def update_variable(
         self,
         name: str,
-        initial_value: float | Derived,
+        initial_value: float | InitialAssignment,
         unit: sympy.Expr | None = None,
         source: str | None = None,
     ) -> Self:
@@ -1035,7 +1150,7 @@ class Model:
         return self
 
     def update_variables(
-        self, variables: Mapping[str, float | Derived | Variable]
+        self, variables: Mapping[str, float | Variable | InitialAssignment]
     ) -> Self:
         """Updates multiple variables in the model.
 
@@ -1104,7 +1219,7 @@ class Model:
         return self
 
     ##########################################################################
-    # Derived - views
+    # Derived
     ##########################################################################
 
     @property
@@ -1713,11 +1828,28 @@ class Model:
             return copy.deepcopy(self._surrogates)
         return self._surrogates
 
-    def get_surrogate_output_names(self) -> list[str]:
-        """Return output names by surrogates."""
+    def get_surrogate_output_names(
+        self,
+        *,
+        include_fluxes: bool = True,
+    ) -> list[str]:
+        """Return output names by surrogates.
+
+        Optionally filter out the names of which surrogate outfluxes are actually
+        fluxes / reactions rather than variables.
+
+        Args:
+            include_fluxes: whether to also include outputs which are reaction
+                            names
+
+        """
         names = []
-        for i in self._surrogates.values():
-            names.extend(i.outputs)
+        if include_fluxes:
+            for i in self._surrogates.values():
+                names.extend(i.outputs)
+        else:
+            for i in self._surrogates.values():
+                names.extend(x for x in i.outputs if x not in i.stoichiometries)
         return names
 
     def get_surrogate_reaction_names(self) -> list[str]:
@@ -1765,7 +1897,8 @@ class Model:
         include_derived_parameters: bool,
         include_derived_variables: bool,
         include_reactions: bool,
-        include_surrogate_outputs: bool,
+        include_surrogate_variables: bool,
+        include_surrogate_fluxes: bool,
         include_readouts: bool,
     ) -> list[str]:
         """Get names of all kinds of model components."""
@@ -1782,8 +1915,10 @@ class Model:
             names.extend(self.get_derived_parameter_names())
         if include_reactions:
             names.extend(self.get_reaction_names())
-        if include_surrogate_outputs:
-            names.extend(self.get_surrogate_output_names())
+        if include_surrogate_variables:
+            names.extend(self.get_surrogate_output_names(include_fluxes=False))
+        if include_surrogate_fluxes:
+            names.extend(self.get_surrogate_reaction_names())
         if include_readouts:
             names.extend(self.get_readout_names())
         return names
@@ -1837,7 +1972,8 @@ class Model:
         include_derived_parameters: bool = True,
         include_derived_variables: bool = True,
         include_reactions: bool = True,
-        include_surrogate_outputs: bool = True,
+        include_surrogate_variables: bool = True,
+        include_surrogate_fluxes: bool = True,
         include_readouts: bool = False,
     ) -> pd.Series:
         """Generate a pandas Series of arguments for the model.
@@ -1864,7 +2000,8 @@ class Model:
             include_derived_parameters: Whether to include derived parameters
             include_derived_variables: Whether to include derived variables
             include_reactions: Whether to include reactions
-            include_surrogate_outputs: Whether to include surrogate outputs
+            include_surrogate_variables: Whether to include derive variables obtained from surrogate
+            include_surrogate_fluxes: Whether to include surrogate fluxes
             include_readouts: Whether to include readouts
 
         Returns:
@@ -1890,7 +2027,8 @@ class Model:
                 include_derived_parameters=include_derived_parameters,
                 include_derived_variables=include_derived_variables,
                 include_reactions=include_reactions,
-                include_surrogate_outputs=include_surrogate_outputs,
+                include_surrogate_variables=include_surrogate_variables,
+                include_surrogate_fluxes=include_surrogate_fluxes,
                 include_readouts=include_readouts,
             )
         ]
@@ -1926,7 +2064,8 @@ class Model:
         include_derived_parameters: bool = True,
         include_derived_variables: bool = True,
         include_reactions: bool = True,
-        include_surrogate_outputs: bool = True,
+        include_surrogate_variables: bool = True,
+        include_surrogate_fluxes: bool = True,
         include_readouts: bool = False,
     ) -> pd.DataFrame:
         """Generate a DataFrame containing time course arguments for model evaluation.
@@ -1949,7 +2088,8 @@ class Model:
             include_derived_parameters: Whether to include derived parameters
             include_derived_variables: Whether to include derived variables
             include_reactions: Whether to include reactions
-            include_surrogate_outputs: Whether to include surrogate outputs
+            include_surrogate_variables: Whether to include variables derived from surrogates
+            include_surrogate_fluxes: Whether to include surrogate fluxes
             include_readouts: Whether to include readouts
 
         Returns:
@@ -1974,7 +2114,8 @@ class Model:
                 include_derived_parameters=include_derived_parameters,
                 include_derived_variables=include_derived_variables,
                 include_reactions=include_reactions,
-                include_surrogate_outputs=include_surrogate_outputs,
+                include_surrogate_variables=include_surrogate_variables,
+                include_surrogate_fluxes=include_surrogate_fluxes,
                 include_readouts=include_readouts,
             ),
         ]
@@ -2011,15 +2152,19 @@ class Model:
             Fluxes: A pandas Series containing the fluxes for each reaction.
 
         """
-        names = self.get_reaction_names()
-        names.extend(self.get_surrogate_reaction_names())
-
-        args = self.get_args(
+        return self.get_args(
             variables=variables,
             time=time,
+            include_time=False,
+            include_variables=False,
+            include_parameters=False,
+            include_derived_parameters=False,
+            include_derived_variables=False,
+            include_reactions=True,
+            include_surrogate_variables=False,
+            include_surrogate_fluxes=True,
             include_readouts=False,
         )
-        return args.loc[names]
 
     def get_fluxes_time_course(self, variables: pd.DataFrame) -> pd.DataFrame:
         """Generate a time course of fluxes for the given reactions and surrogates.
@@ -2029,7 +2174,9 @@ class Model:
                 pd.DataFrame({"v1": [0.1, 0.2], "v2": [0.2, 0.3]})
 
         This method calculates the fluxes for each reaction in the model using the provided
-        arguments and combines them with the outputs from the surrogates to create a complete
+        arguments and combines them wit      names: list[str] = self.get_reaction_names()
+        for surrogate in self._surrogates.values():
+            names.extend(surrogate.stoichiometries)h the outputs from the surrogates to create a complete
         time course of fluxes.
 
         Args:
@@ -2043,21 +2190,23 @@ class Model:
                           the index of the input arguments.
 
         """
-        names: list[str] = self.get_reaction_names()
-        for surrogate in self._surrogates.values():
-            names.extend(surrogate.stoichiometries)
-
-        variables = self.get_args_time_course(
+        return self.get_args_time_course(
             variables=variables,
+            include_variables=False,
+            include_parameters=False,
+            include_derived_parameters=False,
+            include_derived_variables=False,
+            include_reactions=True,
+            include_surrogate_variables=False,
+            include_surrogate_fluxes=True,
             include_readouts=False,
         )
-        return variables.loc[:, names]
 
     ##########################################################################
     # Get rhs
     ##########################################################################
 
-    def __call__(self, /, time: float, variables: Array) -> Array:
+    def __call__(self, /, time: float, variables: Iterable[float]) -> tuple[float, ...]:
         """Simulation version of get_right_hand_side.
 
         Examples:
@@ -2065,7 +2214,7 @@ class Model:
                 np.array([0.1, 0.2])
 
         Warning: Swaps t and y!
-        This can't get kw-only args, as the integrators call it with pos-only
+        This can't get kw args, as the integrators call it with pos-only
 
         Args:
             time: The current time point.
@@ -2091,8 +2240,7 @@ class Model:
             cache=cache,
         )
 
-        dxdt = cache.dxdt
-        dxdt[:] = 0
+        dxdt = dict.fromkeys(cache.var_names, 0.0)
         for k, stoc in cache.stoich_by_cpds.items():
             for flux, n in stoc.items():
                 dxdt[k] += n * dependent[flux]
@@ -2100,7 +2248,7 @@ class Model:
             for flux, dv in sd.items():
                 n = dv.calculate(dependent)
                 dxdt[k] += n * dependent[flux]
-        return cast(Array, dxdt.to_numpy())
+        return tuple(dxdt[i] for i in cache.var_names)
 
     def get_right_hand_side(
         self,
@@ -2148,3 +2296,52 @@ class Model:
                 n = dv.fn(*(dependent[i] for i in dv.args))
                 dxdt[k] += n * dependent[flux]
         return dxdt
+
+    ##########################################################################
+    # Check units
+    ##########################################################################
+
+    def check_units(self, time_unit: Quantity) -> UnitCheck:
+        """Check unit consistency per differential equation and reaction."""
+        units_per_fn = {}
+        for name, rxn in self._reactions.items():
+            unit_per_arg = {}
+            for arg in rxn.args:
+                if (par := self._parameters.get(arg)) is not None:
+                    unit_per_arg[sympy.Symbol(arg)] = par.unit
+                elif (var := self._variables.get(arg)) is not None:
+                    unit_per_arg[sympy.Symbol(arg)] = var.unit
+                else:
+                    raise NotImplementedError
+
+            symbolic_fn = fn_to_sympy(
+                rxn.fn,
+                origin="unit-checking",
+                model_args=list_of_symbols(rxn.args),
+            )
+            units_per_fn[name] = None
+            if symbolic_fn is None:
+                continue
+            if any(i is None for i in unit_per_arg.values()):
+                continue
+            units_per_fn[name] = symbolic_fn.subs(unit_per_arg)
+
+        check_per_variable = {}
+        for name, var in self._variables.items():
+            check_per_rxn = {}
+
+            if (var_unit := var.unit) is None:
+                break
+
+            for rxn in self.get_stoichiometries_of_variable(name):
+                if (rxn_unit := units_per_fn.get(rxn)) is None:
+                    check_per_rxn[rxn] = None
+                elif unit_of(rxn_unit) == var_unit / time_unit:  # type: ignore
+                    check_per_rxn[rxn] = True
+                else:
+                    check_per_rxn[rxn] = Failure(
+                        expected=unit_of(rxn_unit),
+                        obtained=var_unit / time_unit,  # type: ignore
+                    )
+            check_per_variable[name] = check_per_rxn
+        return UnitCheck(check_per_variable)