Add docstrings for converters and better tests

Author: jacobbieker

ocf_datapipes/convert/coordinates.py

@@ -1,5 +1,7 @@
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterDataPipe
+import xarray as xr
+from typing import Union
 
 from ocf_datapipes.utils.geospatial import (
     lat_lon_to_osgb,
@@ -10,10 +12,18 @@
 
 @functional_datapipe("convert_latlon_to_osgb")
 class ConvertLatLonToOSGBIterDataPipe(IterDataPipe):
+    """Convert from Lat/Lon object to OSGB"""
     def __init__(self, source_datapipe: IterDataPipe):
+        """
+        Convert from Lat/Lon to OSGB
+
+        Args:
+            source_datapipe: Datapipe emitting Xarray objects with latitude and longitude data
+        """
         self.source_datapipe = source_datapipe
 
-    def __iter__(self):
+    def __iter__(self) -> Union[xr.DataArray, xr.Dataset]:
+        """Convert from Lat/Lon to OSGB"""
         for xr_data in self.source_datapipe:
             xr_data["x_osgb"], xr_data["y_osgb"] = lat_lon_to_osgb(
                 latitude=xr_data["latitude"], longitude=xr_data["longitude"]
@@ -23,10 +33,18 @@ def __iter__(self):
 
 @functional_datapipe("convert_osgb_to_latlon")
 class ConvertOSGBToLatLonIterDataPipe(IterDataPipe):
+    """Convert from OSGB to Lat/Lon"""
     def __init__(self, source_datapipe: IterDataPipe):
+        """
+        Convert from OSGB to Lat/Lon
+
+        Args:
+            source_datapipe: Datapipe emitting Xarray objects with OSGB data
+        """
         self.source_datapipe = source_datapipe
 
-    def __iter__(self):
+    def __iter__(self) -> Union[xr.DataArray, xr.Dataset]:
+        """Convert and add lat/lon to Xarray object"""
         for xr_data in self.source_datapipe:
             xr_data["latitude"], xr_data["longitude"] = osgb_to_lat_lon(
                 x=xr_data["x_osgb"], y=xr_data["y_osgb"]
@@ -36,10 +54,18 @@ def __iter__(self):
 
 @functional_datapipe("convert_geostationary_to_latlon")
 class ConvertGeostationaryToLatLonIterDataPipe(IterDataPipe):
+    """Convert from geostationary to Lat/Lon points"""
     def __init__(self, source_datapipe: IterDataPipe):
+        """
+        Convert from Geostationary to Lat/Lon points and add to Xarray object
+
+        Args:
+            source_datapipe: Datapipe emitting Xarray object with geostationary points
+        """
         self.source_datapipe = source_datapipe
 
-    def __iter__(self):
+    def __iter__(self) -> Union[xr.DataArray, xr.Dataset]:
+        """Convert from geostationary to Lat/Lon and yield the Xarray object"""
         for xr_data in self.source_datapipe:
             transform = load_geostationary_area_definition_and_transform_latlon(xr_data)
             xr_data["latitude"], xr_data["longitude"] = transform(

ocf_datapipes/convert/gsp.py

@@ -7,12 +7,20 @@
 
 @functional_datapipe("convert_gsp_to_numpy_batch")
 class ConvertGSPToNumpyBatchIterDataPipe(IterDataPipe):
-    def __init__(self, source_dp: IterDataPipe):
+    """Convert GSP Xarray to NumpyBatch"""
+    def __init__(self, source_datapipe: IterDataPipe):
+        """
+        Convert GSP Xarray to NumpyBatch object
+
+        Args:
+            source_datapipe: Datapipe emitting GSP Xarray object
+        """
         super().__init__()
-        self.source_dp = source_dp
+        self.source_datapipe = source_datapipe
 
     def __iter__(self) -> NumpyBatch:
-        for xr_data in self.source_dp:
+        """Convert from Xarray to NumpyBatch"""
+        for xr_data in self.source_datapipe:
             example: NumpyBatch = {
                 BatchKey.gsp: xr_data.values,
                 BatchKey.gsp_t0_idx: xr_data.attrs["t0_idx"],

ocf_datapipes/convert/nwp.py

@@ -8,12 +8,20 @@
 
 @functional_datapipe("convert_nwp_to_numpy_batch")
 class ConvertNWPToNumpyBatchIterDataPipe(IterDataPipe):
-    def __init__(self, source_dp: IterDataPipe):
+    """Convert NWP Xarray objects to NumpyBatch ones"""
+    def __init__(self, source_datapipe: IterDataPipe):
+        """
+        Convert NWP Xarray objecs to NumpyBatch ones
+
+        Args:
+            source_datapipe: Datapipe emitting NWP Xarray objects
+        """
         super().__init__()
-        self.source_dp = source_dp
+        self.source_datapipe = source_datapipe
 
     def __iter__(self) -> NumpyBatch:
-        for xr_data in self.source_dp:
+        """Convert from Xarray to NumpyBatch"""
+        for xr_data in self.source_datapipe:
             example: NumpyBatch = {
                 BatchKey.nwp: xr_data.values,
                 BatchKey.nwp_t0_idx: xr_data.attrs["t0_idx"],

ocf_datapipes/convert/pv.py

@@ -8,12 +8,20 @@
 
 @functional_datapipe("convert_pv_to_numpy_batch")
 class ConvertPVToNumpyBatchIterDataPipe(IterDataPipe):
-    def __init__(self, source_dp: IterDataPipe):
+    """Convert PV Xarray to NumpyBatch"""
+    def __init__(self, source_datapipe: IterDataPipe):
+        """
+        Convert PV Xarray objects to NumpyBatch objects
+
+        Args:
+            source_datapipe: Datapipe emitting PV Xarray objects
+        """
         super().__init__()
-        self.source_dp = source_dp
+        self.source_datapipe = source_datapipe
 
     def __iter__(self) -> NumpyBatch:
-        for xr_data in self.source_dp:
+        """Iterate and convert PV Xarray to NumpyBatch"""
+        for xr_data in self.source_datapipe:
             example: NumpyBatch = {
                 BatchKey.pv: xr_data.values,
                 BatchKey.pv_t0_idx: xr_data.attrs["t0_idx"],

ocf_datapipes/convert/satellite.py

@@ -7,13 +7,22 @@
 
 @functional_datapipe("convert_satellite_to_numpy_batch")
 class ConvertSatelliteToNumpyBatchIterDataPipe(IterDataPipe):
-    def __init__(self, source_dp: IterDataPipe, is_hrv: bool = False):
+    """Converts Xarray Satellite to NumpyBatch object"""
+    def __init__(self, source_datapipe: IterDataPipe, is_hrv: bool = False):
+        """
+        Converts Xarray satellite object to NumpyBatch object
+
+        Args:
+            source_datapipe: Datapipe emitting Xarray satellite objects
+            is_hrv: Whether this is HRV satellite data or non-HRV data
+        """
         super().__init__()
-        self.source_dp = source_dp
+        self.source_datapipe = source_datapipe
         self.is_hrv = is_hrv
 
     def __iter__(self) -> NumpyBatch:
-        for xr_data in self.source_dp:
+        """Convert each example to a NumpyBatch object"""
+        for xr_data in self.source_datapipe:
             if self.is_hrv:
                 example: NumpyBatch = {
                     BatchKey.hrvsatellite_actual: xr_data.values,

tests/convert/test_gsp.py

@@ -2,12 +2,13 @@
 
 from ocf_datapipes.convert import ConvertGSPToNumpyBatch
 from ocf_datapipes.transform.xarray import AddT0IdxAndSamplePeriodDuration
-
+from ocf_datapipes.utils.consts import BatchKey
 
 def test_convert_gsp_to_numpy_batch(gsp_dp):
     gsp_dp = AddT0IdxAndSamplePeriodDuration(
         gsp_dp, sample_period_duration=timedelta(minutes=5), history_duration=timedelta(minutes=60)
     )
     gsp_dp = ConvertGSPToNumpyBatch(gsp_dp)
     data = next(iter(gsp_dp))
-    assert data is not None
+    assert BatchKey.gsp in data
+    assert BatchKey.gsp_id in data

tests/convert/test_nwp.py

@@ -3,6 +3,7 @@
 from ocf_datapipes.convert import ConvertNWPToNumpyBatch
 from ocf_datapipes.select import SelectLiveT0Time
 from ocf_datapipes.transform.xarray import AddT0IdxAndSamplePeriodDuration, ConvertToNWPTargetTime
+from ocf_datapipes.utils.consts import BatchKey
 
 
 def test_convert_nwp_to_numpy_batch(nwp_dp):
@@ -20,4 +21,5 @@ def test_convert_nwp_to_numpy_batch(nwp_dp):
     )
     nwp_dp = ConvertNWPToNumpyBatch(nwp_dp)
     data = next(iter(nwp_dp))
-    assert data is not None
+    assert BatchKey.nwp in data
+    assert BatchKey.nwp_channel_names in data

tests/convert/test_pv.py

@@ -2,6 +2,7 @@
 
 from ocf_datapipes.convert import ConvertPVToNumpyBatch
 from ocf_datapipes.transform.xarray import AddT0IdxAndSamplePeriodDuration
+from ocf_datapipes.utils.consts import BatchKey
 
 
 def test_convert_passiv_to_numpy_batch(passiv_dp):
@@ -12,7 +13,8 @@ def test_convert_passiv_to_numpy_batch(passiv_dp):
     )
     passiv_dp = ConvertPVToNumpyBatch(passiv_dp)
     data = next(iter(passiv_dp))
-    assert data is not None
+    assert BatchKey.pv in data
+    assert BatchKey.pv_t0_idx in data
 
 
 def test_convert_pvoutput_to_numpy_batch(pvoutput_dp):
@@ -25,4 +27,5 @@ def test_convert_pvoutput_to_numpy_batch(pvoutput_dp):
     pvoutput_dp = ConvertPVToNumpyBatch(pvoutput_dp)
 
     data = next(iter(pvoutput_dp))
-    assert data is not None
+    assert BatchKey.pv in data
+    assert BatchKey.pv_t0_idx in data

tests/convert/test_satellite.py

@@ -2,6 +2,7 @@
 
 from ocf_datapipes.convert import ConvertSatelliteToNumpyBatch
 from ocf_datapipes.transform.xarray import AddT0IdxAndSamplePeriodDuration
+from ocf_datapipes.utils.consts import BatchKey
 
 
 def test_convert_satellite_to_numpy_batch(sat_dp):
@@ -11,7 +12,10 @@ def test_convert_satellite_to_numpy_batch(sat_dp):
     )
     sat_dp = ConvertSatelliteToNumpyBatch(sat_dp, is_hrv=False)
     data = next(iter(sat_dp))
-    assert data is not None
+    assert BatchKey.satellite_actual in data
+    assert BatchKey.satellite_t0_idx in data
+    assert BatchKey.hrvsatellite_actual not in data
+    assert BatchKey.hrvsatellite_t0_idx not in data
 
 
 def test_convert_hrvsatellite_to_numpy_batch(sat_dp):
@@ -20,4 +24,7 @@ def test_convert_hrvsatellite_to_numpy_batch(sat_dp):
     )
     sat_dp = ConvertSatelliteToNumpyBatch(sat_dp, is_hrv=True)
     data = next(iter(sat_dp))
-    assert data is not None
+    assert BatchKey.hrvsatellite_actual in data
+    assert BatchKey.hrvsatellite_t0_idx in data
+    assert BatchKey.satellite_actual not in data
+    assert BatchKey.satellite_t0_idx not in data