mxlpy 0.25.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

mxlpy/nn/_torch.py

@@ -15,8 +15,6 @@ 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
 
@@ -29,11 +27,65 @@ __all__ = [
     "LSTM",
     "LossFn",
     "MLP",
+    "cosine_similarity",
+    "mean_abs_error",
+    "mean_absolute_percentage",
+    "mean_error",
+    "mean_squared_error",
+    "mean_squared_logarithmic",
+    "rms_error",
     "train",
 ]
 
 DefaultDevice = torch.device("cpu")
 
+###############################################################################
+# Loss functions
+###############################################################################
+
+
+type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+
+
+def mean_error(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+    """Calculate mean error."""
+    return torch.mean(pred - true)
+
+
+def mean_squared_error(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+    """Calculate mean squared error."""
+    return torch.mean(torch.square(pred - true))
+
+
+def rms_error(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+    """Calculate root mean square error."""
+    return torch.sqrt(torch.mean(torch.square(pred - true)))
+
+
+def mean_abs_error(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+    """Calculate mean absolute error."""
+    return torch.mean(torch.abs(pred - true))
+
+
+def mean_absolute_percentage(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+    """Calculate mean absolute percentag error."""
+    return 100 * torch.mean(torch.abs((true - pred) / pred))
+
+
+def mean_squared_logarithmic(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+    """Calculate root mean square error between model and data."""
+    return torch.mean(torch.square(torch.log(pred + 1) - torch.log(true + 1)))
+
+
+def cosine_similarity(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+    """Calculate root mean square error between model and data."""
+    return -torch.sum(torch.norm(pred, 2) * torch.norm(true, 2))
+
+
+###############################################################################
+# Training routines
+###############################################################################
+
 
 def train(
     model: nn.Module,
@@ -85,6 +137,11 @@ def train(
     return pd.Series(losses, dtype=float)
 
 
+###############################################################################
+# Actual models
+###############################################################################
+
+
 class MLP(nn.Module):
     """Multilayer Perceptron (MLP) for surrogate modeling and neural posterior estimation.
 

mxlpy/npe/__init__.py

@@ -22,11 +22,16 @@ if TYPE_CHECKING:
     import contextlib
 
     with contextlib.suppress(ImportError):
+        from . import _equinox as equinox
         from . import _keras as keras
         from . import _torch as torch
 else:
     from lazy_import import lazy_module
 
+    equinox = lazy_module(
+        "mxlpy.npe._equinox",
+        error_strings={"module": "equinox", "install_name": "mxlpy[equinox]"},
+    )
     keras = lazy_module(
         "mxlpy.npe._keras",
         error_strings={"module": "keras", "install_name": "mxlpy[tf]"},

mxlpy/npe/_equinox.py

@@ -0,0 +1,344 @@
+"""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 typing import TYPE_CHECKING, Self, cast
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import pandas as pd
+
+from mxlpy.nn._equinox import LSTM, MLP, LossFn, mean_abs_error
+from mxlpy.nn._equinox import train as _train
+from mxlpy.types import AbstractEstimator
+
+if TYPE_CHECKING:
+    import equinox as eqx
+
+__all__ = [
+    "SteadyState",
+    "SteadyStateTrainer",
+    "TimeCourse",
+    "TimeCourseTrainer",
+    "train_steady_state",
+    "train_time_course",
+]
+
+
+@dataclass(kw_only=True)
+class SteadyState(AbstractEstimator):
+    """Estimator for steady state data using PyTorch models."""
+
+    model: eqx.Module
+
+    def predict(self, features: pd.Series | pd.DataFrame) -> pd.DataFrame:
+        """Predict the target values for the given features."""
+        # One has to implement __call__ on eqx.Module, so this should
+        # always exist. Should really be abstract on eqx.Module
+        pred = jax.vmap(self.model)(jnp.array(features))  # type: ignore
+        return pd.DataFrame(pred, columns=self.parameter_names)
+
+
+@dataclass(kw_only=True)
+class TimeCourse(AbstractEstimator):
+    """Estimator for time course data using PyTorch models."""
+
+    model: eqx.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_ = jnp.array(
+            np.swapaxes(
+                features.to_numpy().reshape(
+                    (
+                        len(idx.levels[0]),
+                        len(idx.levels[1]),
+                        len(features.columns),
+                    )
+                ),
+                axis1=0,
+                axis2=1,
+            ),
+        )
+        # One has to implement __call__ on eqx.Module, so this should
+        # always exist. Should really be abstract on eqx.Module
+        pred = jax.vmap(self.model)(features_)  # type: ignore
+        return pd.DataFrame(pred, columns=self.parameter_names)
+
+
+@dataclass
+class SteadyStateTrainer:
+    """Trainer for steady state data using PyTorch models."""
+
+    features: pd.DataFrame
+    targets: pd.DataFrame
+    model: eqx.Module
+    optimizer: optax.GradientTransformation
+    losses: list[pd.Series]
+    loss_fn: LossFn
+    seed: int
+
+    def __init__(
+        self,
+        features: pd.DataFrame,
+        targets: pd.DataFrame,
+        model: eqx.Module | None = None,
+        optimizer: optax.GradientTransformation | None = None,
+        loss_fn: LossFn = mean_abs_error,
+        seed: int = 0,
+    ) -> 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
+            model: Predefined neural network model (None to use default MLP)
+            optimizer: Optimizer class to use for training (default: Adam)
+            device: Device to run the training on (default: DefaultDevice)
+            loss_fn: Loss function
+            seed: seed of random initialisation
+
+        """
+        self.features = features
+        self.targets = targets
+
+        if model is None:
+            n_hidden = max(2 * len(features.columns) * len(targets.columns), 10)
+            n_outputs = len(targets.columns)
+            model = MLP(
+                n_inputs=len(features.columns),
+                neurons_per_layer=[n_hidden, n_hidden, n_outputs],
+                key=jax.random.PRNGKey(seed),
+            )
+        self.model = model
+        self.optimizer = (
+            optax.adamw(learning_rate=0.001) if optimizer is None else optimizer
+        )
+        self.loss_fn = loss_fn
+        self.losses = []
+        self.seed = seed
+
+    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(
+            model=self.model,
+            features=jnp.array(self.features),
+            targets=jnp.array(self.targets),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            loss_fn=self.loss_fn,
+        )
+
+        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) -> SteadyState:
+        """Get the trained estimator."""
+        return SteadyState(
+            model=self.model,
+            parameter_names=list(self.targets.columns),
+        )
+
+
+@dataclass
+class TimeCourseTrainer:
+    """Trainer for time course data using PyTorch models."""
+
+    features: pd.DataFrame
+    targets: pd.DataFrame
+    model: eqx.Module
+    optimizer: optax.GradientTransformation
+    losses: list[pd.Series]
+    loss_fn: LossFn
+
+    def __init__(
+        self,
+        features: pd.DataFrame,
+        targets: pd.DataFrame,
+        model: eqx.Module | None = None,
+        optimizer: optax.GradientTransformation | None = None,
+        loss_fn: LossFn = mean_abs_error,
+    ) -> 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
+            model: Predefined neural network model (None to use default LSTM)
+            optimizer: 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 model is None:
+            model = LSTM(
+                n_inputs=len(features.columns),
+                n_outputs=len(targets.columns),
+                n_hidden=1,
+                key=jnp.array([]),
+            )
+        self.model = model
+        self.optimizer = (
+            optax.adamw(learning_rate=0.001) if optimizer is None else optimizer
+        )
+        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(
+            model=self.model,
+            features=jnp.array(
+                np.swapaxes(
+                    self.features.to_numpy().reshape(
+                        (len(self.targets), -1, len(self.features.columns))
+                    ),
+                    axis1=0,
+                    axis2=1,
+                )
+            ),
+            targets=jnp.array(self.targets.to_numpy()),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            loss_fn=self.loss_fn,
+        )
+
+        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) -> TimeCourse:
+        """Get the trained estimator."""
+        return TimeCourse(
+            model=self.model,
+            parameter_names=list(self.targets.columns),
+        )
+
+
+def train_steady_state(
+    features: pd.DataFrame,
+    targets: pd.DataFrame,
+    epochs: int,
+    batch_size: int | None = None,
+    model: eqx.Module | None = None,
+    optimizer: optax.GradientTransformation | None = None,
+) -> tuple[SteadyState, 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)
+        model: Predefined neural network model (None to use default MLP)
+        optimizer: 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 = SteadyStateTrainer(
+        features=features,
+        targets=targets,
+        model=model,
+        optimizer=optimizer,
+    ).train(epochs=epochs, batch_size=batch_size)
+
+    return trainer.get_estimator(), trainer.get_loss()
+
+
+def train_time_course(
+    features: pd.DataFrame,
+    targets: pd.DataFrame,
+    epochs: int,
+    batch_size: int | None = None,
+    model: eqx.Module | None = None,
+    optimizer: optax.GradientTransformation | None = None,
+) -> tuple[TimeCourse, 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)
+        model: Predefined neural network model (None to use default LSTM)
+        optimizer: 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 = TimeCourseTrainer(
+        features=features,
+        targets=targets,
+        model=model,
+        optimizer=optimizer,
+    ).train(epochs=epochs, batch_size=batch_size)
+
+    return trainer.get_estimator(), trainer.get_loss()

mxlpy/npe/_torch.py

@@ -20,7 +20,8 @@ import torch
 from torch import nn
 from torch.optim.adam import Adam
 
-from mxlpy.nn._torch import LSTM, MLP, DefaultDevice, train
+from mxlpy.nn._torch import LSTM, MLP, DefaultDevice, LossFn, mean_abs_error
+from mxlpy.nn._torch import train as _train
 from mxlpy.types import AbstractEstimator
 
 if TYPE_CHECKING:
@@ -29,10 +30,7 @@ if TYPE_CHECKING:
     from torch.optim.optimizer import ParamsT
 
 
-type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
-
 __all__ = [
-    "LossFn",
     "SteadyState",
     "SteadyStateTrainer",
     "TimeCourse",
@@ -42,20 +40,6 @@ __all__ = [
 ]
 
 
-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 SteadyState(AbstractEstimator):
     """Estimator for steady state data using PyTorch models."""
@@ -115,7 +99,7 @@ class SteadyStateTrainer:
         model: nn.Module | None = None,
         optimizer_cls: Callable[[ParamsT], Adam] = Adam,
         device: torch.device = DefaultDevice,
-        loss_fn: LossFn = _mean_abs,
+        loss_fn: LossFn = mean_abs_error,
     ) -> None:
         """Initialize the trainer with features, targets, and model.
 
@@ -156,7 +140,7 @@ class SteadyStateTrainer:
             batch_size: Size of mini-batches for training (None for full-batch)
 
         """
-        losses = train(
+        losses = _train(
             model=self.model,
             features=self.features.to_numpy(),
             targets=self.targets.to_numpy(),
@@ -203,7 +187,7 @@ class TimeCourseTrainer:
         model: nn.Module | None = None,
         optimizer_cls: Callable[[ParamsT], Adam] = Adam,
         device: torch.device = DefaultDevice,
-        loss_fn: LossFn = _mean_abs,
+        loss_fn: LossFn = mean_abs_error,
     ) -> None:
         """Initialize the trainer with features, targets, and model.
 
@@ -243,7 +227,7 @@ class TimeCourseTrainer:
             batch_size: Size of mini-batches for training (None for full-batch)
 
         """
-        losses = train(
+        losses = _train(
             model=self.model,
             features=np.swapaxes(
                 self.features.to_numpy().reshape(

mxlpy/parallel.py

@@ -27,10 +27,7 @@ from tqdm import tqdm
 if TYPE_CHECKING:
     from collections.abc import Callable, Collection, Hashable
 
-__all__ = [
-    "Cache",
-    "parallelise",
-]
+__all__ = ["Cache", "parallelise", "parallelise_keyless"]
 
 
 def _pickle_name(k: Hashable) -> str:
@@ -173,3 +170,75 @@ def parallelise[K: Hashable, Tin, Tout](
         )  # type: ignore
 
     return results
+
+
+def parallelise_keyless[Tin, Tout](
+    fn: Callable[[Tin], Tout],
+    inputs: Collection[Tin],
+    *,
+    parallel: bool = True,
+    max_workers: int | None = None,
+    timeout: float | None = None,
+    disable_tqdm: bool = False,
+    tqdm_desc: str | None = None,
+) -> list[Tout]:
+    """Execute a function in parallel over a collection of inputs.
+
+    Examples:
+        >>> parallelise(square, [("a", 2), ("b", 3), ("c", 4)])
+        {"a": 4, "b": 9, "c": 16}
+
+    Args:
+        fn: Function to execute in parallel. Takes a single input and returns a result.
+        inputs: Collection of (key, input) tuples to process.
+        cache: Optional cache to store and retrieve results.
+        parallel: Whether to execute in parallel (default: True).
+        max_workers: Maximum number of worker processes (default: None, uses all available CPUs).
+        timeout: Maximum time (in seconds) to wait for each worker to complete (default: None).
+        disable_tqdm: Whether to disable the tqdm progress bar (default: False).
+        tqdm_desc: Description for the tqdm progress bar (default: None).
+
+    Returns:
+        dict[Tin, Tout]: Dictionary mapping inputs to their corresponding outputs.
+
+    """
+    if sys.platform in ["win32", "cygwin"]:
+        parallel = False
+
+    results: list[Tout]
+    if parallel:
+        results = []
+        max_workers = (
+            multiprocessing.cpu_count() if max_workers is None else max_workers
+        )
+
+        with (
+            tqdm(
+                total=len(inputs),
+                disable=disable_tqdm,
+                desc=tqdm_desc,
+            ) as pbar,
+            pebble.ProcessPool(max_workers=max_workers) as pool,
+        ):
+            future = pool.map(fn, inputs, timeout=timeout)
+            it = future.result()
+            while True:
+                try:
+                    value = next(it)
+                    pbar.update(1)
+                    results.append(value)
+                except StopIteration:
+                    break
+                except TimeoutError:
+                    pbar.update(1)
+    else:
+        results = list(
+            tqdm(
+                map(fn, inputs),  # type: ignore
+                total=len(inputs),
+                disable=disable_tqdm,
+                desc=tqdm_desc,
+            )
+        )  # type: ignore
+
+    return results

mxlpy/surrogates/__init__.py

@@ -14,11 +14,16 @@ if TYPE_CHECKING:
     import contextlib
 
     with contextlib.suppress(ImportError):
+        from . import _equinox as equinox
         from . import _keras as keras
         from . import _torch as torch
 else:
     from lazy_import import lazy_module
 
+    equinox = lazy_module(
+        "mxlpy.surrogates._equinox",
+        error_strings={"module": "equinox", "install_name": "mxlpy[equinox]"},
+    )
     keras = lazy_module(
         "mxlpy.surrogates._keras",
         error_strings={"module": "keras", "install_name": "mxlpy[tf]"},