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

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,
+            )
+        )

mxlpy/surrogates/_torch.py

@@ -5,21 +5,21 @@ from typing import Self
 import numpy as np
 import pandas as pd
 import torch
-import tqdm
 from torch import nn
 from torch.optim.adam import Adam
 from torch.optim.optimizer import ParamsT
 
 from mxlpy.nn._torch import MLP, DefaultDevice
-from mxlpy.types import AbstractSurrogate
+from mxlpy.nn._torch import train as _train
+from mxlpy.types import AbstractSurrogate, Derived
 
 type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
 
 __all__ = [
     "LossFn",
-    "Torch",
-    "TorchTrainer",
-    "train_torch",
+    "Surrogate",
+    "Trainer",
+    "train",
 ]
 
 
@@ -38,7 +38,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:
@@ -66,12 +66,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]
@@ -81,22 +94,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 = []
@@ -106,28 +119,16 @@ class TorchTrainer:
         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,
-            )
-
+        losses = _train(
+            model=self.model,
+            features=self.features.to_numpy(),
+            targets=self.targets.to_numpy(),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            device=self.device,
+            loss_fn=self.loss_fn,
+        )
         if len(self.losses) > 0:
             losses.index += self.losses[-1].index[-1]
         self.losses.append(losses)
@@ -140,10 +141,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
@@ -152,96 +153,19 @@ class TorchTrainer:
         )
 
 
-def _train_batched(
-    aprox: nn.Module,
-    features: pd.DataFrame,
-    targets: pd.DataFrame,
-    epochs: int,
-    optimizer: Adam,
-    device: torch.device,
-    batch_size: int,
-    loss_fn: LossFn,
-) -> pd.Series:
-    """Train the neural network using mini-batch gradient descent.
-
-    Args:
-        aprox: Neural network model to train.
-        features: Input features as a tensor.
-        targets: Target values as a tensor.
-        epochs: Number of training epochs.
-        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.
-
-    """
-    rng = np.random.default_rng()
-    losses = {}
-    for i in tqdm.trange(epochs):
-        idxs = rng.choice(features.index, size=batch_size)
-        X = torch.Tensor(features.iloc[idxs].to_numpy(), device=device)
-        Y = torch.Tensor(targets.iloc[idxs].to_numpy(), device=device)
-        optimizer.zero_grad()
-        loss = loss_fn(aprox(X), Y)
-        loss.backward()
-        optimizer.step()
-        losses[i] = loss.detach().numpy()
-    return pd.Series(losses, dtype=float)
-
-
-def _train_full(
-    aprox: nn.Module,
-    features: pd.DataFrame,
-    targets: pd.DataFrame,
-    epochs: int,
-    optimizer: Adam,
-    device: torch.device,
-    loss_fn: Callable,
-) -> pd.Series:
-    """Train the neural network using full-batch gradient descent.
-
-    Args:
-        aprox: Neural network model to train.
-        features: Input features as a tensor.
-        targets: Target values as a tensor.
-        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.
-
-    """
-    X = torch.Tensor(features.to_numpy(), device=device)
-    Y = torch.Tensor(targets.to_numpy(), device=device)
-
-    losses = {}
-    for i in tqdm.trange(epochs):
-        optimizer.zero_grad()
-        loss = loss_fn(aprox(X), Y)
-        loss.backward()
-        optimizer.step()
-        losses[i] = loss.detach().numpy()
-    return pd.Series(losses, dtype=float)
-
-
-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:
@@ -263,8 +187,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
 
@@ -272,11 +196,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/symbolic_model.py

@@ -70,6 +70,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,
     )

mxlpy/types.py

@@ -46,10 +46,7 @@ __all__ = [
     "unwrap2",
 ]
 
-# Re-exporting some types here, because their imports have
-# changed between Python versions and I have no interest in
-# fixing it in every file
-from collections.abc import Callable, Iterator, Mapping
+from collections.abc import Callable, Iterable, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, ParamSpec, Protocol, TypeVar, cast
 
 import numpy as np
