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

mxlpy/npe/_torch.py

@@ -12,7 +12,6 @@ Functions:
 from __future__ import annotations
 
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING, Self, cast
 
 import numpy as np
@@ -22,7 +21,6 @@ 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:
@@ -30,19 +28,17 @@ if TYPE_CHECKING:
 
     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",
+    "SteadyState",
+    "SteadyStateTrainer",
+    "TimeCourse",
+    "TimeCourseTrainer",
+    "train_steady_state",
+    "train_time_course",
 ]
 
 
@@ -61,7 +57,7 @@ def _mean_abs(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
 
 @dataclass(kw_only=True)
-class TorchSteadyState(AbstractEstimator):
+class SteadyState(AbstractEstimator):
     """Estimator for steady state data using PyTorch models."""
 
     model: torch.nn.Module
@@ -74,7 +70,7 @@ class TorchSteadyState(AbstractEstimator):
 
 
 @dataclass(kw_only=True)
-class TorchTimeCourse(AbstractEstimator):
+class TimeCourse(AbstractEstimator):
     """Estimator for time course data using PyTorch models."""
 
     model: torch.nn.Module
@@ -101,13 +97,13 @@ class TorchTimeCourse(AbstractEstimator):
 
 
 @dataclass
-class TorchSteadyStateTrainer:
+class SteadyStateTrainer:
     """Trainer for steady state data using PyTorch models."""
 
     features: pd.DataFrame
     targets: pd.DataFrame
-    approximator: nn.Module
-    optimimzer: Adam
+    model: nn.Module
+    optimizer: Adam
     device: torch.device
     losses: list[pd.Series]
     loss_fn: LossFn
@@ -116,8 +112,8 @@ class TorchSteadyStateTrainer:
         self,
         features: pd.DataFrame,
         targets: pd.DataFrame,
-        approximator: nn.Module | None = None,
-        optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
+        model: nn.Module | None = None,
+        optimizer_cls: Callable[[ParamsT], Adam] = Adam,
         device: torch.device = DefaultDevice,
         loss_fn: LossFn = _mean_abs,
     ) -> None:
@@ -126,8 +122,8 @@ class TorchSteadyStateTrainer:
         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)
+            model: Predefined neural network model (None to use default MLP)
+            optimizer_cls: Optimizer class to use for training (default: Adam)
             device: Device to run the training on (default: DefaultDevice)
             loss_fn: Loss function
 
@@ -135,15 +131,15 @@ class TorchSteadyStateTrainer:
         self.features = features
         self.targets = targets
 
-        if approximator is None:
+        if model is None:
             n_hidden = max(2 * len(features.columns) * len(targets.columns), 10)
             n_outputs = len(targets.columns)
