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

mxlpy/npe/_torch.py

@@ -12,18 +12,15 @@ Functions:
 from __future__ import annotations
 
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING, Self, cast
 
 import numpy as np
 import pandas as pd
 import torch
-import tqdm
 from torch import nn
 from torch.optim.adam import Adam
 
-from mxlpy.nn._torch import LSTM, MLP, DefaultDevice
-from mxlpy.parallel import Cache
+from mxlpy.nn._torch import LSTM, MLP, DefaultDevice, train
 from mxlpy.types import AbstractEstimator
 
 if TYPE_CHECKING:
@@ -31,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",
 ]
 
 
@@ -62,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
@@ -75,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
@@ -102,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
@@ -117,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:
@@ -127,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
 
@@ -136,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,28 +156,16 @@ class TorchSteadyStateTrainer:
             batch_size: Size of mini-batches for training (None for full-batch)
 
         """
-        features = torch.Tensor(self.features.to_numpy(), device=self.device)
-        targets = torch.Tensor(self.targets.to_numpy(), device=self.device)
-
-        if batch_size is None:
-            losses = _train_full(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                loss_fn=self.loss_fn,
-            )
-        else:
-            losses = _train_batched(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                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,
+            loss_fn=self.loss_fn,
+            device=self.device,
+        )
 
         if len(self.losses) > 0:
             losses.index += self.losses[-1].index[-1]
@@ -193,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
@@ -217,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:
@@ -227,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
 
@@ -236,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 = []
@@ -260,37 +243,22 @@ class TorchTimeCourseTrainer:
             batch_size: Size of mini-batches for training (None for full-batch)
 
         """
-        features = torch.Tensor(
-            np.swapaxes(
+        losses = train(
+            model=self.model,
+            features=np.swapaxes(
                 self.features.to_numpy().reshape(
                     (len(self.targets), -1, len(self.features.columns))
                 ),
                 axis1=0,
                 axis2=1,
             ),
+            targets=self.targets.to_numpy(),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            loss_fn=self.loss_fn,
             device=self.device,
         )
-        targets = torch.Tensor(self.targets.to_numpy(), device=self.device)
-
-        if batch_size is None:
-            losses = _train_full(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                loss_fn=self.loss_fn,
-            )
-        else:
-            losses = _train_batched(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                batch_size=batch_size,
-                loss_fn=self.loss_fn,
-            )
 
         if len(self.losses) > 0:
             losses.index += self.losses[-1].index[-1]
@@ -301,66 +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_batched(
-    approximator: nn.Module,
-    features: torch.Tensor,
-    targets: torch.Tensor,
-    epochs: int,
-    optimizer: Adam,
-    batch_size: int,
-    loss_fn: LossFn,
-) -> pd.Series:
-    losses = {}
-    for epoch in tqdm.trange(epochs):
-        permutation = torch.randperm(features.size()[0])
-        epoch_loss = 0
-        for i in range(0, features.size()[0], batch_size):
-            optimizer.zero_grad()
-            indices = permutation[i : i + batch_size]
-            loss = loss_fn(approximator(features[indices]), targets[indices])
-            loss.backward()
-            optimizer.step()
-            epoch_loss += loss.detach().numpy()
-
-        losses[epoch] = epoch_loss / (features.size()[0] / batch_size)
-    return pd.Series(losses, dtype=float)
-
-
-def _train_full(
-    approximator: nn.Module,
-    features: torch.Tensor,
-    targets: torch.Tensor,
-    epochs: int,
-    optimizer: Adam,
-    loss_fn: LossFn,
-) -> pd.Series:
-    losses = {}
-    for i in tqdm.trange(epochs):
-        optimizer.zero_grad()
-        loss = loss_fn(approximator(features), targets)
-        loss.backward()
-        optimizer.step()
-        losses[i] = loss.detach().numpy()
-    return pd.Series(losses, dtype=float)
-
-
-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
@@ -375,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
@@ -417,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)