mxlpy 0.18.0__py3-none-any.whl → 0.20.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/parallel.py

@@ -24,11 +24,14 @@ 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
 
+__all__ = [
+    "Cache",
+    "parallelise",
+]
+
 
 def _pickle_name(k: Hashable) -> str:
     return f"{k}.p"

mxlpy/parameterise.py

@@ -1,11 +1,19 @@
 """Module to parameterise models."""
 
-from pathlib import Path
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
-import pandas as pd
 from parameteriser.brenda.v0 import Brenda
 
-__all__ = ["get_km_and_kcat_from_brenda"]
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    import pandas as pd
+
+__all__ = [
+    "get_km_and_kcat_from_brenda",
+]
 
 
 def get_km_and_kcat_from_brenda(