-            approximator = MLP(
+            model = 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.model = model.to(device)
+        self.optimizer = optimizer_cls(model.parameters())
         self.device = device
         self.loss_fn = loss_fn
         self.losses = []
@@ -161,7 +157,7 @@ class TorchSteadyStateTrainer:
 
         """
         losses = train(
-            aprox=self.approximator,
+            model=self.model,
             features=self.features.to_numpy(),
             targets=self.targets.to_numpy(),
             epochs=epochs,
@@ -180,22 +176,22 @@ class TorchSteadyStateTrainer:
         """Get the loss history of the training process."""
         return pd.concat(self.losses)
 
-    def get_estimator(self) -> TorchSteadyState:
+    def get_estimator(self) -> SteadyState:
         """Get the trained estimator."""
-        return TorchSteadyState(
-            model=self.approximator,
+        return SteadyState(
+            model=self.model,
             parameter_names=list(self.targets.columns),
         )
 
 
 @dataclass
-class TorchTimeCourseTrainer:
+class TimeCourseTrainer:
     """Trainer for time course data using PyTorch models."""
 
     features: pd.DataFrame
     targets: pd.DataFrame
-    approximator: nn.Module
-    optimimzer: Adam
+    model: nn.Module
+    optimizer: Adam
     device: torch.device
     losses: list[pd.Series]
     loss_fn: LossFn
@@ -204,8 +200,8 @@ class TorchTimeCourseTrainer:
         self,
         features: pd.DataFrame,
         targets: pd.DataFrame,
-        approximator: nn.Module | None = None,
-        optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
+        model: nn.Module | None = None,
+        optimizer_cls: Callable[[ParamsT], Adam] = Adam,
         device: torch.device = DefaultDevice,
         loss_fn: LossFn = _mean_abs,
     ) -> None:
@@ -214,8 +210,8 @@ class TorchTimeCourseTrainer:
         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)
+            model: Predefined neural network model (None to use default LSTM)
+            optimizer_cls: Optimizer class to use for training (default: Adam)
             device: Device to run the training on (default: DefaultDevice)
             loss_fn: Loss function
 
@@ -223,14 +219,14 @@ class TorchTimeCourseTrainer:
         self.features = features
         self.targets = targets
 
-        if approximator is None:
-            approximator = LSTM(
+        if model is None:
+            model = 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.model = model.to(device)
+        self.optimizer = optimizer_cls(model.parameters())
         self.device = device
         self.loss_fn = loss_fn
         self.losses = []
@@ -248,7 +244,7 @@ class TorchTimeCourseTrainer:
 
         """
         losses = train(
-            aprox=self.approximator,
+            model=self.model,
             features=np.swapaxes(
                 self.features.to_numpy().reshape(
                     (len(self.targets), -1, len(self.features.columns))
@@ -273,23 +269,23 @@ class TorchTimeCourseTrainer:
         """Get the loss history of the training process."""
         return pd.concat(self.losses)
 
-    def get_estimator(self) -> TorchTimeCourse:
+    def get_estimator(self) -> TimeCourse:
         """Get the trained estimator."""
-        return TorchTimeCourse(
-            model=self.approximator,
+        return TimeCourse(
+            model=self.model,
             parameter_names=list(self.targets.columns),
         )
 
 
-def train_torch_steady_state(
+def train_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,
+    model: nn.Module | None = None,
+    optimizer_cls: Callable[[ParamsT], Adam] = Adam,
     device: torch.device = DefaultDevice,
-) -> tuple[TorchSteadyState, pd.Series]:
+) -> tuple[SteadyState, pd.Series]:
     """Train a PyTorch steady state estimator.
 
     This function trains a neural network model to estimate steady state data
@@ -304,34 +300,34 @@ def train_torch_steady_state(
         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)
+        model: Predefined neural network model (None to use default MLP)
+        optimizer_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(
+    trainer = SteadyStateTrainer(
         features=features,
         targets=targets,
-        approximator=approximator,
-        optimimzer_cls=optimimzer_cls,
+        model=model,
+        optimizer_cls=optimizer_cls,
         device=device,
     ).train(epochs=epochs, batch_size=batch_size)
 
     return trainer.get_estimator(), trainer.get_loss()
 
 
-def train_torch_time_course(
+def train_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,
+    model: nn.Module | None = None,
+    optimizer_cls: Callable[[ParamsT], Adam] = Adam,
     device: torch.device = DefaultDevice,
-) -> tuple[TorchTimeCourse, pd.Series]:
+) -> tuple[TimeCourse, pd.Series]:
     """Train a PyTorch time course estimator.
 
     This function trains a neural network model to estimate time course data
@@ -346,19 +342,19 @@ def train_torch_time_course(
         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)
+        model: Predefined neural network model (None to use default LSTM)
+        optimizer_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(
+    trainer = TimeCourseTrainer(
         features=features,
         targets=targets,
-        approximator=approximator,
-        optimimzer_cls=optimimzer_cls,
+        model=model,
+        optimizer_cls=optimizer_cls,
         device=device,
     ).train(epochs=epochs, batch_size=batch_size)
 

mxlpy/plot.py

@@ -712,11 +712,11 @@ def _create_heatmap(
     if title is not None:
         ax.set_title(title)
     ax.set_xticks(
-        np.arange(0, len(df.columns), 1) + 0.5,
+        np.arange(0, len(df.columns), 1, dtype=float) + 0.5,
         labels=xticklabels,
     )
     ax.set_yticks(
-        np.arange(0, len(df.index), 1) + 0.5,
+        np.arange(0, len(df.index), 1, dtype=float) + 0.5,
         labels=yticklabels,
     )
 

mxlpy/sbml/_export.py

@@ -447,7 +447,14 @@ def _create_sbml_variables(
         cpd.setConstant(False)
         cpd.setBoundaryCondition(False)
         cpd.setHasOnlySubstanceUnits(False)
-        cpd.setInitialAmount(float(value))
+        if isinstance(value, Derived):
+            ar = sbml_model.createInitialAssignment()
+            ar.setId(_convert_id_to_sbml(id_=name, prefix="IA"))
+            ar.setName(_convert_id_to_sbml(id_=name, prefix="IA"))
+            ar.setVariable(_convert_id_to_sbml(id_=name, prefix="IA"))
+            ar.setMath(_sbmlify_fn(value.fn, value.args))
+        else:
+            cpd.setInitialAmount(float(value))
 
 
 def _create_sbml_derived_variables(*, model: Model, sbml_model: libsbml.Model) -> None:

mxlpy/surrogates/__init__.py

@@ -4,29 +4,37 @@ This module provides classes and functions for creating and training surrogate m
 for metabolic simulations. It includes functionality for both steady-state and time-series
 data using neural networks.
 
-Classes:
-    AbstractSurrogate: Abstract base class for surrogate models.
-    TorchSurrogate: Surrogate model using PyTorch.
-    Approximator: Neural network approximator for surrogate modeling.
-
-Functions:
-    train_torch_surrogate: Train a PyTorch surrogate model.
-    train_torch_time_course_estimator: Train a PyTorch time course estimator.
 """
 
 from __future__ import annotations
 
-import contextlib
+from typing import TYPE_CHECKING
 
-with contextlib.suppress(ImportError):
-    from ._torch import Torch, TorchTrainer, train_torch
+if TYPE_CHECKING:
+    import contextlib
 
-from ._poly import Polynomial, train_polynomial
+    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.surrogates._keras",
+        error_strings={"module": "keras", "install_name": "mxlpy[tf]"},
+    )
+    torch = lazy_module(
+        "mxlpy.surrogates._torch",
+        error_strings={"module": "torch", "install_name": "mxlpy[torch]"},
+    )
+
+
+from . import _poly as poly
+from . import _qss as qss
 
 __all__ = [
-    "Polynomial",
-    "Torch",
-    "TorchTrainer",
-    "train_polynomial",
-    "train_torch",
+    "keras",
+    "poly",
+    "qss",
+    "torch",
 ]

mxlpy/surrogates/_keras.py

@@ -0,0 +1,137 @@
+from dataclasses import dataclass
+from typing import Self, cast
+
+import keras
+import numpy as np
+import pandas as pd
+
+from mxlpy.nn._keras import MLP
+from mxlpy.nn._keras import train as _train
+from mxlpy.types import AbstractSurrogate, Array, Derived
+
+__all__ = [
+    "DefaultLoss",
+    "DefaultOptimizer",
+    "LossFn",
+    "Optimizer",
+    "Surrogate",
+    "Trainer",
+    "train",
+]
+
+type Optimizer = keras.optimizers.Optimizer | str
+type LossFn = keras.losses.Loss | str
+
+DefaultOptimizer = keras.optimizers.Adam()
+DefaultLoss = keras.losses.MeanAbsoluteError()
+
+
+@dataclass(kw_only=True)
+class Surrogate(AbstractSurrogate):
+    model: keras.Model
+
+    def predict_raw(self, y: Array) -> Array:
+        return np.atleast_1d(np.squeeze(self.model.predict(y)))
+
+    def predict(
+        self, args: dict[str, float | pd.Series | pd.DataFrame]
+    ) -> dict[str, float]:
+        return dict(
+            zip(
+                self.outputs,
+                self.predict_raw(np.array([args[arg] for arg in self.args])),
+                strict=True,
+            )
+        )
+
+
+@dataclass(init=False)
+class Trainer:
+    features: pd.DataFrame
+    targets: pd.DataFrame
+    model: keras.Model
+    optimizer: Optimizer | str
+    losses: list[pd.Series]
+    loss_fn: LossFn
+
+    def __init__(
+        self,
+        features: pd.DataFrame,
+        targets: pd.DataFrame,
+        model: keras.Model | None = None,
+        optimizer: Optimizer = DefaultOptimizer,
+        loss: LossFn = DefaultLoss,
+    ) -> None:
+        self.features = features
+        self.targets = targets
+        if model is None:
+            model = MLP(
+                n_inputs=len(features.columns),
+                neurons_per_layer=[50, 50, len(targets.columns)],
+            )
+        self.model = model
+        model.compile(optimizer=cast(str, optimizer), loss=loss)
+
+        self.losses = []
+
+    def train(self, epochs: int, batch_size: int | None = None) -> Self:
+        losses = _train(
+            model=self.model,
+            features=self.features,
+            targets=self.targets,
+            epochs=epochs,
+            batch_size=batch_size,
+        )
+
+        if len(self.losses) > 0:
+            losses.index += self.losses[-1].index[-1]
+        self.losses.append(losses)
+
+        return self
+
+    def get_loss(self) -> pd.Series:
+        return pd.concat(self.losses)
+
+    def get_surrogate(
+        self,
+        surrogate_args: list[str] | None = None,
+        surrogate_outputs: list[str] | None = None,
+        surrogate_stoichiometries: dict[str, dict[str, float | Derived]] | None = None,
+    ) -> Surrogate:
+        return Surrogate(
+            model=self.model,
+            args=surrogate_args if surrogate_args is not None else [],
+            outputs=surrogate_outputs if surrogate_outputs is not None else [],
+            stoichiometries=surrogate_stoichiometries
+            if surrogate_stoichiometries is not None
+            else {},
+        )
+
+
+def train(
+    features: pd.DataFrame,
+    targets: pd.DataFrame,
+    epochs: int,
+    surrogate_args: list[str] | None = None,
+    surrogate_outputs: list[str] | None = None,
+    surrogate_stoichiometries: dict[str, dict[str, float | Derived]] | None = None,
+    batch_size: int | None = None,
+    model: keras.Model | None = None,
+    optimizer: Optimizer = DefaultOptimizer,
+    loss: LossFn = DefaultLoss,
+) -> tuple[Surrogate, pd.Series]:
+    trainer = Trainer(
+        features=features,
+        targets=targets,
+        model=model,
+        optimizer=optimizer,
+        loss=loss,
+    ).train(
+        epochs=epochs,
+        batch_size=batch_size,
+    )
+    return trainer.get_surrogate(
+        surrogate_args=surrogate_args,
+        surrogate_outputs=surrogate_outputs,
+        surrogate_stoichiometries=surrogate_stoichiometries,
+    ), trainer.get_loss()

mxlpy/surrogates/_poly.py

@@ -6,12 +6,12 @@ import numpy as np
 import pandas as pd
 from numpy import polynomial
 
-from mxlpy.types import AbstractSurrogate, ArrayLike
+from mxlpy.types import AbstractSurrogate, ArrayLike, Derived
 
 __all__ = [
-    "Polynomial",
     "PolynomialExpansion",
-    "train_polynomial",
+    "Surrogate",
+    "train",
 ]
 
 # define custom type
@@ -26,14 +26,25 @@ PolynomialExpansion = (
 
 
 @dataclass(kw_only=True)
-class Polynomial(AbstractSurrogate):
+class Surrogate(AbstractSurrogate):
     model: PolynomialExpansion
 
     def predict_raw(self, y: np.ndarray) -> np.ndarray:
         return self.model(y)
 
+    def predict(
+        self, args: dict[str, float | pd.Series | pd.DataFrame]
+    ) -> dict[str, float]:
+        return dict(
+            zip(
+                self.outputs,
+                self.model(np.array([args[arg] for arg in self.args])),
+                strict=True,
+            )
+        )
 
-def train_polynomial(
+
+def train(
     feature: ArrayLike | pd.Series,
     target: ArrayLike | pd.Series,
     series: Literal[
@@ -42,8 +53,8 @@ def train_polynomial(
     degrees: Iterable[int] = (1, 2, 3, 4, 5, 6, 7),
     surrogate_args: list[str] | None = None,
     surrogate_outputs: list[str] | None = None,
-    surrogate_stoichiometries: dict[str, dict[str, float]] | None = None,
-) -> tuple[Polynomial, pd.DataFrame]:
+    surrogate_stoichiometries: dict[str, dict[str, float | Derived]] | None = None,
+) -> tuple[Surrogate, pd.DataFrame]:
     """Train a surrogate model based on function series expansion.
 
     Args:
@@ -85,7 +96,7 @@ def train_polynomial(
     # Choose the model with the lowest AIC
     model = models[np.argmin(score)]
     return (
-        Polynomial(
+        Surrogate(
             model=model,
             args=surrogate_args if surrogate_args is not None else [],
             outputs=surrogate_outputs if surrogate_outputs is not None else [],

mxlpy/surrogates/_qss.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from mxlpy.types import AbstractSurrogate, Array
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+__all__ = ["QSSFn", "Surrogate"]
+
+type QSSFn = Callable[..., Iterable[float] | Array]
+
+
+@dataclass(kw_only=True)
+class Surrogate(AbstractSurrogate):
+    model: QSSFn
+
+    def predict(
+        self,
+        args: dict[str, float | pd.Series | pd.DataFrame],
+    ) -> dict[str, float]:
+        return dict(
+            zip(
+                self.outputs,
+                self.model(*(args[arg] for arg in self.args)),
+                strict=True,
+            )
+        )