modelbase2 0.1.79__py3-none-any.whl → 0.3.0__py3-none-any.whl

modelbase2/nnarchitectures.py

@@ -0,0 +1,128 @@
+"""Neural network architectures.
+
+This module provides implementations of neural network architectures used for mechanistic learning.
+
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+__all__ = ["DefaultDevice", "LSTMnn", "MLP"]
+
+DefaultDevice = torch.device("cpu")
+
+
+class MLP(nn.Module):
+    """Multilayer Perceptron (MLP) for surrogate modeling and neural posterior estimation.
+
+    Attributes:
+        net: Sequential neural network model.
+
+    Methods:
+        forward: Forward pass through the neural network.
+
+    """
+
+    def __init__(
+        self,
+        n_inputs: int,
+        layers: list[int],
+        activation: Callable | None = nn.ReLU(),
+        output_activation: Callable | None = None,
+    ) -> None:
+        """Initializes the MLP with the given number of inputs and list of (hidden) layers.
+
+        Args:
+            n_inputs (int): The number of input features.
+            n_outputs list(int): A list containing the number of neurons in hidden and output layer.
+            activation Callable | None (default nn.ReLU()): The activation function to be applied after each hidden layer
+            activation Callable | None (default None): The activation function to be applied after the final (output) layer
+
+        For instance, MLP(10, layers = [50, 50, 10]) initializes a neural network with the following architecture:
+        - Linear layer with `n_inputs` inputs and 50 outputs
+        - ReLU activation
+        - Linear layer with 50 inputs and 50 outputs
+        - ReLU activation
+        - Linear layer with 50 inputs and 10 outputs
+
+        The weights of the linear layers are initialized with a normal distribution
+        (mean=0, std=0.1) and the biases are initialized to 0.
+
+        """
+        super().__init__()
+        self.layers = layers
+        self.activation = activation
+        self.output_activation = output_activation
+
+        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)
+
+            previous_neurons = neurons
+
+        self.net = nn.Sequential(*levels)
+
+        for m in self.net.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, mean=0, std=0.1)
+                nn.init.constant_(m.bias, val=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the neural network.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor.
+
+        """
+        return self.net(x)
+
+
+class LSTMnn(nn.Module):
+    """Default LSTM neural network model for time-series approximation."""
+
+    def __init__(self, n_inputs: int, n_outputs: int, n_hidden: int) -> None:
+        """Initializes the neural network model.
+
+        Args:
+            n_inputs (int): Number of input features.
+            n_outputs (int): Number of output features.
+            n_hidden (int): Number of hidden units in the LSTM layer.
+
+        """
+        super().__init__()
+
+        self.n_hidden = n_hidden
+
+        self.lstm = nn.LSTM(n_inputs, n_hidden)
+        self.to_out = nn.Linear(n_hidden, n_outputs)
+
+        nn.init.normal_(self.to_out.weight, mean=0, std=0.1)
+        nn.init.constant_(self.to_out.bias, val=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the neural network."""
+        # lstm_out, (hidden_state, cell_state)
+        _, (hn, _) = self.lstm(x)
+        return cast(torch.Tensor, self.to_out(hn[-1]))  # Use last hidden state

modelbase2/npe.py

@@ -0,0 +1,271 @@
+"""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
+
+__all__ = [
+    "AbstractEstimator",
+    "DefaultCache",
+    "TorchSSEstimator",
+    "TorchTimeCourseEstimator",
+    "train_torch_ss_estimator",
+    "train_torch_time_course_estimator",
+]
+
+from abc import abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import cast
+
+import numpy as np
+import pandas as pd
+import torch
+import tqdm
+from torch import nn
+from torch.optim.adam import Adam
+
+from modelbase2.nnarchitectures import MLP, DefaultDevice, LSTMnn
+from modelbase2.parallel import Cache
+
+DefaultCache = Cache(Path(".cache"))
+
+
+@dataclass(kw_only=True)
+class AbstractEstimator:
+    """Abstract class for parameter estimation using neural networks."""
+
+    parameter_names: list[str]
+
+    @abstractmethod
+    def predict(self, features: pd.Series | pd.DataFrame) -> pd.DataFrame:
+        """Predict the target values for the given features."""
+
+
+@dataclass(kw_only=True)
+class TorchSSEstimator(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 TorchTimeCourseEstimator(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)
+
+
+def _train_batched(
+    approximator: nn.Module,
+    features: torch.Tensor,
+    targets: torch.Tensor,
+    epochs: int,
+    optimizer: Adam,
+    batch_size: int,
+) -> pd.Series:
+    losses = {}
+
+    for epoch in tqdm.trange(epochs):
+        permutation = torch.randperm(features.size()[0])
+        epoch_loss = 0
+        for i in range(0, features.size()[0], batch_size):
+            optimizer.zero_grad()
+            indices = permutation[i : i + batch_size]
+
+            loss = torch.mean(
+                torch.abs(approximator(features[indices]) - targets[indices])
+            )
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.detach().numpy()
+
+        losses[epoch] = epoch_loss / (features.size()[0] / batch_size)
+    return pd.Series(losses, dtype=float)
+
+
+def _train_full(
+    approximator: nn.Module,
+    features: torch.Tensor,
+    targets: torch.Tensor,
+    epochs: int,
+    optimizer: Adam,
+) -> pd.Series:
+    losses = {}
+    for i in tqdm.trange(epochs):
+        optimizer.zero_grad()
+        loss = torch.mean(torch.abs(approximator(features) - targets))
+        loss.backward()
+        optimizer.step()
+        losses[i] = loss.detach().numpy()
+    return pd.Series(losses, dtype=float)
+
+
+def train_torch_ss_estimator(
+    features: pd.DataFrame,
+    targets: pd.DataFrame,
+    epochs: int,
+    batch_size: int | None = None,
+    approximator: nn.Module | None = None,
+    optimimzer_cls: type[Adam] = Adam,
+    device: torch.device = DefaultDevice,
+) -> tuple[TorchSSEstimator, 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
+
+    """
+    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), layers=[n_hidden, n_hidden, n_outputs]
+        ).to(device)
+
+    features_ = torch.Tensor(features.to_numpy(), device=device)
+    targets_ = torch.Tensor(targets.to_numpy(), device=device)
+
+    optimizer = optimimzer_cls(approximator.parameters())
+    if batch_size is None:
+        losses = _train_full(
+            approximator=approximator,
+            features=features_,
+            targets=targets_,
+            epochs=epochs,
+            optimizer=optimizer,
+        )
+    else:
+        losses = _train_batched(
+            approximator=approximator,
+            features=features_,
+            targets=targets_,
+            epochs=epochs,
+            optimizer=optimizer,
+            batch_size=batch_size,
+        )
+
+    return TorchSSEstimator(
+        model=approximator,
+        parameter_names=list(targets.columns),
+    ), losses
+
+
+def train_torch_time_course_estimator(
+    features: pd.DataFrame,
+    targets: pd.DataFrame,
+    epochs: int,
+    batch_size: int | None = None,
+    approximator: nn.Module | None = None,
+    optimimzer_cls: type[Adam] = Adam,
+    device: torch.device = DefaultDevice,
+) -> tuple[TorchTimeCourseEstimator, 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
+
+    """
+    if approximator is None:
+        approximator = LSTMnn(
+            n_inputs=len(features.columns),
+            n_outputs=len(targets.columns),
+            n_hidden=1,
+        ).to(device)
+
+    optimizer = optimimzer_cls(approximator.parameters())
+    features_ = torch.Tensor(
+        np.swapaxes(
+            features.to_numpy().reshape((len(targets), -1, len(features.columns))),
+            axis1=0,
+            axis2=1,
+        ),
+        device=device,
+    )
+    targets_ = torch.Tensor(targets.to_numpy(), device=device)
+    if batch_size is None:
+        losses = _train_full(
+            approximator=approximator,
+            features=features_,
+            targets=targets_,
+            epochs=epochs,
+            optimizer=optimizer,
+        )
+    else:
+        losses = _train_batched(
+            approximator=approximator,
+            features=features_,
+            targets=targets_,
+            epochs=epochs,
+            optimizer=optimizer,
+            batch_size=batch_size,
+        )
+    return TorchTimeCourseEstimator(
+        model=approximator,
+        parameter_names=list(targets.columns),
+    ), losses

modelbase2/parallel.py

@@ -0,0 +1,171 @@
+"""Parallel Execution Module.
+
+This module provides functions and classes for parallel execution and caching of
+computation results. It includes functionality for parallel processing and result
+caching using multiprocessing and pickle.
+
+Classes:
+    Cache: Cache class for storing and retrieving computation results.
+
+Functions:
+    parallelise: Execute a function in parallel over a collection of inputs.
+"""
+
+from __future__ import annotations
+
+import multiprocessing
+import pickle
+import sys
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import pebble
+from tqdm import tqdm
+
+__all__ = ["Cache", "parallelise"]
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Collection, Hashable
+
+
+def _pickle_name(k: Hashable) -> str:
+    return f"{k}.p"
+
+
+def _pickle_load(file: Path) -> Any:
+    with file.open("rb") as fp:
+        return pickle.load(fp)  # nosec
+
+
+def _pickle_save(file: Path, data: Any) -> None:
+    with file.open("wb") as fp:
+        pickle.dump(data, fp)
+
+
+@dataclass
+class Cache:
+    """Cache class for storing and retrieving computation results.
+
+    Attributes:
+        tmp_dir: Directory to store cache files.
+        name_fn: Function to generate file names from keys.
+        load_fn: Function to load data from files.
+        save_fn: Function to save data to files.
+
+    """
+
+    tmp_dir: Path = Path(".cache")
+    name_fn: Callable[[Any], str] = _pickle_name
+    load_fn: Callable[[Path], Any] = _pickle_load
+    save_fn: Callable[[Path, Any], None] = _pickle_save
+
+
+def _load_or_run[K: Hashable, Tin, Tout](
+    inp: tuple[K, Tin],
+    fn: Callable[[Tin], Tout],
+    cache: Cache | None,
+) -> tuple[K, Tout]:
+    """Load data from cache or execute function and save result.
+
+    Args:
+        inp: Tuple containing a key and input value.
+        fn: Function to execute if result is not in cache.
+        cache: Optional cache to store and retrieve results.
+
+    Returns:
+        tuple[K, Tout]: Tuple containing the key and the result of the function.
+
+    """
+    k, v = inp
+    if cache is None:
+        res = fn(v)
+    else:
+        file = cache.tmp_dir / cache.name_fn(k)
+        if file.exists():
+            return k, cast(Tout, cache.load_fn(file))
+        res = fn(v)
+        cache.save_fn(file, res)
+    return k, res
+
+
+def parallelise[K: Hashable, Tin, Tout](
+    fn: Callable[[Tin], Tout],
+    inputs: Collection[tuple[K, Tin]],
+    *,
+    cache: Cache | None = None,
+    parallel: bool = True,
+    max_workers: int | None = None,
+    timeout: float | None = None,
+    disable_tqdm: bool = False,
+    tqdm_desc: str | None = None,
+) -> dict[Tin, 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 cache is not None:
+        cache.tmp_dir.mkdir(parents=True, exist_ok=True)
+
+    if sys.platform in ["win32", "cygwin"]:
+        parallel = False
+
+    worker: Callable[[K, Tin], tuple[K, Tout]] = partial(
+        _load_or_run,
+        fn=fn,
+        cache=cache,
+    )  # type: ignore
+
+    results: dict[Tin, 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(worker, inputs, timeout=timeout)
+            it = future.result()
+            while True:
+                try:
+                    key, value = next(it)
+                    pbar.update(1)
+                    results[key] = value
+                except StopIteration:
+                    break
+                except TimeoutError:
+                    pbar.update(1)
+    else:
+        results = dict(
+            tqdm(
+                map(worker, inputs),  # type: ignore
+                total=len(inputs),
+                disable=disable_tqdm,
+                desc=tqdm_desc,
+            )  # type: ignore
+        )  # type: ignore
+    return results

modelbase2/parameterise.py

@@ -0,0 +1,28 @@
+"""Module to parameterise models."""
+
+from pathlib import Path
+
+import pandas as pd
+from parameteriser.brenda.v0 import Brenda
+
+__all__ = ["get_km_and_kcat_from_brenda"]
+
+
+def get_km_and_kcat_from_brenda(
+    ec: str,
+    brenda_path: Path,
+) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """Obtain michaelis and catalytic constants for given ec number.
+
+    You can obtain the database from https://www.brenda-enzymes.org/download.php
+    """
+    brenda = Brenda()
+    if brenda_path is not None:
+        brenda.read_database(brenda_path)
+
+    kms, kcats = brenda.get_kms_and_kcats(
+        ec=ec,
+        filter_mutant=True,
+        filter_missing_sequences=True,
+    )
+    return kms, kcats

modelbase2/paths.py

@@ -0,0 +1,36 @@
+"""Shared paths between the modelbase2 package."""
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+__all__ = [
+    "default_tmp_dir",
+]
+
+
+def default_tmp_dir(tmp_dir: Path | None, *, remove_old_cache: bool) -> Path:
+    """Returns the default temporary directory path.
+
+    If `tmp_dir` is None, it defaults to the user's home directory under ".cache/modelbase".
+    Optionally removes old cache if specified.
+
+    Args:
+        tmp_dir (Path | None): The temporary directory path. If None, defaults to
+                               Path.home() / ".cache" / "modelbase".
+        remove_old_cache (bool): If True, removes the old cache directory if it exists.
+                                 Defaults to False.
+
+    Returns:
+        Path: The path to the temporary directory.
+
+    """
+    if tmp_dir is None:
+        tmp_dir = Path.home() / ".cache" / "modelbase"
+
+    if tmp_dir.exists() and remove_old_cache:
+        shutil.rmtree(tmp_dir)
+
+    tmp_dir.mkdir(exist_ok=True, parents=True)
+    return tmp_dir