Data sampler updates (#316)

Data sampler migration related updates

Author: felix-e-h-p

pvnet/data/__init__.py

@@ -1,4 +1,3 @@
 """Data parts"""
 from .site_datamodule import SiteDataModule
 from .uk_regional_datamodule import DataModule
-from .utils import BatchSplitter

pvnet/data/base_datamodule.py

@@ -4,9 +4,9 @@
 
 from lightning.pytorch import LightningDataModule
 from ocf_data_sampler.numpy_sample.collate import stack_np_samples_into_batch
-from ocf_data_sampler.sample.base import SampleBase
-from ocf_datapipes.batch import (
+from ocf_data_sampler.sample.base import (
     NumpyBatch,
+    SampleBase,
     TensorBatch,
     batch_to_tensor,
 )

pvnet/data/site_datamodule.py

@@ -1,16 +1,14 @@
 """ Data module for pytorch lightning """
 
 from ocf_data_sampler.sample.site import SiteSample
-from ocf_data_sampler.torch_datasets.datasets.site import (
-    SitesDataset,
-)
+from ocf_data_sampler.torch_datasets.datasets.site import SitesDataset
 from torch.utils.data import Dataset
 
 from pvnet.data.base_datamodule import BaseDataModule, PremadeSamplesDataset
 
 
 class SiteDataModule(BaseDataModule):
-    """Datamodule for training pvnet and using pvnet pipeline in `ocf_datapipes`."""
+    """Datamodule for training pvnet and using pvnet pipeline in `ocf-data-sampler`."""
 
     def __init__(
         self,

pvnet/data/utils.py

@@ -13,7 +13,7 @@
 import xarray as xr
 from lightning.pytorch.loggers import Logger
 from lightning.pytorch.utilities import rank_zero_only
-from ocf_datapipes.utils import Location
+from ocf_data_sampler.select.location import Location
 from omegaconf import DictConfig, OmegaConf
 
 

pvnet/utils.py

@@ -6,7 +6,7 @@ dynamic = ["version", "readme"]
 license={file="LICENCE"}
 
 dependencies = [
-    "ocf_data_sampler==0.1.2",
+    "ocf_data_sampler==0.1.7",
     "ocf_datapipes>=3.3.34",
     "ocf_ml_metrics>=0.0.11",
     "numpy",

pyproject.toml

@@ -26,8 +26,6 @@ def test_iter():
         val_period=[None, None],
     )
 
-    batch = next(iter(dm.train_dataloader()))
-
 
 def test_iter_multiprocessing():
     dm = DataModule(
@@ -55,12 +53,12 @@ def test_iter_multiprocessing():
 def test_site_init_sample_dir():
     dm = SiteDataModule(
         configuration=None,
+        sample_dir="tests/test_data/presaved_site_samples",
         batch_size=2,
         num_workers=0,
         prefetch_factor=None,
         train_period=[None, None],
         val_period=[None, None],
-        sample_dir="tests/test_data/presaved_site_samples",
     )
 
 
@@ -74,21 +72,3 @@ def test_site_init_config():
         val_period=[None, None],
         sample_dir=None,
     )
-
-
-def test_worker_configuration():
-    dm = DataModule(
-        sample_dir="tests/test_data/presaved_samples_uk_regional",
-        batch_size=2,
-        num_workers=4,
-        prefetch_factor=2,
-    )
-
-    # Iterate through dataloader - assert multi-processing functions fine
-    batches_processed = 0
-    for batch in dm.train_dataloader():
-        batches_processed += 1
-        if batches_processed >= 5:
-            break
-
-    assert batches_processed > 0, "No batches were processed"