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

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

mxlpy/nn/_torch.py

@@ -8,17 +8,77 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, cast
 
+import numpy as np
+import pandas as pd
 import torch
+import tqdm
 from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-__all__ = ["DefaultDevice", "LSTM", "MLP"]
+    from torch.optim.adam import Adam
+
+    from mxlpy.types import Array
+
+__all__ = ["DefaultDevice", "LSTM", "LossFn", "MLP", "train"]
 
 DefaultDevice = torch.device("cpu")
 
 
+def train(
+    model: nn.Module,
+    features: Array,
+    targets: Array,
+    epochs: int,
+    optimizer: Adam,
+    device: torch.device,
+    batch_size: int | None,
+    loss_fn: LossFn,
+) -> 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.
+
+    """
+    losses = {}
+
+    data = TensorDataset(
+        torch.tensor(features.astype(np.float32), dtype=torch.float32, device=device),
+        torch.tensor(targets.astype(np.float32), dtype=torch.float32, device=device),
+    )
+    data_loader = DataLoader(
+        data,
+        batch_size=len(features) if batch_size is None else batch_size,
+        shuffle=True,
+    )
+
+    for i in tqdm.trange(epochs):
+        epoch_loss = 0
+        for xb, yb in data_loader:
+            optimizer.zero_grad()
+            loss = loss_fn(model(xb), yb)
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item() * xb.size(0)
+        losses[i] = epoch_loss / len(data_loader.dataset)  # type: ignore
+    return pd.Series(losses, dtype=float)
+
+
 class MLP(nn.Module):
     """Multilayer Perceptron (MLP) for surrogate modeling and neural posterior estimation.
 
@@ -65,21 +125,17 @@ class MLP(nn.Module):
         levels = []
         previous_neurons = n_inputs
 
-        for idx, neurons in enumerate(self.layers):
-            if idx == (len(self.layers) - 1):
-                levels.append(nn.Linear(previous_neurons, neurons))
-
-                if self.output_activation:
-                    levels.append(self.output_activation)
-
-            else:
-                levels.append(nn.Linear(previous_neurons, neurons))
-
-                if self.activation:
-                    levels.append(self.activation)
-
+        for neurons in self.layers[:-1]:
+            levels.append(nn.Linear(previous_neurons, neurons))
+            if self.activation:
+                levels.append(self.activation)
             previous_neurons = neurons
 
+        # Output layer
+        levels.append(nn.Linear(previous_neurons, self.layers[-1]))
+        if self.output_activation:
+            levels.append(self.output_activation)
+
         self.net = nn.Sequential(*levels)
 
         for m in self.net.modules():
@@ -103,7 +159,12 @@ class MLP(nn.Module):
 class LSTM(nn.Module):
     """Default LSTM neural network model for time-series approximation."""
 
-    def __init__(self, n_inputs: int, n_outputs: int, n_hidden: int) -> None:
+    def __init__(
+        self,
+        n_inputs: int,
+        n_outputs: int,
+        n_hidden: int,
+    ) -> None:
         """Initializes the neural network model.
 
         Args: