scaler options

Author: ahinterl94-th

src/spotPython/data/lightdatamodule.py

@@ -1,9 +1,10 @@
 import lightning as L
 import torch
-from torch.utils.data import DataLoader, random_split
+from torch.utils.data import DataLoader, random_split, TensorDataset
 from typing import Optional
 
 
+
 class LightDataModule(L.LightningDataModule):
     """
     A LightningDataModule for handling data.
@@ -25,6 +26,8 @@ class LightDataModule(L.LightningDataModule):
             The test seed. Defaults to 42.
         num_workers (int):
             The number of workers. Defaults to 0.
+        scaler (object):
+            The spot scaler object (e.g. TorchStandardScaler). Defaults to None.
 
     Attributes:
         batch_size (int): The batch size.
@@ -79,13 +82,16 @@ def __init__(
         test_size: float,
         test_seed: int = 42,
         num_workers: int = 0,
+        scaler: Optional[object] = None,
     ):
         super().__init__()
         self.batch_size = batch_size
         self.data_full = dataset
         self.test_size = test_size
         self.test_seed = test_seed
         self.num_workers = num_workers
+        self.scaler = scaler
+
 
     def prepare_data(self) -> None:
         """Prepares the data for use."""
@@ -98,6 +104,7 @@ def setup(self, stage: Optional[str] = None) -> None:
         Uses torch.utils.data.random_split() to split the data.
         Splitting is based on the test_size and test_seed.
         The test_size can be a float or an int.
+        If a spotPython scaler object is defined, the data will be scaled.
 
         Args:
             stage (Optional[str]):
@@ -140,14 +147,32 @@ def setup(self, stage: Optional[str] = None) -> None:
             self.data_train, self.data_val, _ = random_split(
                 self.data_full, [train_size, val_size, test_size], generator=generator_fit
             )
-
+            if self.scaler is not None:
+                # Fit the scaler on training data and transform both train and val data
+                train_val_data = torch.cat([self.data_train[i][0] for i in range(len(self.data_train))])
+                self.scaler.fit(train_val_data)
+                self.data_train = [(self.scaler.transform(data), target) for data, target in self.data_train]
+                data_tensors_train = [torch.tensor(data, dtype=torch.float32) for data, target in self.data_train]
+                target_tensors_train = [torch.tensor(target, dtype=torch.float32) for data, target in self.data_train]
+                self.data_train = TensorDataset(torch.stack(data_tensors_train), torch.stack(target_tensors_train))
+                #print(self.data_train)
+                self.data_val = [(self.scaler.transform(data), target) for data, target in self.data_val]
+                data_tensors_val = [torch.tensor(data, dtype=torch.float32) for data, target in self.data_val]
+                target_tensors_val = [torch.tensor(target, dtype=torch.float32) for data, target in self.data_val]
+                self.data_val = TensorDataset(torch.stack(data_tensors_val), torch.stack(target_tensors_val))
+               
         # Assign test dataset for use in dataloader(s)
         if stage == "test" or stage is None:
             print(f"test_size: {test_size} used for test dataset.")
             # get test data set as test_abs percent of the full dataset
             generator_test = torch.Generator().manual_seed(self.test_seed)
             self.data_test, _ = random_split(self.data_full, [test_size, full_train_size], generator=generator_test)
-
+            if self.scaler is not None:
+                self.data_test = [(self.scaler.transform(data), target) for data, target in self.data_test]
+                data_tensors_test = [torch.tensor(data, dtype=torch.float32) for data, target in self.data_test]
+                target_tensors_test = [torch.tensor(target, dtype=torch.float32) for data, target in self.data_test]
+                self.data_test = TensorDataset(torch.stack(data_tensors_test), torch.stack(target_tensors_test))
+               
         # if stage == "predict" or stage is None:
         #     print(f"test_size, full_train_size: {test_size}, {full_train_size}")
         #     generator_predict = torch.Generator().manual_seed(self.test_seed)
@@ -165,6 +190,11 @@ def setup(self, stage: Optional[str] = None) -> None:
             self.data_predict, _ = random_split(
                 self.data_full, [test_size, full_train_size], generator=generator_predict
             )
+            if self.scaler is not None:
+                self.data_predict = [(self.scaler.transform(data), target) for data, target in self.data_predict]
+                data_tensors_predict= [torch.tensor(data, dtype=torch.float32) for data, target in self.data_predict]
+                target_tensors_predict = [torch.tensor(target, dtype=torch.float32) for data, target in self.data_predict]
+                self.data_predict = TensorDataset(torch.stack(data_tensors_predict), torch.stack(target_tensors_predict))
 
     def train_dataloader(self) -> DataLoader:
         """
