mxlpy 0.15.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

mxlpy/scan.py

@@ -51,7 +51,7 @@ if TYPE_CHECKING:
     from mxlpy.types import Array
 
 
-def _update_parameters_and[T](
+def _update_parameters_and_initial_conditions[T](
     pars: pd.Series,
     fn: Callable[[Model], T],
     model: Model,
@@ -67,7 +67,9 @@ def _update_parameters_and[T](
         Result of the function execution.
 
     """
-    model.update_parameters(pars.to_dict())
+    pd = pars.to_dict()
+    model.update_variables({k: v for k, v in pd.items() if k in model._variables})  # noqa: SLF001
+    model.update_parameters({k: v for k, v in pd.items() if k in model._parameters})  # noqa: SLF001
     return fn(model)
 
 
@@ -282,7 +284,6 @@ class SteadyStateWorker(Protocol):
     def __call__(
         self,
         model: Model,
-        y0: dict[str, float] | None,
         *,
         rel_norm: bool,
         integrator: IntegratorType,
@@ -297,7 +298,6 @@ class TimeCourseWorker(Protocol):
     def __call__(
         self,
         model: Model,
-        y0: dict[str, float] | None,
         time_points: Array,
         *,
         integrator: IntegratorType,
@@ -312,7 +312,6 @@ class ProtocolWorker(Protocol):
     def __call__(
         self,
         model: Model,
-        y0: dict[str, float] | None,
         protocol: pd.DataFrame,
         *,
         integrator: IntegratorType,
@@ -324,7 +323,6 @@ class ProtocolWorker(Protocol):
 
 def _steady_state_worker(
     model: Model,
-    y0: dict[str, float] | None,
     *,
     rel_norm: bool,
     integrator: IntegratorType,
@@ -343,7 +341,7 @@ def _steady_state_worker(
     """
     try:
         res = (
-            Simulator(model, y0=y0, integrator=integrator)
+            Simulator(model, integrator=integrator)
             .simulate_to_steady_state(rel_norm=rel_norm)
             .get_result()
         )
@@ -354,7 +352,6 @@ def _steady_state_worker(
 
 def _time_course_worker(
     model: Model,
-    y0: dict[str, float] | None,
     time_points: Array,
     integrator: IntegratorType,
 ) -> TimeCourse:
@@ -372,7 +369,7 @@ def _time_course_worker(
     """
     try:
         res = (
-            Simulator(model, y0=y0, integrator=integrator)
+            Simulator(model, integrator=integrator)
             .simulate_time_course(time_points=time_points)
             .get_result()
         )
@@ -387,7 +384,6 @@ def _time_course_worker(
 
 def _protocol_worker(
     model: Model,
-    y0: dict[str, float] | None,
     protocol: pd.DataFrame,
     *,
     integrator: IntegratorType = DefaultIntegrator,
@@ -408,7 +404,7 @@ def _protocol_worker(
     """
     try:
         res = (
-            Simulator(model, y0=y0, integrator=integrator)
+            Simulator(model, integrator=integrator)
             .simulate_over_protocol(
                 protocol=protocol,
                 time_points_per_step=time_points_per_step,
@@ -432,20 +428,20 @@ def _protocol_worker(
 
 def steady_state(
     model: Model,
-    parameters: pd.DataFrame,
-    y0: dict[str, float] | None = None,
     *,
+    to_scan: pd.DataFrame,
+    y0: dict[str, float] | None = None,
     parallel: bool = True,
     rel_norm: bool = False,
     cache: Cache | None = None,
     worker: SteadyStateWorker = _steady_state_worker,
     integrator: IntegratorType = DefaultIntegrator,
 ) -> SteadyStates:
-    """Get steady-state results over supplied parameters.
+    """Get steady-state results over supplied values.
 
     Args:
         model: Model instance to simulate.
-        parameters: DataFrame containing parameter values to scan.
+        to_scan: DataFrame containing parameter or initial values to scan.
         y0: Initial conditions as a dictionary {variable: value}.
         parallel: Whether to execute in parallel (default: True).
         rel_norm: Whether to use relative normalization (default: False).
@@ -478,39 +474,41 @@ def steady_state(
         | (2, 4) | 0.5  |   2 |
 
     """
+    if y0 is not None:
+        model.update_variables(y0)
+
     res = parallelise(
         partial(
-            _update_parameters_and,
+            _update_parameters_and_initial_conditions,
             fn=partial(
                 worker,
-                y0=y0,
                 rel_norm=rel_norm,
                 integrator=integrator,
             ),
             model=model,
         ),
-        inputs=list(parameters.iterrows()),
+        inputs=list(to_scan.iterrows()),
         cache=cache,
         parallel=parallel,
     )
     concs = pd.DataFrame({k: v.variables.T for k, v in res.items()}).T
     fluxes = pd.DataFrame({k: v.fluxes.T for k, v in res.items()}).T
     idx = (
-        pd.Index(parameters.iloc[:, 0])
-        if parameters.shape[1] == 1
-        else pd.MultiIndex.from_frame(parameters)
+        pd.Index(to_scan.iloc[:, 0])
+        if to_scan.shape[1] == 1
+        else pd.MultiIndex.from_frame(to_scan)
     )
     concs.index = idx
     fluxes.index = idx
-    return SteadyStates(variables=concs, fluxes=fluxes, parameters=parameters)
+    return SteadyStates(variables=concs, fluxes=fluxes, parameters=to_scan)
 
 
 def time_course(
     model: Model,
-    parameters: pd.DataFrame,
+    *,
+    to_scan: pd.DataFrame,
     time_points: Array,
     y0: dict[str, float] | None = None,
-    *,
     parallel: bool = True,
     cache: Cache | None = None,
     worker: TimeCourseWorker = _time_course_worker,
@@ -521,7 +519,7 @@ def time_course(
     Examples:
         >>> time_course(
         >>>     model,
-        >>>     parameters=pd.DataFrame({"k1": [1, 1.5, 2]}),
+        >>>     to_scan=pd.DataFrame({"k1": [1, 1.5, 2]}),
         >>>     time_points=np.linspace(0, 1, 3)
         >>> ).variables
 
@@ -539,7 +537,7 @@ def time_course(
 
         >>> time_course(
         >>>     model,
-        >>>     parameters=cartesian_product({"k1": [1, 2], "k2": [3, 4]}),
+        >>>     to_scan=cartesian_product({"k1": [1, 2], "k2": [3, 4]}),
         >>>     time_points=[0.0, 0.5, 1.0],
         >>> ).variables
 
@@ -553,7 +551,7 @@ def time_course(
 
     Args:
         model: Model instance to simulate.
-        parameters: DataFrame containing parameter values to scan.
+        to_scan: DataFrame containing parameter or initial values to scan.
         time_points: Array of time points for the simulation.
         y0: Initial conditions as a dictionary {variable: value}.
         cache: Optional cache to store and retrieve results.
@@ -566,25 +564,27 @@ def time_course(
 
 
     """
+    if y0 is not None:
+        model.update_variables(y0)
+
     res = parallelise(
         partial(
-            _update_parameters_and,
+            _update_parameters_and_initial_conditions,
             fn=partial(
                 worker,
                 time_points=time_points,
-                y0=y0,
                 integrator=integrator,
             ),
             model=model,
         ),
-        inputs=list(parameters.iterrows()),
+        inputs=list(to_scan.iterrows()),
         cache=cache,
         parallel=parallel,
     )
     concs = cast(dict, {k: v.variables for k, v in res.items()})
     fluxes = cast(dict, {k: v.fluxes for k, v in res.items()})
     return TimeCourseByPars(
-        parameters=parameters,
+        parameters=to_scan,
         variables=pd.concat(concs, names=["n", "time"]),
         fluxes=pd.concat(fluxes, names=["n", "time"]),
     )
@@ -592,11 +592,11 @@ def time_course(
 
 def time_course_over_protocol(
     model: Model,
-    parameters: pd.DataFrame,
+    *,
+    to_scan: pd.DataFrame,
     protocol: pd.DataFrame,
     time_points_per_step: int = 10,
     y0: dict[str, float] | None = None,
-    *,
     parallel: bool = True,
     cache: Cache | None = None,
     worker: ProtocolWorker = _protocol_worker,
@@ -618,7 +618,7 @@ def time_course_over_protocol(
 
     Args:
         model: Model instance to simulate.
-        parameters: DataFrame containing parameter values to scan.
+        to_scan: DataFrame containing parameter or initial values to scan.
         protocol: Protocol to follow for the simulation.
         time_points_per_step: Number of time points per protocol step (default: 10).
         y0: Initial conditions as a dictionary {variable: value}.
@@ -631,26 +631,28 @@ def time_course_over_protocol(
         TimeCourseByPars: Protocol series results for each parameter set.
 
     """
+    if y0 is not None:
+        model.update_variables(y0)
+
     res = parallelise(
         partial(
-            _update_parameters_and,
+            _update_parameters_and_initial_conditions,
             fn=partial(
                 worker,
                 protocol=protocol,
-                y0=y0,
                 time_points_per_step=time_points_per_step,
                 integrator=integrator,
             ),
             model=model,
         ),
-        inputs=list(parameters.iterrows()),
+        inputs=list(to_scan.iterrows()),
         cache=cache,
         parallel=parallel,
     )
     concs = cast(dict, {k: v.variables for k, v in res.items()})
     fluxes = cast(dict, {k: v.fluxes for k, v in res.items()})
     return ProtocolByPars(
-        parameters=parameters,
+        parameters=to_scan,
         protocol=protocol,
         variables=pd.concat(concs, names=["n", "time"]),
         fluxes=pd.concat(fluxes, names=["n", "time"]),

mxlpy/surrogates/__init__.py

@@ -19,13 +19,14 @@ from __future__ import annotations
 import contextlib
 
 with contextlib.suppress(ImportError):
-    from ._torch import TorchSurrogate, train_torch_surrogate
+    from ._torch import Torch, TorchTrainer, train_torch
 
-from ._poly import PolySurrogate, train_polynomial_surrogate
+from ._poly import Polynomial, train_polynomial
 
 __all__ = [
-    "PolySurrogate",
-    "TorchSurrogate",
-    "train_polynomial_surrogate",
-    "train_torch_surrogate",
+    "Polynomial",
+    "Torch",
+    "TorchTrainer",
+    "train_polynomial",
+    "train_torch",
 ]

mxlpy/surrogates/_poly.py

@@ -9,9 +9,9 @@ from numpy import polynomial
 from mxlpy.types import AbstractSurrogate, ArrayLike
 
 __all__ = [
-    "PolySurrogate",
+    "Polynomial",
     "PolynomialExpansion",
-    "train_polynomial_surrogate",
+    "train_polynomial",
 ]
 
 # define custom type
@@ -26,23 +26,24 @@ PolynomialExpansion = (
 
 
 @dataclass(kw_only=True)
-class PolySurrogate(AbstractSurrogate):
+class Polynomial(AbstractSurrogate):
     model: PolynomialExpansion
 
     def predict_raw(self, y: np.ndarray) -> np.ndarray:
         return self.model(y)
 
 
-def train_polynomial_surrogate(
-    feature: ArrayLike,
-    target: ArrayLike,
+def train_polynomial(
+    feature: ArrayLike | pd.Series,
+    target: ArrayLike | pd.Series,
     series: Literal[
         "Power", "Chebyshev", "Legendre", "Laguerre", "Hermite", "HermiteE"
     ] = "Power",
     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[PolySurrogate, pd.DataFrame]:
+) -> tuple[Polynomial, pd.DataFrame]:
     """Train a surrogate model based on function series expansion.
 
     Args:
@@ -51,7 +52,8 @@ def train_polynomial_surrogate(
         series: Base functions for the surrogate model
         degrees: Degrees of the polynomial to fit to the data.
         surrogate_args: Additional arguments for the surrogate model.
-        surrogate_stoichiometries: Stoichiometries for the surrogate model.
+        surrogate_outputs: Names of the surrogate model outputs.
+        surrogate_stoichiometries: Mapping of variables to their stoichiometries
 
     Returns:
         PolySurrogate: Polynomial surrogate model.
@@ -83,9 +85,10 @@ def train_polynomial_surrogate(
     # Choose the model with the lowest AIC
     model = models[np.argmin(score)]
     return (
-        PolySurrogate(
+        Polynomial(
             model=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 {},

mxlpy/surrogates/_torch.py

@@ -1,5 +1,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Self
 
 import numpy as np
 import pandas as pd
@@ -12,14 +13,32 @@ from torch.optim.optimizer import ParamsT
 from mxlpy.nn._torch import MLP, DefaultDevice
 from mxlpy.types import AbstractSurrogate
 
+type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+
 __all__ = [
-    "TorchSurrogate",
-    "train_torch_surrogate",
+    "LossFn",
+    "Torch",
+    "TorchTrainer",
+    "train_torch",
 ]
 
 
+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 TorchSurrogate(AbstractSurrogate):
+class Torch(AbstractSurrogate):
     """Surrogate model using PyTorch.
 
     Attributes:
@@ -48,6 +67,91 @@ class TorchSurrogate(AbstractSurrogate):
             ).numpy()
 
 
+@dataclass(init=False)
+class TorchTrainer:
+    features: pd.DataFrame
+    targets: pd.DataFrame
+    approximator: nn.Module
+    optimizer: 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:
+        self.features = features
+        self.targets = targets
+
+        if approximator is None:
+            approximator = MLP(
+                n_inputs=len(features.columns),
+                neurons_per_layer=[50, 50, len(targets.columns)],
+            )
+        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:
+        if batch_size is None:
+            losses = _train_full(
+                aprox=self.approximator,
+                features=self.features,
+                targets=self.targets,
+                epochs=epochs,
+                optimizer=self.optimizer,
+                device=self.device,
+                loss_fn=self.loss_fn,
+            )
+        else:
+            losses = _train_batched(
+                aprox=self.approximator,
+                features=self.features,
+                targets=self.targets,
+                epochs=epochs,
+                optimizer=self.optimizer,
+                device=self.device,
+                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:
+        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]] | None = None,
+    ) -> Torch:
+        return Torch(
+            model=self.approximator,
+            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_batched(
     aprox: nn.Module,
     features: pd.DataFrame,
@@ -56,6 +160,7 @@ def _train_batched(
     optimizer: Adam,
     device: torch.device,
     batch_size: int,
+    loss_fn: LossFn,
 ) -> pd.Series:
     """Train the neural network using mini-batch gradient descent.
 
@@ -67,6 +172,7 @@ def _train_batched(
         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.
@@ -79,7 +185,7 @@ def _train_batched(
         X = torch.Tensor(features.iloc[idxs].to_numpy(), device=device)
         Y = torch.Tensor(targets.iloc[idxs].to_numpy(), device=device)
         optimizer.zero_grad()
-        loss = torch.mean(torch.abs(aprox(X) - Y))
+        loss = loss_fn(aprox(X), Y)
         loss.backward()
         optimizer.step()
         losses[i] = loss.detach().numpy()
@@ -93,6 +199,7 @@ def _train_full(
     epochs: int,
     optimizer: Adam,
     device: torch.device,
+    loss_fn: Callable,
 ) -> pd.Series:
     """Train the neural network using full-batch gradient descent.
 
@@ -103,6 +210,7 @@ def _train_full(
         epochs: Number of training epochs.
         optimizer: Optimizer for training.
         device: Torch device
+        loss_fn: Loss function
 
     Returns:
         pd.Series: Series containing the training loss history.
@@ -114,24 +222,26 @@ def _train_full(
     losses = {}
     for i in tqdm.trange(epochs):
         optimizer.zero_grad()
-        loss = torch.mean(torch.abs(aprox(X) - Y))
+        loss = loss_fn(aprox(X), Y)
         loss.backward()
         optimizer.step()
         losses[i] = loss.detach().numpy()
     return pd.Series(losses, dtype=float)
 
 
-def train_torch_surrogate(
+def train_torch(
     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,
     batch_size: int | None = None,
     approximator: nn.Module | None = None,
     optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
     device: torch.device = DefaultDevice,
-) -> tuple[TorchSurrogate, pd.Series]:
+    loss_fn: LossFn = _mean_abs,
+) -> tuple[Torch, pd.Series]:
     """Train a PyTorch surrogate model.
 
     Examples:
@@ -143,54 +253,38 @@ def train_torch_surrogate(
         ...     surrogate_stoichiometries={
         ...         "v1": {"x1": -1, "x2": 1, "ATP": -1},
         ...     },
-        ...)
+        ...)surrogate_stoichiometries
 
     Args:
         features: DataFrame containing the input features for training.
         targets: DataFrame containing the target values for training.
         epochs: Number of training epochs.
-        surrogate_args: List of input variable names for the surrogate model.
-        surrogate_stoichiometries: Dictionary mapping reaction names to stoichiometries.
+        surrogate_args: Names of inputs arguments for the surrogate model.
+        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).
         device: Device to run the training on (default: DefaultDevice).
+        loss_fn: Custom loss function or instance of torch loss object
 
     Returns:
         tuple[TorchSurrogate, pd.Series]: Trained surrogate model and loss history.
 
     """
-    if approximator is None:
-        approximator = MLP(
-            n_inputs=len(features.columns),
-            neurons_per_layer=[50, 50, len(targets.columns)],
-        ).to(device)
-
-    optimizer = optimimzer_cls(approximator.parameters())
-    if batch_size is None:
-        losses = _train_full(
-            aprox=approximator,
-            features=features,
-            targets=targets,
-            epochs=epochs,
-            optimizer=optimizer,
-            device=device,
-        )
-    else:
-        losses = _train_batched(
-            aprox=approximator,
-            features=features,
-            targets=targets,
-            epochs=epochs,
-            optimizer=optimizer,
-            device=device,
-            batch_size=batch_size,
-        )
-    surrogate = TorchSurrogate(
-        model=approximator,
-        args=surrogate_args if surrogate_args is not None else [],
-        stoichiometries=surrogate_stoichiometries
-        if surrogate_stoichiometries is not None
-        else {},
+    trainer = TorchTrainer(
+        features=features,
+        targets=targets,
+        approximator=approximator,
+        optimimzer_cls=optimimzer_cls,
+        device=device,
+        loss_fn=loss_fn,
+    ).train(
+        epochs=epochs,
+        batch_size=batch_size,
     )
-    return surrogate, losses
+    return trainer.get_surrogate(
+        surrogate_args=surrogate_args,
+        surrogate_outputs=surrogate_outputs,
+        surrogate_stoichiometries=surrogate_stoichiometries,
+    ), trainer.get_loss()

mxlpy/symbolic/strikepy.py

@@ -27,9 +27,7 @@ import tqdm
 from sympy import Matrix
 from sympy.matrices import zeros
 
-from mxlpy.model import Model
-
-from .symbolic_model import SymbolicModel, to_symbolic_model
+from .symbolic_model import SymbolicModel
 
 __all__ = [
     "Options",