mxlpy 0.18.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

mxlpy/__init__.py

@@ -48,18 +48,15 @@ from . import (
     fns,
     mc,
     mca,
-    npe,
     plot,
     report,
     sbml,
-    surrogates,
 )
 from .integrators import DefaultIntegrator, Scipy
 from .label_map import LabelMapper
 from .linear_label_map import LinearLabelMapper
 from .mc import Cache
 from .model import Model
-from .nn import _torch
 from .scan import (
     steady_state,
     time_course,
@@ -75,6 +72,19 @@ with contextlib.suppress(ImportError):
 if TYPE_CHECKING:
     from mxlpy.types import ArrayLike
 
+    from . import (
+        nn,
+        npe,
+        surrogates,
+    )
+else:
+    from lazy_import import lazy_module
+
+    nn = lazy_module("mxlpy.nn")
+    npe = lazy_module("mxlpy.npe")
+    surrogates = lazy_module("mxlpy.surrogates")
+
+
 __all__ = [
     "Assimulo",
     "Cache",
@@ -87,7 +97,6 @@ __all__ = [
     "Scipy",
     "Simulator",
     "SymbolicModel",
-    "_torch",
     "cartesian_product",
     "distributions",
     "experimental",
@@ -96,6 +105,7 @@ __all__ = [
     "make_protocol",
     "mc",
     "mca",
+    "nn",
     "npe",
     "plot",
     "report",

mxlpy/experimental/diff.py

@@ -181,7 +181,7 @@ def model_diff(m1: Model, m2: Model) -> ModelDiff:
         if (v2 := m2._variables.get(k)) is None:  # noqa: SLF001
             diff.missing_variables.add(k)
         elif v1 != v2:
-            diff.different_variables[k] = (v1, v2)
+            diff.different_variables[k] = (v1, v2)  # type: ignore
 
     for k, v1 in m1._readouts.items():  # noqa: SLF001
         if (v2 := m2._readouts.get(k)) is None:  # noqa: SLF001

mxlpy/integrators/int_assimulo.py

@@ -97,7 +97,11 @@ class Assimulo:
         if steps is None:
             steps = 0
         try:
-            return self.integrator.simulate(t_end, steps)  # type: ignore
+            t, y = self.integrator.simulate(t_end, steps)
+            return (
+                np.atleast_1d(np.array(t, dtype=float)),
+                np.atleast_2d(np.array(y, dtype=float)),
+            )
         except CVodeError:
             return None, None
 
@@ -116,8 +120,11 @@ class Assimulo:
 
         """
         try:
-            t, y = self.integrator.simulate(time_points[-1], 0, time_points)  # type: ignore
-            return np.array(t, dtype=float), np.array(y, dtype=float)
+            t, y = self.integrator.simulate(time_points[-1], 0, time_points)
+            return (
+                np.atleast_1d(np.array(t, dtype=float)),
+                np.atleast_2d(np.array(y, dtype=float)),
+            )
         except CVodeError:
             return None, None
 

mxlpy/integrators/int_scipy.py

@@ -108,10 +108,14 @@ class Scipy:
             rtol=self.rtol,
             method="LSODA",
         )
+
         if res.success:
-            self.t0 = time_points[-1]
-            self.y0 = res.y[:, -1]
-            return np.array(time_points, dtype=float), res.y.T
+            t = np.atleast_1d(np.array(res.t, dtype=float))
+            y = np.atleast_2d(np.array(res.y, dtype=float).T)
+
+            self.t0 = t[-1]
+            self.y0 = y[-1]
+            return t, y
         return None, None
 
     def integrate_to_steady_state(

mxlpy/label_map.py

@@ -29,6 +29,8 @@ __all__ = ["LabelMapper"]
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
 
+    from mxlpy.types import Derived
+
 
 def _total_concentration(*args: float) -> float:
     """Calculate sum of isotopomer concentrations.
@@ -552,7 +554,7 @@ class LabelMapper:
         for name, dp in self.model.derived_parameters.items():
             m.add_derived(name, fn=dp.fn, args=dp.args)
 
-        variables: dict[str, float] = {}
+        variables: dict[str, float | Derived] = {}
         for k, v in self.model.variables.items():
             if (isos := isotopomers.get(k)) is None:
                 variables[k] = v

mxlpy/meta/codegen_latex.py

@@ -694,7 +694,7 @@ class TexExport:
 def _to_tex_export(self: Model) -> TexExport:
     return TexExport(
         parameters=self.parameters,
-        variables=self.variables,
+        variables=self.get_initial_conditions(),  # FIXME: think about this later
         derived=self.derived,
         reactions={k: TexReaction(v.fn, v.args) for k, v in self.reactions.items()},
         stoichiometries={k: v.stoichiometry for k, v in self.reactions.items()},

mxlpy/meta/source_tools.py

@@ -434,5 +434,5 @@ def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr:
                 ctx=ctx.updated(parent_module=imports[module_name.id]),
             )
 
-    msg = f"Onsupported function type {node.func}"
+    msg = f"Unsupported function type {node.func}"
     raise NotImplementedError(msg)

mxlpy/model.py

@@ -250,17 +250,17 @@ class ModelCache:
         stoich_by_cpds: A dictionary mapping compound names to their stoichiometric coefficients.
         dyn_stoich_by_cpds: A dictionary mapping compound names to their dynamic stoichiometric coefficients.
         dxdt: A pandas Series representing the rate of change of variables.
+        initial_conditions: calculated initial conditions
 
     """
 
     var_names: list[str]
-    order: list[str]
+    dyn_order: list[str]
     all_parameter_values: dict[str, float]
-    derived_parameter_names: list[str]
-    derived_variable_names: list[str]
     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]
 
 
 @dataclass(slots=True)
@@ -276,17 +276,19 @@ class Model:
         _reactions: Dictionary of reactions in the model.
         _surrogates: Dictionary of surrogate models.
         _cache: Cache for storing model-related data structures.
+        _data: Named references to data sets
 
     """
 
     _ids: dict[str, str] = field(default_factory=dict)
-    _variables: dict[str, float] = field(default_factory=dict)
+    _variables: dict[str, float | Derived] = field(default_factory=dict)
     _parameters: dict[str, float] = field(default_factory=dict)
     _derived: dict[str, Derived] = field(default_factory=dict)
     _readouts: dict[str, Readout] = field(default_factory=dict)
     _reactions: dict[str, Reaction] = field(default_factory=dict)
     _surrogates: dict[str, AbstractSurrogate] = field(default_factory=dict)
     _cache: ModelCache | None = None
+    _data: dict[str, pd.Series | pd.DataFrame] = field(default_factory=dict)
 
     ###########################################################################
     # Cache
@@ -317,9 +319,17 @@ class Model:
                 raise ArityMismatchError(name, el.fn, el.args)
 
         # Sort derived & reactions
-        to_sort = self._derived | self._reactions | self._surrogates
+        to_sort = (
+            self._derived
+            | self._reactions
+            | self._surrogates
+            | {k: v for k, v in self._variables.items() if isinstance(v, Derived)}
+        )
         order = _sort_dependencies(
-            available=set(self._parameters) | set(self._variables) | {"time"},
+            available=all_parameter_names
+            | {k for k, v in self._variables.items() if not isinstance(v, Derived)}
+            | set(self._data)
+            | {"time"},
             elements=[
                 Dependency(name=k, required=set(v.args), provided={k})
                 if not isinstance(v, AbstractSurrogate)
@@ -328,35 +338,42 @@ class Model:
             ],
         )
 
-        # Split derived into parameters and variables
-        # for user convenience
-        derived_variable_names: list[str] = []
-        derived_parameter_names: list[str] = []
+        # 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: v for k, v in self._variables.items() if not isinstance(v, Derived)}
+            | {"time": 0.0}
+        )
+        for name in order:
+            to_sort[name].calculate_inpl(name, dependent)
+
+        # Split derived into static and dynamic variables
+        static_order = []
+        dyn_order = []
         for name in order:
             if name in self._reactions or name in self._surrogates:
-                continue
-            derived = self._derived[name]
-            if all(i in all_parameter_names for i in derived.args):
-                all_parameter_names.add(name)
-                derived_parameter_names.append(name)
-                all_parameter_values[name] = float(
-                    derived.fn(*(all_parameter_values[i] for i in derived.args))
-                )
+                dyn_order.append(name)
+            elif name in self._variables:
+                static_order.append(name)
             else:
-                derived_variable_names.append(name)
+                derived = self._derived[name]
+                if all(i in all_parameter_names for i in derived.args):
+                    static_order.append(name)
+                    all_parameter_names.add(name)
+                else:
+                    dyn_order.append(name)
 
+        # Calculate dynamic and static stochiometries
         stoich_by_compounds: dict[str, dict[str, float]] = {}
         dyn_stoich_by_compounds: dict[str, dict[str, Derived]] = {}
-
         for rxn_name, rxn in self._reactions.items():
             for cpd_name, factor in rxn.stoichiometry.items():
                 d_static = stoich_by_compounds.setdefault(cpd_name, {})
-
                 if isinstance(factor, Derived):
                     if all(i in all_parameter_names for i in factor.args):
-                        d_static[rxn_name] = float(
-                            factor.fn(*(all_parameter_values[i] for i in factor.args))
-                        )
+                        d_static[rxn_name] = factor.calculate(dependent)
                     else:
                         dyn_stoich_by_compounds.setdefault(cpd_name, {})[rxn_name] = (
                             factor
@@ -367,20 +384,40 @@ class Model:
         for surrogate in self._surrogates.values():
             for rxn_name, rxn in surrogate.stoichiometries.items():
                 for cpd_name, factor in rxn.items():
-                    stoich_by_compounds.setdefault(cpd_name, {})[rxn_name] = factor
+                    d_static = stoich_by_compounds.setdefault(cpd_name, {})
+                    if isinstance(factor, Derived):
+                        if all(i in all_parameter_names for i in factor.args):
+                            d_static[rxn_name] = factor.calculate(dependent)
+                        else:
+                            dyn_stoich_by_compounds.setdefault(cpd_name, {})[
+                                rxn_name
+                            ] = factor
+                    else:
+                        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: v for k, v in self._variables.items() if not isinstance(v, Derived)
+        }
+        for name in static_order:
+            if name in self._variables:
+                initial_conditions[name] = cast(float, dependent[name])
+            elif 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(
             var_names=var_names,
-            order=order,
+            dyn_order=dyn_order,
             all_parameter_values=all_parameter_values,
             stoich_by_cpds=stoich_by_compounds,
             dyn_stoich_by_cpds=dyn_stoich_by_compounds,
-            derived_variable_names=derived_variable_names,
-            derived_parameter_names=derived_parameter_names,
             dxdt=dxdt,
+            initial_conditions=initial_conditions,
         )
         return self._cache
 
@@ -662,7 +699,7 @@ class Model:
     ##########################################################################
 
     @property
-    def variables(self) -> dict[str, float]:
+    def variables(self) -> dict[str, float | Derived]:
         """Returns a copy of the variables dictionary.
 
         Examples:
@@ -679,7 +716,7 @@ class Model:
         return self._variables.copy()
 
     @_invalidate_cache
-    def add_variable(self, name: str, initial_condition: float) -> Self:
+    def add_variable(self, name: str, initial_condition: float | Derived) -> Self:
         """Adds a variable to the model with the given name and initial condition.
 
         Examples:
@@ -697,7 +734,7 @@ class Model:
         self._variables[name] = initial_condition
         return self
 
-    def add_variables(self, variables: dict[str, float]) -> Self:
+    def add_variables(self, variables: Mapping[str, float | Derived]) -> Self:
         """Adds multiple variables to the model with their initial conditions.
 
         Examples:
@@ -751,7 +788,7 @@ class Model:
         return self
 
     @_invalidate_cache
-    def update_variable(self, name: str, initial_condition: float) -> Self:
+    def update_variable(self, name: str, initial_condition: float | Derived) -> Self:
         """Updates the value of a variable in the model.
 
         Examples:
@@ -771,7 +808,7 @@ class Model:
         self._variables[name] = initial_condition
         return self
 
-    def update_variables(self, variables: dict[str, float]) -> Self:
+    def update_variables(self, variables: Mapping[str, float | Derived]) -> Self:
         """Updates multiple variables in the model.
 
         Examples:
@@ -812,7 +849,9 @@ class Model:
             initial_conditions: A dictionary where the keys are variable names and the values are their initial conditions.
 
         """
-        return self._variables
+        if (cache := self._cache) is None:
+            cache = self._create_cache()
+        return cache.initial_conditions
 
     def make_variable_static(self, name: str, value: float | None = None) -> Self:
         """Converts a variable to a static parameter.
@@ -833,9 +872,12 @@ class Model:
             Self: The instance of the class for method chaining.
 
         """
-        value = self._variables[name] if value is None else value
+        value_or_derived = self._variables[name] if value is None else value
         self.remove_variable(name)
-        self.add_parameter(name, value)
+        if isinstance(value_or_derived, Derived):
+            self.add_derived(name, value_or_derived.fn, args=value_or_derived.args)
+        else:
+            self.add_parameter(name, value_or_derived)
 
         # Remove from stoichiometries
         for reaction in self._reactions.values():
@@ -886,7 +928,8 @@ class Model:
         if (cache := self._cache) is None:
             cache = self._create_cache()
         derived = self._derived
-        return {k: derived[k] for k in cache.derived_variable_names}
+
+        return {k: v for k, v in derived.items() if k not in cache.all_parameter_values}
 
     @property
     def derived_parameters(self) -> dict[str, Derived]:
@@ -905,7 +948,7 @@ class Model:
         if (cache := self._cache) is None:
             cache = self._create_cache()
         derived = self._derived
-        return {k: derived[k] for k in cache.derived_parameter_names}
+        return {k: v for k, v in derived.items() if k in cache.all_parameter_values}
 
     @_invalidate_cache
     def add_derived(
@@ -1049,6 +1092,33 @@ class Model:
                 )
         return pd.DataFrame(stoich_by_cpds).T.fillna(0)
 
+    def get_stoichiometries_of_variable(
+        self,
+        variable: str,
+        variables: dict[str, float] | None = None,
+        time: float = 0.0,
+    ) -> dict[str, float]:
+        """Retrieve the stoichiometry of a specific variable.
+
+        Examples:
+            >>> model.get_stoichiometries_of_variable("x1")
+                {"v1": -1, "v2": 1}
+
+        Args:
+            variable: The name of the variable for which to retrieve the stoichiometry.
+            variables: A dictionary of variable names and their values.
+            time: The time point at which to evaluate the stoichiometry.
+
+        """
+        if (cache := self._cache) is None:
+            cache = self._create_cache()
+        args = self.get_dependent(variables=variables, time=time)
+
+        stoich = copy.deepcopy(cache.stoich_by_cpds[variable])
+        for rxn, derived in cache.dyn_stoich_by_cpds[variable].items():
+            stoich[rxn] = float(derived.fn(*(args[i] for i in derived.args)))
+        return stoich
+
     @_invalidate_cache
     def add_reaction(
         self,
@@ -1250,7 +1320,7 @@ class Model:
         surrogate: AbstractSurrogate,
         args: list[str] | None = None,
         outputs: list[str] | None = None,
-        stoichiometries: dict[str, dict[str, float]] | None = None,
+        stoichiometries: dict[str, dict[str, float | Derived]] | None = None,
     ) -> Self:
         """Adds a surrogate model to the current instance.
 
@@ -1285,7 +1355,7 @@ class Model:
         name: str,
         surrogate: AbstractSurrogate | None = None,
         args: list[str] | None = None,
-        stoichiometries: dict[str, dict[str, float]] | None = None,
+        stoichiometries: dict[str, dict[str, float | Derived]] | None = None,
     ) -> Self:
         """Update a surrogate model in the model.
 
@@ -1337,6 +1407,27 @@ class Model:
             names.extend(i.stoichiometries)
         return names
 
+    ##########################################################################
+    # Datasets
+    ##########################################################################
+
+    def add_data(self, name: str, data: pd.Series | pd.DataFrame) -> Self:
+        """Add named data set to model."""
+        self._insert_id(name=name, ctx="data")
+        self._data[name] = data
+        return self
+
+    def update_data(self, name: str, data: pd.Series | pd.DataFrame) -> Self:
+        """Update named data set."""
+        self._data[name] = data
+        return self
+
+    def remove_data(self, name: str) -> Self:
+        """Remove data set from model."""
+        self._remove_id(name=name)
+        self._data.pop(name)
+        return self
+
     ##########################################################################
     # Get dependent values. This includes
     # - derived parameters
@@ -1371,14 +1462,17 @@ class Model:
                 with their respective names as keys and their calculated values as values.
 
         """
-        args: dict[str, float] = cache.all_parameter_values | variables
+        args = cache.all_parameter_values | variables | self._data
         args["time"] = time
 
         containers = self._derived | self._reactions | self._surrogates
-        for name in cache.order:
+        for name in cache.dyn_order:
             containers[name].calculate_inpl(name, args)
 
-        return args
+        for k in self._data:
+            args.pop(k)
+
+        return cast(dict[str, float], args)
 
     def get_dependent(
         self,
@@ -1454,29 +1548,16 @@ class Model:
             derived variables, and optionally readout variables, with time as an additional column.
 
         """
-        if (cache := self._cache) is None:
-            cache = self._create_cache()
-
-        pars_df = pd.DataFrame(
-            np.full(
-                (len(variables), len(cache.all_parameter_values)),
-                np.fromiter(cache.all_parameter_values.values(), dtype=float),
-            ),
-            index=variables.index,
-            columns=list(cache.all_parameter_values),
-        )
-
-        args = pd.concat((variables, pars_df), axis=1)
-        args["time"] = args.index
-
-        containers = self._derived | self._reactions | self._surrogates
-        for name in cache.order:
-            containers[name].calculate_inpl_time_course(name, args)
+        args = {
+            time: self.get_dependent(
+                variables=values.to_dict(),
+                time=cast(float, time),
+                include_readouts=include_readouts,
+            )
+            for time, values in variables.iterrows()
+        }
 
-        if include_readouts:
-            for name, ro in self._readouts.items():
-                args[name] = ro.fn(*args.loc[:, ro.args].to_numpy().T)
-        return args
+        return pd.DataFrame(args, dtype=float).T
 
     ##########################################################################
     # Get args
@@ -1569,28 +1650,6 @@ class Model:
     # Get fluxes
     ##########################################################################
 
-    def _get_fluxes(self, args: dict[str, float]) -> dict[str, float]:
-        """Calculate the fluxes for the given arguments.
-
-        Examples:
-            >>> model._get_fluxes({"x1": 1.0, "x2": 2.0, "k1": 0.1, "time": 0.0})
-                {"r1": 0.1, "r2": 0.2}
-
-        Args:
-            args (dict[str, float]): A dictionary where the keys are argument names and the values are their corresponding float values.
-
-        Returns:
-            dict[str, float]: A dictionary where the keys are reaction names and the values are the calculated fluxes.
-
-        """
-        fluxes: dict[str, float] = {}
-        for name, rxn in self._reactions.items():
-            fluxes[name] = cast(float, rxn.fn(*(args[arg] for arg in rxn.args)))
-
-        for surrogate in self._surrogates.values():
-            fluxes |= surrogate.predict(np.array([args[arg] for arg in surrogate.args]))
-        return fluxes
-
     def get_fluxes(
         self,
         variables: dict[str, float] | None = None,

mxlpy/nn/__init__.py

@@ -1,10 +1,29 @@
 """Collection of neural network architectures."""
 
-import contextlib
+from __future__ import annotations
 
-__all__ = ["tensorflow", "torch"]
+from typing import TYPE_CHECKING
 
-with contextlib.suppress(ImportError):
-    from . import _torch as torch
+if TYPE_CHECKING:
+    import contextlib
 
-from . import _tensorflow as tensorflow
+    with contextlib.suppress(ImportError):
+        from . import _keras as keras
+        from . import _torch as torch
+else:
+    from lazy_import import lazy_module
+
+    keras = lazy_module(
+        "mxlpy.nn._keras",
+        error_strings={"module": "keras", "install_name": "mxlpy[tf]"},
+    )
+    torch = lazy_module(
+        "mxlpy.nn._torch",
+        error_strings={"module": "torch", "install_name": "mxlpy[torch]"},
+    )
+
+
+__all__ = [
+    "keras",
+    "torch",
+]

mxlpy/nn/_keras.py

@@ -0,0 +1,85 @@
+from typing import cast
+
+import keras
+import pandas as pd
+from tqdm.keras import TqdmCallback
+
+from mxlpy.types import Array
+
+__all__ = ["LSTM", "MLP", "train"]
+
+
+def train(
+    model: keras.Model,
+    features: pd.DataFrame | Array,
+    targets: pd.DataFrame | Array,
+    epochs: int,
+    batch_size: int | None,
+) -> pd.Series:
+    """Train the neural network using mini-batch gradient descent.
+
+    Args:
+        model: Neural network model to train.
+        features: Input features as a tensor.
+        targets: Target values as a tensor.
+        epochs: Number of training epochs.
+        optimizer: Optimizer for training.
+        device: torch device
+        batch_size: Size of mini-batches for training.
+        loss_fn: Loss function
+
+    Returns:
+        pd.Series: Series containing the training loss history.
+
+    """
+    history = model.fit(
+        features,
+        targets,
+        batch_size=batch_size,
+        epochs=epochs,
+        verbose=cast(str, 0),
+        callbacks=[TqdmCallback()],
+    )
+    return pd.Series(history.history["loss"])
+
+
+def MLP(  # noqa: N802
+    n_inputs: int,
+    neurons_per_layer: list[int],
+    activation: None = None,
+    output_activation: None = None,
+) -> keras.Sequential:
+    """Multilayer Perceptron (MLP) for surrogate modeling and neural posterior estimation.
+
+    Methods:
+        forward: Forward pass through the neural network.
+
+    """
+    model = keras.Sequential([keras.Input(shape=(n_inputs,))])
+    for neurons in neurons_per_layer[:-1]:
+        model.add(keras.layers.Dense(neurons, activation=activation))
+    model.add(keras.layers.Dense(neurons_per_layer[-1], activation=output_activation))
+    return model
+
+
+def LSTM(  # noqa: N802
+    n_inputs: int,
+    n_outputs: int,
+    n_hidden: int,
+) -> keras.Sequential:
+    """Long Short-Term Memory (LSTM) network for time series modeling.
+
+    Methods:
+        forward: Forward pass through the neural network.
+
+    """
+    model = keras.Sequential(
+        [
+            keras.Input(
+                shape=(n_inputs),
+            )
+        ]
+    )
+    model.add(keras.layers.LSTM(n_hidden))
+    model.add(keras.layers.Dense(n_outputs))
+    return model