@@ -265,3 +295,5 @@ def predict_dataloader(self) -> DataLoader:
         # apply fit_transform to the val data
         return DataLoader(self.data_test, batch_size=self.batch_size, num_workers=self.num_workers)
         return DataLoader(self.data_predict, batch_size=len(self.data_predict), num_workers=self.num_workers)
+
+

src/spotPython/utils/init.py

@@ -407,6 +407,7 @@ def fun_control_init(
         "prep_model_name": prep_model_name,
         "progress_file": progress_file,
         "save_model": False,
+        "scaler":scaler,
         "scenario": scenario,
         "seed": seed,
         "show_batch_interval": 1_000_000,

src/spotPython/utils/scaler.py

@@ -0,0 +1,100 @@
+import torch
+
+class TorchStandardScaler:
+    """
+    A class for scaling data using standardization with torch tensors.
+    """
+    def fit(self, x):
+        """
+        Compute the mean and standard deviation of the input tensor.
+        Args:
+            x (torch.Tensor): The input tensor.
+        Raises:
+            TypeError: If the input is not a torch tensor.
+        """
+        if not torch.is_tensor(x):
+            raise TypeError("Input should be a torch tensor")
+        self.mean = x.mean(0, keepdim=True)
+        self.std = x.std(0, unbiased=False, keepdim=True)
+    
+    def transform(self, x):
+        """
+        Scale the input tensor using the computed mean and standard deviation.
+        Args:
+            x (torch.Tensor): The input tensor.
+        Returns:
+            torch.Tensor: The scaled tensor.
+        Raises:
+            TypeError: If the input is not a torch tensor.
+            RuntimeError: If the scaler has not been fitted before transforming data.
+        """
+        if not torch.is_tensor(x):
+            raise TypeError("Input should be a torch tensor")
+        if not hasattr(self, 'mean') or not hasattr(self, 'std'):
+            raise RuntimeError("Must fit scaler before transforming data")
+        x = (x - self.mean) / (self.std + 1e-7)
+        return x
+    
+    def fit_transform(self, x):
+        """
+        Fit the scaler to the input tensor and then scale the tensor.
+        Args:
+            x (torch.Tensor): The input tensor.
+        Returns:
+            torch.Tensor: The scaled tensor.
+        """
+        self.fit(x)
+        return self.transform(x)
+
+
+
+class TorchMinMaxScaler:
+    """
+    A class for scaling data using min-max normalization with PyTorch tensors.
+    """
+    def fit(self, x):
+        """
+        Fit the scaler to the input data.
+        Parameters:
+        - x: torch.Tensor
+            The input data to fit the scaler to.
+        Raises:
+        - TypeError: If the input is not a torch tensor.
+        """
+        if not torch.is_tensor(x):
+            raise TypeError("Input should be a torch tensor")
+        self.min = x.min(0, keepdim=True).values
+        self.max = x.max(0, keepdim=True).values
+    
+    def transform(self, x):
+        """
+        Transform the input data using the fitted scaler.
+        Parameters:
+        - x: torch.Tensor
+            The input data to transform.
+        Returns:
+        - torch.Tensor: The transformed data.
+        Raises:
+        - TypeError: If the input is not a torch tensor.
+        - RuntimeError: If the scaler has not been fitted before transforming data.
+        """
+        if not torch.is_tensor(x):
+            raise TypeError("Input should be a torch tensor")
+        if not hasattr(self, 'min') or not hasattr(self, 'max'):
+            raise RuntimeError("Must fit scaler before transforming data")
+        x = (x - self.min) / (self.max - self.min + 1e-7)
+        return x
+    
+    def fit_transform(self, x):
+        """
+        Fit the scaler to the input data and transform it.
+        Parameters:
+        - x: torch.Tensor
+            The input data to fit and transform.
+        Returns:
+        - torch.Tensor: The transformed data.
+        Raises:
+        - TypeError: If the input is not a torch tensor.
+        """
+        self.fit(x)
+        return self.transform(x)