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

mxlpy/simulator.py

@@ -21,14 +21,17 @@ from sympy import lambdify
 from mxlpy.integrators import DefaultIntegrator
 from mxlpy.symbolic import to_symbolic_model
 
-__all__ = ["Result", "Simulator"]
-
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
     from mxlpy.model import Model
     from mxlpy.types import Array, ArrayLike, IntegratorProtocol, IntegratorType
 
+__all__ = [
+    "Result",
+    "Simulator",
+]
+
 
 def _normalise_split_results(
     results: list[pd.DataFrame],
@@ -500,6 +503,35 @@ class Simulator:
         self._handle_simulation_results(time, results, skipfirst=True)
         return self
 
+    def simulate_over_protocol(
+        self,
+        protocol: pd.DataFrame,
+        time_points_per_step: int = 10,
+    ) -> Self:
+        """Simulate the model over a given protocol.
+
+        Examples:
+            >>> Simulator(model).simulate_over_protocol(
+            ...     protocol,
+            ...     time_points_per_step=10
+            ... )
+
+        Args:
+            protocol: DataFrame containing the protocol.
+            time_points_per_step: Number of time points per step.
+
+        Returns:
+            The Simulator instance with updated results.
+
+        """
+        for t_end, pars in protocol.iterrows():
+            t_end = cast(pd.Timedelta, t_end)
+            self.model.update_parameters(pars.to_dict())
+            self.simulate(t_end.total_seconds(), steps=time_points_per_step)
+            if self.variables is None:
+                break
+        return self
+
     def simulate_to_steady_state(
         self,
         tolerance: float = 1e-6,
@@ -535,35 +567,6 @@ class Simulator:
         )
         return self
 
-    def simulate_over_protocol(
-        self,
-        protocol: pd.DataFrame,
-        time_points_per_step: int = 10,
-    ) -> Self:
-        """Simulate the model over a given protocol.
-
-        Examples:
-            >>> Simulator(model).simulate_over_protocol(
-            ...     protocol,
-            ...     time_points_per_step=10
-            ... )
-
-        Args:
-            protocol: DataFrame containing the protocol.
-            time_points_per_step: Number of time points per step.
-
-        Returns:
-            The Simulator instance with updated results.
-
-        """
-        for t_end, pars in protocol.iterrows():
-            t_end = cast(pd.Timedelta, t_end)
-            self.model.update_parameters(pars.to_dict())
-            self.simulate(t_end.total_seconds(), steps=time_points_per_step)
-            if self.variables is None:
-                break
-        return self
-
     def get_result(self) -> Result | None:
         """Get result of the simulation.
 

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,139 @@
+from __future__ import annotations
+
+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

@@ -1,17 +1,21 @@
-from collections.abc import Iterable
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 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
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
 
 __all__ = [
-    "Polynomial",
     "PolynomialExpansion",
-    "train_polynomial",
+    "Surrogate",
+    "train",
 ]
 
 # define custom type
@@ -26,14 +30,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 +57,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 +100,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,34 @@
+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,
+            )
+        )

mxlpy/surrogates/_torch.py

@@ -1,24 +1,29 @@
+from __future__ import annotations
+
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Self
+from typing import TYPE_CHECKING, Self
 
 import numpy as np
 import pandas as pd
 import torch
 from torch import nn
 from torch.optim.adam import Adam
-from torch.optim.optimizer import ParamsT
 
-from mxlpy.nn._torch import MLP, DefaultDevice, train
-from mxlpy.types import AbstractSurrogate
+from mxlpy.nn._torch import MLP, DefaultDevice
+from mxlpy.nn._torch import train as _train
+from mxlpy.types import AbstractSurrogate, Derived
+
+if TYPE_CHECKING:
+    from torch.optim.optimizer import ParamsT
 
 type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
 
 __all__ = [
     "LossFn",
-    "Torch",
-    "TorchTrainer",
-    "train_torch",
+    "Surrogate",
+    "Trainer",
+    "train",
 ]
 
 
@@ -37,7 +42,7 @@ def _mean_abs(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
 
 @dataclass(kw_only=True)
-class Torch(AbstractSurrogate):
+class Surrogate(AbstractSurrogate):
     """Surrogate model using PyTorch.
 
     Attributes:
@@ -65,12 +70,25 @@ class Torch(AbstractSurrogate):
                 torch.tensor(y, dtype=torch.float32),
             ).numpy()
 
+    def predict(
+        self,
+        args: dict[str, float | pd.Series | pd.DataFrame],
+    ) -> dict[str, float]:
+        """Predict outputs based on input data."""
+        return dict(
+            zip(
+                self.outputs,
+                self.predict_raw(np.array([args[arg] for arg in self.args])),
+                strict=True,
+            )
+        )
+
 
 @dataclass(init=False)
-class TorchTrainer:
+class Trainer:
     features: pd.DataFrame
     targets: pd.DataFrame
-    approximator: nn.Module
+    model: nn.Module
     optimizer: Adam
     device: torch.device
     losses: list[pd.Series]
