0.18.10

tests added

Author: bartzbeielstein

pyproject.toml

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

src/spotpython/data/lightdatamodule.py

@@ -143,30 +143,33 @@ def _setup_full_data_provided(self, stage) -> None:
 
         # Assign train/val datasets for use in dataloaders
         if stage == "fit" or stage is None:
-            if self.verbosity > 0:
-                print(f"train_size: {train_size}, val_size: {val_size} used for train & val data.")
             generator_fit = torch.Generator().manual_seed(self.test_seed)
             self.data_train, self.data_val, _ = random_split(self.data_full, [train_size, val_size, test_size], generator=generator_fit)
+            if self.verbosity > 0:
+                print(f"train_size: {train_size}, val_size: {val_size}, test_sie: {test_size} for splitting train & val data.")
+                print(f"train samples: {len(self.data_train)}, val samples: {len(self.data_val)} generated for train & val data.")
             # Handle scaling and transformation if scaler is provided
             if self.scaler is not None:
                 self.handle_scaling_and_transform()
 
         # 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.")
             generator_test = torch.Generator().manual_seed(self.test_seed)
             self.data_test, _, _ = random_split(self.data_full, [test_size, train_size, val_size], generator=generator_test)
+            if self.verbosity > 0:
+                print(f"train_size: {train_size}, val_size: {val_size}, test_sie: {test_size} for splitting test data.")
+                print(f"test samples: {len(self.data_test)} generated for test data.")
             if self.scaler is not None:
                 # 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.")
             generator_predict = torch.Generator().manual_seed(self.test_seed)
             self.data_predict, _, _ = random_split(self.data_full, [test_size, train_size, val_size], generator=generator_predict)
+            if self.verbosity > 0:
+                print(f"train_size: {train_size}, val_size: {val_size}, test_size (= predict_size): {test_size} for splitting predict data.")
+                print(f"predict samples: {len(self.data_predict)} generated for train & val data.")
             if self.scaler is not None:
                 # Transform the predict data
                 self.data_predict = self.transform_dataset(self.data_predict)

test/test_lightdatamodule_with_setup.py

@@ -0,0 +1,53 @@
+import pytest
+import torch
+from torch.utils.data import TensorDataset
+from lightning import seed_everything
+from spotpython.data.lightdatamodule import LightDataModule
+
+# Assuming LightDataModule is already imported from the provided code.
+
+# Define a mock scaler for testing purpose.
+class MockScaler:
+    def fit(self, data):
+        pass
+    
+    def transform(self, data):
+        return data
+
+# Define a simple dataset for testing.
+def create_mock_dataset(size=12):
+    data = torch.arange(size).float().view(-1, 1)
+    target = torch.arange(size).long()
+    return TensorDataset(data, target)
+
+# Test initialization and data splits
+@pytest.mark.parametrize("test_size, expected_train_size, expected_val_size, expected_test_size", [
+    (0.5, 3, 3, 6),  # Split 12 items into 3 train, 3 val, 6 test
+    (0.4, 5, 3, 5),  # Split 12 items into 5 train, 3 val, 5 test
+])
+def test_data_splitting(test_size, expected_train_size, expected_val_size, expected_test_size):
+    dataset = create_mock_dataset()
+    data_module = LightDataModule(dataset=dataset, batch_size=2, test_size=test_size, verbosity=1)
+    data_module.setup()
+
+    assert len(data_module.data_train) == expected_train_size
+    assert len(data_module.data_val) == expected_val_size
+    assert len(data_module.data_test) == expected_test_size
+
+
+# Test DataLoader
+def test_dataloader():
+    dataset = create_mock_dataset()
+    data_module = LightDataModule(dataset=dataset, batch_size=2, test_size=0.5, verbosity=1)
+    data_module.setup()
+
+    train_loader = data_module.train_dataloader()
+    val_loader = data_module.val_dataloader()
+    test_loader = data_module.test_dataloader()
+
+    assert len(train_loader.dataset) == len(data_module.data_train)
+    assert len(val_loader.dataset) == len(data_module.data_val)
+    assert len(test_loader.dataset) == len(data_module.data_test)
+
+if __name__ == "__main__":
+    pytest.main([__file__])