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

mxlpy/model.py

@@ -18,6 +18,7 @@ import pandas as pd
 
 from mxlpy import fns
 from mxlpy.types import (
+    AbstractSurrogate,
     Array,
     Derived,
     Reaction,
@@ -27,6 +28,7 @@ from mxlpy.types import (
 __all__ = [
     "ArityMismatchError",
     "CircularDependencyError",
+    "Dependency",
     "MissingDependenciesError",
     "Model",
     "ModelCache",
@@ -36,7 +38,16 @@ if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
     from inspect import FullArgSpec
 
-    from mxlpy.types import AbstractSurrogate, Callable, Param, RateFn, RetType
+    from mxlpy.types import Callable, Param, RateFn, RetType
+
+
+@dataclass
+class Dependency:
+    """Container class for building dependency tree."""
+
+    name: str
+    required: set[str]
+    provided: set[str]
 
 
 class MissingDependenciesError(Exception):
@@ -145,30 +156,33 @@ def _invalidate_cache(method: Callable[Param, RetType]) -> Callable[Param, RetTy
 
 def _check_if_is_sortable(
     available: set[str],
-    elements: list[tuple[str, set[str]]],
+    elements: list[Dependency],
 ) -> None:
     all_available = available.copy()
-    for name, _ in elements:
-        all_available.add(name)
+    for dependency in elements:
+        all_available.update(dependency.provided)
 
     # Check if it can be sorted in the first place
     not_solvable = {}
-    for name, args in elements:
-        if not args.issubset(all_available):
-            not_solvable[name] = sorted(args.difference(all_available))
+    for dependency in elements:
+        if not dependency.required.issubset(all_available):
+            not_solvable[dependency.name] = sorted(
+                dependency.required.difference(all_available)
+            )
 
     if not_solvable:
         raise MissingDependenciesError(not_solvable=not_solvable)
 
 
 def _sort_dependencies(
-    available: set[str], elements: list[tuple[str, set[str]]]
+    available: set[str],
+    elements: list[Dependency],
 ) -> list[str]:
     """Sort model elements topologically based on their dependencies.
 
     Args:
         available: Set of available component names
-        elements: List of (name, dependencies) tuples to sort
+        elements: List of (name, dependencies, supplier) tuples to sort
 
     Returns:
         List of element names in dependency order
@@ -184,26 +198,27 @@ def _sort_dependencies(
     order = []
     # FIXME: what is the worst case here?
     max_iterations = len(elements) ** 2
-    queue: SimpleQueue[tuple[str, set[str]]] = SimpleQueue()
-    for k, v in elements:
-        queue.put((k, v))
+    queue: SimpleQueue[Dependency] = SimpleQueue()
+    for dependency in elements:
+        queue.put(dependency)
 
     last_name = None
     i = 0
     while True:
         try:
-            new, args = queue.get_nowait()
+            dependency = queue.get_nowait()
         except Empty:
             break
-        if args.issubset(available):
-            available.add(new)
-            order.append(new)
+        if dependency.required.issubset(available):
+            available.update(dependency.provided)
+            order.append(dependency.name)
+
         else:
-            if last_name == new:
-                order.append(new)
+            if last_name == dependency.name:
+                order.append(last_name)
                 break
-            queue.put((new, args))
-            last_name = new
+            queue.put(dependency)
+            last_name = dependency.name
         i += 1
 
         # Failure case
@@ -211,11 +226,13 @@ def _sort_dependencies(
             unsorted = []
             while True:
                 try:
-                    unsorted.append(queue.get_nowait()[0])
+                    unsorted.append(queue.get_nowait().name)
                 except Empty:
                     break
 
-            mod_to_args: dict[str, set[str]] = dict(elements)
+            mod_to_args: dict[str, set[str]] = {
+                dependency.name: dependency.required for dependency in elements
+            }
             missing = {k: mod_to_args[k].difference(available) for k in unsorted}
             raise CircularDependencyError(missing=missing)
     return order
@@ -303,7 +320,12 @@ class Model:
         to_sort = self._derived | self._reactions | self._surrogates
         order = _sort_dependencies(
             available=set(self._parameters) | set(self._variables) | {"time"},
-            elements=[(k, set(v.args)) for k, v in to_sort.items()],
+            elements=[
+                Dependency(name=k, required=set(v.args), provided={k})
+                if not isinstance(v, AbstractSurrogate)
+                else Dependency(name=k, required=set(v.args), provided=set(v.outputs))
+                for k, v in to_sort.items()
+            ],
         )
 
         # Split derived into parameters and variables
@@ -1227,6 +1249,7 @@ class Model:
         name: str,
         surrogate: AbstractSurrogate,
         args: list[str] | None = None,
+        outputs: list[str] | None = None,
         stoichiometries: dict[str, dict[str, float]] | None = None,
     ) -> Self:
         """Adds a surrogate model to the current instance.
@@ -1237,7 +1260,8 @@ class Model:
         Args:
             name (str): The name of the surrogate model.
             surrogate (AbstractSurrogate): The surrogate model instance to be added.
-            args: A list of arguments for the surrogate model.
+            args: Names of the values passed for the surrogate model.
+            outputs: Names of values produced by the surrogate model.
             stoichiometries: A dictionary mapping reaction names to stoichiometries.
 
         Returns:
@@ -1248,6 +1272,8 @@ class Model:
 
         if args is not None:
             surrogate.args = args
+        if outputs is not None:
+            surrogate.outputs = outputs
         if stoichiometries is not None:
             surrogate.stoichiometries = stoichiometries
 

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(
+    aprox: 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:
+        aprox: 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(aprox(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.
 

mxlpy/npe/__init__.py

@@ -0,0 +1,38 @@
+"""Neural Process Estimation (NPE) module.
+
+This module provides classes and functions for estimating metabolic processes using
+neural networks. It includes functionality for both steady-state and time-course data.
+
+Classes:
+    TorchSteadyState: Class for steady-state neural network estimation.
+    TorchSteadyStateTrainer: Class for training steady-state neural networks.
+    TorchTimeCourse: Class for time-course neural network estimation.
+    TorchTimeCourseTrainer: Class for training time-course neural networks.
+
+Functions:
+    train_torch_steady_state: Train a PyTorch steady-state neural network.
+    train_torch_time_course: Train a PyTorch time-course neural network.
+"""
+
+from __future__ import annotations
+
+import contextlib
+
+with contextlib.suppress(ImportError):
+    from ._torch import (
+        TorchSteadyState,
+        TorchSteadyStateTrainer,
+        TorchTimeCourse,
+        TorchTimeCourseTrainer,
+        train_torch_steady_state,
+        train_torch_time_course,
+    )
+
+__all__ = [
+    "TorchSteadyState",
+    "TorchSteadyStateTrainer",
+    "TorchTimeCourse",
+    "TorchTimeCourseTrainer",
+    "train_torch_steady_state",
+    "train_torch_time_course",
+]

mxlpy/npe/_torch.py

@@ -0,0 +1,365 @@
+"""Neural Network Parameter Estimation (NPE) Module.
+
+This module provides classes and functions for training neural network models to estimate
+parameters in metabolic models. It includes functionality for both steady-state and
+time-series data.
+
+Functions:
+    train_torch_surrogate: Train a PyTorch surrogate model
+    train_torch_time_course_estimator: Train a PyTorch time course estimator
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Self, cast
+
+import numpy as np
+import pandas as pd
+import torch
+from torch import nn
+from torch.optim.adam import Adam
+
+from mxlpy.nn._torch import LSTM, MLP, DefaultDevice, train
+from mxlpy.parallel import Cache
+from mxlpy.types import AbstractEstimator
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.optim.optimizer import ParamsT
+
+DefaultCache = Cache(Path(".cache"))
+
+type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+
+__all__ = [
+    "DefaultCache",
+    "LossFn",
+    "TorchSteadyState",
+    "TorchSteadyStateTrainer",
+    "TorchTimeCourse",
+    "TorchTimeCourseTrainer",
+    "train_torch_steady_state",
+    "train_torch_time_course",
+]
+
+
+def _mean_abs(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """Standard loss for surrogates.
+
+    Args:
+        x: Predictions of a model.
+        y: Targets.
+
+    Returns:
+        torch.Tensor: loss.
+
+    """
+    return torch.mean(torch.abs(x - y))
+
+
+@dataclass(kw_only=True)
+class TorchSteadyState(AbstractEstimator):
+    """Estimator for steady state data using PyTorch models."""
+
+    model: torch.nn.Module
+
+    def predict(self, features: pd.Series | pd.DataFrame) -> pd.DataFrame:
+        """Predict the target values for the given features."""
+        with torch.no_grad():
+            pred = self.model(torch.tensor(features.to_numpy(), dtype=torch.float32))
+            return pd.DataFrame(pred, columns=self.parameter_names)
+
+
+@dataclass(kw_only=True)
+class TorchTimeCourse(AbstractEstimator):
+    """Estimator for time course data using PyTorch models."""
+
+    model: torch.nn.Module
+
+    def predict(self, features: pd.Series | pd.DataFrame) -> pd.DataFrame:
+        """Predict the target values for the given features."""
+        idx = cast(pd.MultiIndex, features.index)
+        features_ = torch.Tensor(
+            np.swapaxes(
+                features.to_numpy().reshape(
+                    (
+                        len(idx.levels[0]),
+                        len(idx.levels[1]),
+                        len(features.columns),
+                    )
+                ),
+                axis1=0,
+                axis2=1,
+            ),
+        )
+        with torch.no_grad():
+            pred = self.model(features_)
+            return pd.DataFrame(pred, columns=self.parameter_names)
+
+
+@dataclass
+class TorchSteadyStateTrainer:
+    """Trainer for steady state data using PyTorch models."""
+
+    features: pd.DataFrame
+    targets: pd.DataFrame
+    approximator: nn.Module
+    optimimzer: Adam
+    device: torch.device
+    losses: list[pd.Series]
+    loss_fn: LossFn
+
+    def __init__(
+        self,
+        features: pd.DataFrame,
+        targets: pd.DataFrame,
+        approximator: nn.Module | None = None,
+        optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
+        device: torch.device = DefaultDevice,
+        loss_fn: LossFn = _mean_abs,
+    ) -> None:
+        """Initialize the trainer with features, targets, and model.
+
+        Args:
+            features: DataFrame containing the input features for training
+            targets: DataFrame containing the target values for training
+            approximator: Predefined neural network model (None to use default MLP)
+            optimimzer_cls: Optimizer class to use for training (default: Adam)
+            device: Device to run the training on (default: DefaultDevice)
+            loss_fn: Loss function
+
+        """
+        self.features = features
+        self.targets = targets
+
+        if approximator is None:
+            n_hidden = max(2 * len(features.columns) * len(targets.columns), 10)
+            n_outputs = len(targets.columns)
+            approximator = MLP(
+                n_inputs=len(features.columns),
+                neurons_per_layer=[n_hidden, n_hidden, n_outputs],
+            )
+        self.approximator = approximator.to(device)
+        self.optimizer = optimimzer_cls(approximator.parameters())
+        self.device = device
+        self.loss_fn = loss_fn
+        self.losses = []
+
+    def train(
+        self,
+        epochs: int,
+        batch_size: int | None = None,
+    ) -> Self:
+        """Train the model using the provided features and targets.
+
+        Args:
+            epochs: Number of training epochs
+            batch_size: Size of mini-batches for training (None for full-batch)
+
+        """
+        losses = train(
+            aprox=self.approximator,
+            features=self.features.to_numpy(),
+            targets=self.targets.to_numpy(),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            loss_fn=self.loss_fn,
+            device=self.device,
+        )
+
+        if len(self.losses) > 0:
+            losses.index += self.losses[-1].index[-1]
+        self.losses.append(losses)
+        return self
+
+    def get_loss(self) -> pd.Series:
+        """Get the loss history of the training process."""
+        return pd.concat(self.losses)
+
+    def get_estimator(self) -> TorchSteadyState:
+        """Get the trained estimator."""
+        return TorchSteadyState(
+            model=self.approximator,
+            parameter_names=list(self.targets.columns),
+        )
+
+
+@dataclass
+class TorchTimeCourseTrainer:
+    """Trainer for time course data using PyTorch models."""
+
+    features: pd.DataFrame
+    targets: pd.DataFrame
+    approximator: nn.Module
+    optimimzer: Adam
+    device: torch.device
+    losses: list[pd.Series]
+    loss_fn: LossFn
+
+    def __init__(
+        self,
+        features: pd.DataFrame,
+        targets: pd.DataFrame,
+        approximator: nn.Module | None = None,
+        optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
+        device: torch.device = DefaultDevice,
+        loss_fn: LossFn = _mean_abs,
+    ) -> None:
+        """Initialize the trainer with features, targets, and model.
+
+        Args:
+            features: DataFrame containing the input features for training
+            targets: DataFrame containing the target values for training
+            approximator: Predefined neural network model (None to use default LSTM)
+            optimimzer_cls: Optimizer class to use for training (default: Adam)
+            device: Device to run the training on (default: DefaultDevice)
+            loss_fn: Loss function
+
+        """
+        self.features = features
+        self.targets = targets
+
+        if approximator is None:
+            approximator = LSTM(
+                n_inputs=len(features.columns),
+                n_outputs=len(targets.columns),
+                n_hidden=1,
+            ).to(device)
+        self.approximator = approximator.to(device)
+        self.optimizer = optimimzer_cls(approximator.parameters())
+        self.device = device
+        self.loss_fn = loss_fn
+        self.losses = []
+
+    def train(
+        self,
+        epochs: int,
+        batch_size: int | None = None,
+    ) -> Self:
+        """Train the model using the provided features and targets.
+
+        Args:
+            epochs: Number of training epochs
+            batch_size: Size of mini-batches for training (None for full-batch)
+
+        """
+        losses = train(
+            aprox=self.approximator,
+            features=np.swapaxes(
+                self.features.to_numpy().reshape(
+                    (len(self.targets), -1, len(self.features.columns))
+                ),
+                axis1=0,
+                axis2=1,
+            ),
+            targets=self.targets.to_numpy(),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            loss_fn=self.loss_fn,
+            device=self.device,
+        )
+
+        if len(self.losses) > 0:
+            losses.index += self.losses[-1].index[-1]
+        self.losses.append(losses)
+        return self
+
+    def get_loss(self) -> pd.Series:
+        """Get the loss history of the training process."""
+        return pd.concat(self.losses)
+
+    def get_estimator(self) -> TorchTimeCourse:
+        """Get the trained estimator."""
+        return TorchTimeCourse(
+            model=self.approximator,
+            parameter_names=list(self.targets.columns),
+        )
+
+
+def train_torch_steady_state(
+    features: pd.DataFrame,
+    targets: pd.DataFrame,
+    epochs: int,
+    batch_size: int | None = None,
+    approximator: nn.Module | None = None,
+    optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
+    device: torch.device = DefaultDevice,
+) -> tuple[TorchSteadyState, pd.Series]:
+    """Train a PyTorch steady state estimator.
+
+    This function trains a neural network model to estimate steady state data
+    using the provided features and targets. It supports both full-batch and
+    mini-batch training.
+
+    Examples:
+        >>> train_torch_ss_estimator(features, targets, epochs=100)
+
+    Args:
+        features: DataFrame containing the input features for training
+        targets: DataFrame containing the target values for training
+        epochs: Number of training epochs
+        batch_size: Size of mini-batches for training (None for full-batch)
+        approximator: Predefined neural network model (None to use default MLP)
+        optimimzer_cls: Optimizer class to use for training (default: Adam)
+        device: Device to run the training on (default: DefaultDevice)
+
+    Returns:
+        tuple[TorchTimeSeriesEstimator, pd.Series]: Trained estimator and loss history
+
+    """
+    trainer = TorchSteadyStateTrainer(
+        features=features,
+        targets=targets,
+        approximator=approximator,
+        optimimzer_cls=optimimzer_cls,
+        device=device,
+    ).train(epochs=epochs, batch_size=batch_size)
+
+    return trainer.get_estimator(), trainer.get_loss()
+
+
+def train_torch_time_course(
+    features: pd.DataFrame,
+    targets: pd.DataFrame,
+    epochs: int,
+    batch_size: int | None = None,
+    approximator: nn.Module | None = None,
+    optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
+    device: torch.device = DefaultDevice,
+) -> tuple[TorchTimeCourse, pd.Series]:
+    """Train a PyTorch time course estimator.
+
+    This function trains a neural network model to estimate time course data
+    using the provided features and targets. It supports both full-batch and
+    mini-batch training.
+
+    Examples:
+        >>> train_torch_time_course_estimator(features, targets, epochs=100)
+
+    Args:
+        features: DataFrame containing the input features for training
+        targets: DataFrame containing the target values for training
+        epochs: Number of training epochs
+        batch_size: Size of mini-batches for training (None for full-batch)
+        approximator: Predefined neural network model (None to use default LSTM)
+        optimimzer_cls: Optimizer class to use for training (default: Adam)
+        device: Device to run the training on (default: DefaultDevice)
+
+    Returns:
+        tuple[TorchTimeSeriesEstimator, pd.Series]: Trained estimator and loss history
+
+    """
+    trainer = TorchTimeCourseTrainer(
+        features=features,
+        targets=targets,
+        approximator=approximator,
+        optimimzer_cls=optimimzer_cls,
+        device=device,
+    ).train(epochs=epochs, batch_size=batch_size)
+
+    return trainer.get_estimator(), trainer.get_loss()