@@ -80,22 +98,22 @@ class TorchTrainer:
         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:
         self.features = features
         self.targets = targets
 
-        if approximator is None:
-            approximator = MLP(
+        if model is None:
+            model = MLP(
                 n_inputs=len(features.columns),
                 neurons_per_layer=[50, 50, len(targets.columns)],
             )
-        self.approximator = approximator.to(device)
+        self.model = model.to(device)
 
-        self.optimizer = optimimzer_cls(approximator.parameters())
+        self.optimizer = optimizer_cls(model.parameters())
         self.device = device
         self.loss_fn = loss_fn
         self.losses = []
@@ -105,8 +123,8 @@ class TorchTrainer:
         epochs: int,
         batch_size: int | None = None,
     ) -> Self:
-        losses = train(
-            aprox=self.approximator,
+        losses = _train(
+            model=self.model,
             features=self.features.to_numpy(),
             targets=self.targets.to_numpy(),
             epochs=epochs,
@@ -127,10 +145,10 @@ class TorchTrainer:
         self,
         surrogate_args: list[str] | None = None,
         surrogate_outputs: list[str] | None = None,
-        surrogate_stoichiometries: dict[str, dict[str, float]] | None = None,
-    ) -> Torch:
-        return Torch(
-            model=self.approximator,
+        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
@@ -139,19 +157,19 @@ class TorchTrainer:
         )
 
 
-def train_torch(
+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]] | None = None,
+    surrogate_stoichiometries: dict[str, dict[str, float | Derived]] | None = None,
     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,
     loss_fn: LossFn = _mean_abs,
-) -> tuple[Torch, pd.Series]:
+) -> tuple[Surrogate, pd.Series]:
     """Train a PyTorch surrogate model.
 
     Examples:
@@ -173,8 +191,8 @@ def train_torch(
         surrogate_outputs: Names of output arguments from the surrogate.
         surrogate_stoichiometries: Mapping of variables to their stoichiometries
         batch_size: Size of mini-batches for training (None for full-batch).
-        approximator: Predefined neural network model (None to use default MLP features-50-50-output).
-        optimimzer_cls: Optimizer class to use for training (default: Adam).
+        model: Predefined neural network model (None to use default MLP features-50-50-output).
+        optimizer_cls: Optimizer class to use for training (default: Adam).
         device: Device to run the training on (default: DefaultDevice).
         loss_fn: Custom loss function or instance of torch loss object
 
@@ -182,11 +200,11 @@ def train_torch(
         tuple[TorchSurrogate, pd.Series]: Trained surrogate model and loss history.
 
     """
-    trainer = TorchTrainer(
+    trainer = Trainer(
         features=features,
         targets=targets,
-        approximator=approximator,
-        optimimzer_cls=optimimzer_cls,
+        model=model,
+        optimizer_cls=optimizer_cls,
         device=device,
         loss_fn=loss_fn,
     ).train(

mxlpy/symbolic/__init__.py

@@ -1,10 +1,12 @@
 """Symbolic utilities."""
 
+from __future__ import annotations
+
+from .strikepy import check_identifiability
+from .symbolic_model import SymbolicModel, to_symbolic_model
+
 __all__ = [
     "SymbolicModel",
     "check_identifiability",
     "to_symbolic_model",
 ]
-
-from .strikepy import check_identifiability
-from .symbolic_model import SymbolicModel, to_symbolic_model

mxlpy/symbolic/strikepy.py

@@ -10,13 +10,15 @@ FIXME:
 - performance issues of generic_rank
 """
 
+from __future__ import annotations
+
 import textwrap
 from concurrent.futures import ProcessPoolExecutor
 from dataclasses import dataclass, field
 from functools import partial
 from math import ceil, inf
 from time import time
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -27,7 +29,8 @@ import tqdm
 from sympy import Matrix
 from sympy.matrices import zeros
 
-from .symbolic_model import SymbolicModel
+if TYPE_CHECKING:
+    from .symbolic_model import SymbolicModel
 
 __all__ = [
     "Options",

mxlpy/symbolic/symbolic_model.py

@@ -1,15 +1,24 @@
 # ruff: noqa: D100, D101, D102, D103, D104, D105, D106, D107, D200, D203, D400, D401
 
 
-from collections.abc import Iterable
+from __future__ import annotations
+
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import sympy
 
 from mxlpy.meta.source_tools import fn_to_sympy
-from mxlpy.model import Model
 
-__all__ = ["SymbolicModel", "to_symbolic_model"]
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from mxlpy.model import Model
+
+__all__ = [
+    "SymbolicModel",
+    "to_symbolic_model",
+]
 
 
 @dataclass
@@ -70,6 +79,6 @@ def to_symbolic_model(model: Model) -> SymbolicModel:
         variables=variables,
         parameters=parameters,
         eqs=[eqs[i] for i in cache.var_names],
-        initial_conditions=model.variables.copy(),
-        parameter_values=model.parameters.copy(),
+        initial_conditions=model.get_initial_conditions(),
+        parameter_values=model.parameters,
     )