0.15.26

Cleanup lightdatamodule

Author: bartzbeielstein

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spotpython"
-version = "0.15.25"
+version = "0.15.26"
 authors = [
   { name="T. Bartz-Beielstein", email="tbb@bartzundbartz.de" }
 ]

src/spotpython/data/lightdatamodule.py

@@ -2,6 +2,7 @@
 import torch
 from torch.utils.data import DataLoader, random_split, TensorDataset
 from typing import Optional
+from spotpython.utils.split import calculate_data_split
 
 
 class LightDataModule(L.LightningDataModule):
@@ -95,6 +96,13 @@ def __init__(
         self.scaler = scaler
         self.verbosity = verbosity
 
+    def transform_dataset(self, dataset):
+        """Applies the scaler transformation to the dataset."""
+        transformed_data = [(self.scaler.transform(data), target) for data, target in dataset]
+        data_tensors = [data.clone().detach() for data, target in transformed_data]
+        target_tensors = [target.clone().detach() for data, target in transformed_data]
+        return TensorDataset(torch.stack(data_tensors).squeeze(1), torch.stack(target_tensors))
+
     def prepare_data(self) -> None:
         """Prepares the data for use."""
         # download
@@ -124,25 +132,12 @@ def setup(self, stage: Optional[str] = None) -> None:
                 Training set size: 3
 
         """
-        # if test_size is float, then train_size is 1 - test_size
-        test_size = self.test_size
-        if isinstance(self.test_size, float):
-            full_train_size = round(1.0 - test_size, 2)
-            val_size = round(full_train_size * test_size, 2)
-            train_size = round(full_train_size - val_size, 2)
-        else:
-            # if test_size is int, then train_size is len(data_full) - test_size
-            full_train_size = len(self.data_full) - test_size
-            val_size = int(full_train_size * test_size / len(self.data_full))
-            train_size = full_train_size - val_size
-
-        if self.verbosity > 0:
-            print(f"LightDataModule.setup(): stage: {stage}")
-        if self.verbosity > 1:
-            print(f"LightDataModule setup(): full_train_size: {full_train_size}")
-            print(f"LightDataModule setup(): val_size: {val_size}")
-            print(f"LightDataModule setup(): train_size: {train_size}")
-            print(f"LightDataModule setup(): test_size: {test_size}")
+        full_train_size, val_size, train_size, test_size = calculate_data_split(
+            test_size=self.test_size,
+            full_size=len(self.data_full),
+            verbosity=self.verbosity,
+            stage=stage,
+        )
 
         # Assign train/val datasets for use in dataloaders
         if stage == "fit" or stage is None:
@@ -153,64 +148,37 @@ def setup(self, stage: Optional[str] = None) -> None:
                 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
+                # Fit the scaler on training data
                 scaler_train_data = torch.stack([self.data_train[i][0] for i in range(len(self.data_train))]).squeeze(1)
-                # train_val_data = self.data_train[:,0]
                 if self.verbosity > 0:
                     print(scaler_train_data.shape)
                 self.scaler.fit(scaler_train_data)
-                self.data_train = [(self.scaler.transform(data), target) for data, target in self.data_train]
-                data_tensors_train = [data.clone().detach() for data, target in self.data_train]
-                target_tensors_train = [target.clone().detach() for data, target in self.data_train]
-                self.data_train = TensorDataset(
-                    torch.stack(data_tensors_train).squeeze(1), 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 = [data.clone().detach() for data, target in self.data_val]
-                target_tensors_val = [target.clone().detach() for data, target in self.data_val]
-                self.data_val = TensorDataset(torch.stack(data_tensors_val).squeeze(1), torch.stack(target_tensors_val))
+                # Transform the training data
+                self.data_train = self.transform_dataset(self.data_train)
+                # Transform the validation data
+                self.data_val = self.transform_dataset(self.data_val)
 
         # Assign test dataset for use in dataloader(s)
         if stage == "test" or stage is None:
             if self.verbosity > 0:
                 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 = [data.clone().detach() for data, target in self.data_test]
-                target_tensors_test = [target.clone().detach() for data, target in self.data_test]
-                self.data_test = TensorDataset(
-                    torch.stack(data_tensors_test).squeeze(1), 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)
-        #     full_data_predict, _ = random_split(
-        #         self.data_full, [test_size, full_train_size], generator=generator_predict
-        #     )
-        #     # Only keep the features for prediction
-        #     self.data_predict = [x for x, _ in full_data_predict]
+                # Transform the test data
+                self.data_test = self.transform_dataset(self.data_test)
 
         # Assign pred dataset for use in dataloader(s)
         if stage == "predict" or stage is None:
             if self.verbosity > 0:
                 print(f"test_size: {test_size} used for predict dataset.")
-            # get test data set as test_abs percent of the full dataset
             generator_predict = torch.Generator().manual_seed(self.test_seed)
             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 = [data.clone().detach() for data, target in self.data_predict]
-                target_tensors_predict = [target.clone().detach() for data, target in self.data_predict]
-                self.data_predict = TensorDataset(
-                    torch.stack(data_tensors_predict).squeeze(1), torch.stack(target_tensors_predict)
-                )
+                # Transform the predict data
+                self.data_predict = self.transform_dataset(self.data_predict)
 
     def train_dataloader(self) -> DataLoader:
         """
@@ -235,7 +203,6 @@ def train_dataloader(self) -> DataLoader:
             print(f"LightDataModule.train_dataloader(). data_train size: {len(self.data_train)}")
         # print(f"LightDataModule: train_dataloader(). batch_size: {self.batch_size}")
         # print(f"LightDataModule: train_dataloader(). num_workers: {self.num_workers}")
-        # apply fit_transform to the training data
         return DataLoader(self.data_train, batch_size=self.batch_size, num_workers=self.num_workers)
 
     def val_dataloader(self) -> DataLoader:
@@ -260,7 +227,6 @@ def val_dataloader(self) -> DataLoader:
             print(f"LightDataModule.val_dataloader(). Val. set size: {len(self.data_val)}")
         # print(f"LightDataModule: val_dataloader(). batch_size: {self.batch_size}")
         # print(f"LightDataModule: val_dataloader(). num_workers: {self.num_workers}")
-        # apply fit_transform to the val data
         return DataLoader(self.data_val, batch_size=self.batch_size, num_workers=self.num_workers)
 
     def test_dataloader(self) -> DataLoader:
@@ -312,6 +278,4 @@ def predict_dataloader(self) -> DataLoader:
             print(f"LightDataModule.predict_dataloader(). Predict set size: {len(self.data_predict)}")
         # print(f"LightDataModule: predict_dataloader(). batch_size: {self.batch_size}")
         # print(f"LightDataModule: predict_dataloader(). num_workers: {self.num_workers}")
-        # apply fit_transform to the val data
-
         return DataLoader(self.data_predict, batch_size=len(self.data_predict), num_workers=self.num_workers)

src/spotpython/utils/split.py

@@ -0,0 +1,37 @@
+def calculate_data_split(test_size, full_size, verbosity=0, stage=None) -> tuple:
+    """
+    Calculates the split sizes for training, validation, and test datasets.
+
+    Args:
+        test_size (float or int):
+            The size of the test set.
+            Can be a float for proportion or an int for absolute number of test samples.
+        full_size (int):
+            The size of the full dataset.
+        verbosity (int, optional):
+            The level of verbosity for debug output. Defaults to 0.
+        stage (str, optional):
+            The stage of setup, for debug output if needed.
+
+    Returns:
+        tuple: A tuple containing the sizes (full_train_size, val_size, train_size, test_size).
+    """
+    if isinstance(test_size, float):
+        full_train_size = round(1.0 - test_size, 2)
+        val_size = round(full_train_size * test_size, 2)
+        train_size = round(full_train_size - val_size, 2)
+    else:
+        # test_size is considered an int, training size calculation directly based on it
+        full_train_size = full_size - test_size
+        val_size = int(full_train_size * test_size / full_size)
+        train_size = full_train_size - val_size
+
+    if verbosity > 0:
+        print(f"stage: {stage}")
+    if verbosity > 1:
+        print(f"full_sizefull_train_size: {full_train_size}")
+        print(f"full_sizeval_size: {val_size}")
+        print(f"full_sizetrain_size: {train_size}")
+        print(f"full_sizetest_size: {test_size}")
+
+    return full_train_size, val_size, train_size, test_size

test/test_calculate_split.py

@@ -0,0 +1,43 @@
+from spotpython.utils.split import calculate_data_split
+import pytest
+
+def test_calculate_data_split_float():
+    full_size = 100
+    test_size = 0.2
+    expected_full_train_size = 0.8
+    expected_val_size = 0.16
+    expected_train_size = 0.64
+
+    result = calculate_data_split(test_size, full_size)
+
+    assert result == (expected_full_train_size, expected_val_size, expected_train_size, test_size), \
+           f"Result was {result}, expected {(expected_full_train_size, expected_val_size, expected_train_size, test_size)}"
+
+def test_calculate_data_split_int():
+    full_size = 100
+    test_size = 20
+    expected_full_train_size = 80
+    expected_val_size = 16  # Calculated as 80 * 20 / 100
+    expected_train_size = 64  # 80 - 16
+
+    result = calculate_data_split(test_size, full_size)
+
+    assert result == (expected_full_train_size, expected_val_size, expected_train_size, test_size), \
+           f"Result was {result}, expected {(expected_full_train_size, expected_val_size, expected_train_size, test_size)}"
+
+def test_calculate_data_split_verbosity():
+    full_size = 100
+    test_size = 0.2
+
+    # Ideally, we'd capture the output here as well
+    # For now, we just confirm it runs without error
+    result = calculate_data_split(test_size, full_size, verbosity=2, stage='test')
+
+    expected_full_train_size = 0.8
+    expected_val_size = 0.16
+    expected_train_size = 0.64
+
+    assert result == (expected_full_train_size, expected_val_size, expected_train_size, test_size)
+
+if __name__ == "__main__":
+    pytest.main()

test/test_transform_dataset.py

@@ -0,0 +1,55 @@
+import pytest
+import torch
+from torch.utils.data import TensorDataset
+from unittest.mock import MagicMock
+
+# Assuming the class containing transform_dataset is named MyDataModule
+class MyDataModule:
+    def __init__(self, scaler):
+        self.scaler = scaler
+
+    def transform_dataset(self, dataset):
+        transformed_data = [(self.scaler.transform(data), target) for data, target in dataset]
+        data_tensors = [data.clone().detach() for data, target in transformed_data]
+        target_tensors = [target.clone().detach() for data, target in transformed_data]
+        return TensorDataset(torch.stack(data_tensors).squeeze(1), torch.stack(target_tensors))
+
+# Test function for transform_dataset
+@pytest.fixture
+def setup_data():
+    # Mock dataset
+    input_data = torch.randn(3, 4)  # Mock input data
+    target_data = torch.tensor([0, 1, 2])  # Mock target data
+
+    dataset = [(input_data[i], target_data[i]) for i in range(len(target_data))]
+    
+    # Mock scaler with a simple transform logic
+    mock_scaler = MagicMock()
+    mock_scaler.transform = lambda x: 2 * x  # Example transformation: multiply by 2
+
+    return mock_scaler, dataset
+
+def test_transform_dataset(setup_data):
+    mock_scaler, dataset = setup_data
+    data_module = MyDataModule(mock_scaler)
+    
+    transformed_dataset = data_module.transform_dataset(dataset)
+
+    # Check that transform_dataset returns a TensorDataset
+    assert isinstance(transformed_dataset, TensorDataset)
+
+    # Extract transformed data and targets
+    transformed_data, transformed_targets = transformed_dataset.tensors
+
+    # Verify the shape
+    assert transformed_data.shape == torch.Size([3, 4])
+    assert transformed_targets.shape == torch.Size([3])
+
+    # Verify that the data was transformed correctly (i.e., multiplied by 2)
+    expected_data = torch.stack([mock_scaler.transform(d[0]) for d in dataset]).squeeze(1)
+    for td, ed in zip(transformed_data, expected_data):
+        assert torch.allclose(td, ed)
+
+    # Verify that the targets were unchanged
+    expected_targets = torch.tensor([d[1] for d in dataset])
+    assert torch.equal(transformed_targets, expected_targets)