@@ -162,7 +159,7 @@ class Derived:
     fn: RateFn
     args: list[str]
 
-    def calculate(self, dependent: dict[str, float]) -> float:
+    def calculate(self, dependent: dict[str, Any]) -> float:
         """Calculate the derived value.
 
         Args:
@@ -174,7 +171,7 @@ class Derived:
         """
         return cast(float, self.fn(*(dependent[arg] for arg in self.args)))
 
-    def calculate_inpl(self, name: str, dependent: dict[str, float]) -> None:
+    def calculate_inpl(self, name: str, dependent: dict[str, Any]) -> None:
         """Calculate the derived value in place.
 
         Args:
@@ -184,20 +181,6 @@ class Derived:
         """
         dependent[name] = cast(float, self.fn(*(dependent[arg] for arg in self.args)))
 
-    def calculate_inpl_time_course(self, name: str, dependent: pd.DataFrame) -> None:
-        """Calculate the derived value in place.
-
-        Args:
-            name: Name of the derived variable.
-            dependent: Dictionary of dependent variables.
-
-        """
-        try:
-            dependent[name] = self.fn(*dependent.loc[:, self.args].to_numpy().T)
-        except ValueError:  # e.g. numpy.where
-            sub = dependent.loc[:, self.args].to_numpy()
-            dependent[name] = [self.fn(*row) for row in sub]
-
 
 @dataclass(kw_only=True, slots=True)
 class Readout:
@@ -206,7 +189,7 @@ class Readout:
     fn: RateFn
     args: list[str]
 
-    def calculate(self, dependent: dict[str, float]) -> float:
+    def calculate(self, dependent: dict[str, Any]) -> float:
         """Calculate the derived value.
 
         Args:
@@ -218,8 +201,8 @@ class Readout:
         """
         return cast(float, self.fn(*(dependent[arg] for arg in self.args)))
 
-    def calculate_inpl(self, name: str, dependent: dict[str, float]) -> None:
-        """Calculate the readout in place.
+    def calculate_inpl(self, name: str, dependent: dict[str, Any]) -> None:
+        """Calculate the reaction in place.
 
         Args:
             name: Name of the derived variable.
@@ -228,20 +211,6 @@ class Readout:
         """
         dependent[name] = cast(float, self.fn(*(dependent[arg] for arg in self.args)))
 
-    def calculate_inpl_time_course(self, name: str, dependent: pd.DataFrame) -> None:
-        """Calculate the derived value in place.
-
-        Args:
-            name: Name of the derived variable.
-            dependent: Dictionary of dependent variables.
-
-        """
-        try:
-            dependent[name] = self.fn(*dependent.loc[:, self.args].to_numpy().T)
-        except ValueError:  # e.g. numpy.where
-            sub = dependent.loc[:, self.args].to_numpy()
-            dependent[name] = [self.fn(*row) for row in sub]
-
 
 @dataclass(kw_only=True, slots=True)
 class Reaction:
@@ -258,7 +227,7 @@ class Reaction:
 
         return [k for k in self.args if k in include and k not in exclude]
 
-    def calculate(self, dependent: dict[str, float]) -> float:
+    def calculate(self, dependent: dict[str, Any]) -> float:
         """Calculate the derived value.
 
         Args:
@@ -270,7 +239,7 @@ class Reaction:
         """
         return cast(float, self.fn(*(dependent[arg] for arg in self.args)))
 
-    def calculate_inpl(self, name: str, dependent: dict[str, float]) -> None:
+    def calculate_inpl(self, name: str, dependent: dict[str, Any]) -> None:
         """Calculate the reaction in place.
 
         Args:
@@ -280,19 +249,57 @@ class Reaction:
         """
         dependent[name] = cast(float, self.fn(*(dependent[arg] for arg in self.args)))
 
-    def calculate_inpl_time_course(self, name: str, dependent: pd.DataFrame) -> None:
-        """Calculate the derived value in place.
 
-        Args:
-            name: Name of the derived variable.
-            dependent: Dictionary of dependent variables.
+@dataclass(kw_only=True)
+class AbstractSurrogate:
+    """Abstract base class for surrogate models.
 
-        """
-        try:
-            dependent[name] = self.fn(*dependent.loc[:, self.args].to_numpy().T)
-        except ValueError:  # e.g. numpy.where
-            sub = dependent.loc[:, self.args].to_numpy()
-            dependent[name] = [self.fn(*row) for row in sub]
+    Attributes:
+        inputs: List of input variable names.
+        stoichiometries: Dictionary mapping reaction names to stoichiometries.
+
+    Methods:
+        predict: Abstract method to predict outputs based on input data.
+
+    """
+
+    args: list[str]
+    outputs: list[str]
+    stoichiometries: dict[str, dict[str, float | Derived]] = field(default_factory=dict)
+
+    @abstractmethod
+    def predict(
+        self, args: dict[str, float | pd.Series | pd.DataFrame]
+    ) -> dict[str, float]:
+        """Predict outputs based on input data."""
+
+    def calculate_inpl(
+        self,
+        name: str,  # noqa: ARG002, for API compatibility
+        args: dict[str, float | pd.Series | pd.DataFrame],
+    ) -> None:
+        """Predict outputs based on input data."""
+        args |= self.predict(args=args)
+
+
+@dataclass(kw_only=True)
+class MockSurrogate(AbstractSurrogate):
+    """Mock surrogate model for testing purposes."""
+
+    fn: Callable[..., Iterable[float]]
+
+    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.fn(*(args[i] for i in self.args)),
+                strict=True,
+            )
+        )  # type: ignore
 
 
 @dataclass(kw_only=True, slots=True)
@@ -432,70 +439,6 @@ class ProtocolByPars:
         return cast(pd.DataFrame, mean.unstack().T)
 
 
-@dataclass(kw_only=True)
-class AbstractSurrogate:
-    """Abstract base class for surrogate models.
-
-    Attributes:
-        inputs: List of input variable names.
-        stoichiometries: Dictionary mapping reaction names to stoichiometries.
-
-    Methods:
-        predict: Abstract method to predict outputs based on input data.
-
-    """
-
-    args: list[str]
-    outputs: list[str]
-    stoichiometries: dict[str, dict[str, float]] = field(default_factory=dict)
-
-    @abstractmethod
-    def predict_raw(self, y: np.ndarray) -> np.ndarray:
-        """Predict outputs based on input data."""
-
-    def predict(self, y: np.ndarray) -> dict[str, float]:
-        """Predict outputs based on input data."""
-        return dict(
-            zip(
-                self.outputs,
-                self.predict_raw(y),
-                strict=True,
-            )
-        )
-
-    def calculate_inpl(
-        self,
-        name: str,  # noqa: ARG002, for API compatibility
-        args: dict[str, float],
-    ) -> None:
-        """Predict outputs based on input data."""
-        args |= self.predict(np.array([args[arg] for arg in self.args]))
-
-    def calculate_inpl_time_course(
-        self,
-        name: str,  # noqa: ARG002, for API compatibility
-        args: pd.DataFrame,
-    ) -> None:
-        """Predict outputs based on input data."""
-        args[self.outputs] = pd.DataFrame(
-            [self.predict(y) for y in args.loc[:, self.args].to_numpy()],
-            index=args.index,
-            dtype=float,
-        )
-
-
-@dataclass(kw_only=True)
-class MockSurrogate(AbstractSurrogate):
-    """Mock surrogate model for testing purposes."""
-
-    def predict(
-        self,
-        y: np.ndarray,
-    ) -> dict[str, float]:
-        """Predict outputs based on input data."""
-        return dict(zip(self.outputs, y, strict=True))
-
-
 @dataclass(kw_only=True)
 class AbstractEstimator:
     """Abstract class for parameter estimation using neural networks."""