PVNet concurrent datapipe

ocf_datapipes/convert/numpy_batch/__init__.py

@@ -1 +1,7 @@
 """Conversion from Xarray to NumpyBatch"""
+from .gsp import convert_gsp_to_numpy_batch
+from .nwp import convert_nwp_to_numpy_batch
+from .pv import convert_pv_to_numpy_batch
+from .satellite import convert_satellite_to_numpy_batch
+from .sensor import convert_sensor_to_numpy_batch
+from .wind import convert_wind_to_numpy_batch

ocf_datapipes/convert/numpy_batch/gsp.py

@@ -9,6 +9,31 @@
 logger = logging.getLogger(__name__)
 
 
+def convert_gsp_to_numpy_batch(xr_data):
+    """Convert from Xarray to NumpyBatch"""
+
+    example: NumpyBatch = {
+        BatchKey.gsp: xr_data.values,
+        BatchKey.gsp_t0_idx: xr_data.attrs["t0_idx"],
+        BatchKey.gsp_id: xr_data.gsp_id.values,
+        BatchKey.gsp_nominal_capacity_mwp: xr_data.isel(time_utc=0)["nominal_capacity_mwp"].values,
+        BatchKey.gsp_effective_capacity_mwp: (
+            xr_data.isel(time_utc=0)["effective_capacity_mwp"].values
+        ),
+        BatchKey.gsp_time_utc: datetime64_to_float(xr_data["time_utc"].values),
+    }
+
+    # Coordinates
+    for batch_key, dataset_key in (
+        (BatchKey.gsp_y_osgb, "y_osgb"),
+        (BatchKey.gsp_x_osgb, "x_osgb"),
+    ):
+        if dataset_key in xr_data.coords.keys():
+            example[batch_key] = xr_data[dataset_key].values
+
+    return example
+
+
 @functional_datapipe("convert_gsp_to_numpy_batch")
 class ConvertGSPToNumpyBatchIterDataPipe(IterDataPipe):
     """Convert GSP Xarray to NumpyBatch"""
@@ -25,29 +50,5 @@ def __init__(self, source_datapipe: IterDataPipe):
 
     def __iter__(self) -> NumpyBatch:
         """Convert from Xarray to NumpyBatch"""
-        logger.debug("Converting GSP to numpy to batch")
         for xr_data in self.source_datapipe:
-            example: NumpyBatch = {
-                BatchKey.gsp: xr_data.values,
-                BatchKey.gsp_t0_idx: xr_data.attrs["t0_idx"],
-                BatchKey.gsp_id: xr_data.gsp_id.values,
-                BatchKey.gsp_nominal_capacity_mwp: xr_data.isel(time_utc=0)[
-                    "nominal_capacity_mwp"
-                ].values,
-                BatchKey.gsp_effective_capacity_mwp: (
-                    xr_data.isel(time_utc=0)["effective_capacity_mwp"].values
-                ),
-                BatchKey.gsp_time_utc: datetime64_to_float(xr_data["time_utc"].values),
-            }
-
-            # Coordinates
-            for batch_key, dataset_key in (
-                (BatchKey.gsp_y_osgb, "y_osgb"),
-                (BatchKey.gsp_x_osgb, "x_osgb"),
-            ):
-                if dataset_key in xr_data.coords.keys():
-                    values = xr_data[dataset_key].values
-                    # Expand dims so AddFourierSpaceTime works!
-                    example[batch_key] = values  # np.expand_dims(values, axis=1)
-
-            yield example
+            yield convert_gsp_to_numpy_batch(xr_data)

ocf_datapipes/convert/numpy_batch/nwp.py

@@ -6,6 +6,31 @@
 from ocf_datapipes.utils.utils import datetime64_to_float
 
 
+def convert_nwp_to_numpy_batch(xr_data):
+    """Convert from Xarray to NWPBatchKey"""
+
+    example: NWPNumpyBatch = {
+        NWPBatchKey.nwp: xr_data.values,
+        NWPBatchKey.nwp_t0_idx: xr_data.attrs["t0_idx"],
+        NWPBatchKey.nwp_channel_names: xr_data.channel.values,
+        NWPBatchKey.nwp_init_time_utc: datetime64_to_float(xr_data.init_time_utc.values),
+        NWPBatchKey.nwp_step: (xr_data.step.values / np.timedelta64(1, "h")).astype(np.int64),
+    }
+
+    if "target_time_utc" in xr_data.coords:
+        target_time = xr_data.target_time_utc.values
+        example[NWPBatchKey.nwp_target_time_utc] = datetime64_to_float(target_time)
+
+    for batch_key, dataset_key in (
+        (NWPBatchKey.nwp_y_osgb, "y_osgb"),
+        (NWPBatchKey.nwp_x_osgb, "x_osgb"),
+    ):
+        if dataset_key in xr_data.coords:
+            example[batch_key] = xr_data[dataset_key].values
+
+    return example
+
+
 @functional_datapipe("convert_nwp_to_numpy_batch")
 class ConvertNWPToNumpyBatchIterDataPipe(IterDataPipe):
     """Convert NWP Xarray objects to NWPNumpyBatch"""
@@ -23,26 +48,4 @@ def __init__(self, source_datapipe: IterDataPipe):
     def __iter__(self) -> NWPNumpyBatch:
         """Convert from Xarray to NWPBatchKey"""
         for xr_data in self.source_datapipe:
-            example: NWPNumpyBatch = {
-                NWPBatchKey.nwp: xr_data.values,
-                NWPBatchKey.nwp_t0_idx: xr_data.attrs["t0_idx"],
-            }
-            if "target_time_utc" in xr_data.coords:
-                target_time = xr_data.target_time_utc.values
-                example[NWPBatchKey.nwp_target_time_utc] = datetime64_to_float(target_time)
-            example[NWPBatchKey.nwp_channel_names] = xr_data.channel.values
-            example[NWPBatchKey.nwp_step] = (xr_data.step.values / np.timedelta64(1, "h")).astype(
-                np.int64
-            )
-            example[NWPBatchKey.nwp_init_time_utc] = datetime64_to_float(
-                xr_data.init_time_utc.values
-            )
-
-            for batch_key, dataset_key in (
-                (NWPBatchKey.nwp_y_osgb, "y_osgb"),
-                (NWPBatchKey.nwp_x_osgb, "x_osgb"),
-            ):
-                if dataset_key in xr_data.coords.keys():
-                    example[batch_key] = xr_data[dataset_key].values
-
-            yield example
+            yield convert_nwp_to_numpy_batch(xr_data)

ocf_datapipes/convert/numpy_batch/pv.py

@@ -10,6 +10,23 @@
 logger = logging.getLogger(__name__)
 
 
+def convert_pv_to_numpy_batch(xr_data):
+    """Convert PV Xarray to NumpyBatch"""
+    example: NumpyBatch = {
+        BatchKey.pv: xr_data.values,
+        BatchKey.pv_t0_idx: xr_data.attrs["t0_idx"],
+        BatchKey.pv_ml_id: xr_data["ml_id"].values,
+        BatchKey.pv_id: xr_data["pv_system_id"].values.astype(np.float32),
+        BatchKey.pv_observed_capacity_wp: (xr_data["observed_capacity_wp"].values),
+        BatchKey.pv_nominal_capacity_wp: (xr_data["nominal_capacity_wp"].values),
+        BatchKey.pv_time_utc: datetime64_to_float(xr_data["time_utc"].values),
+        BatchKey.pv_latitude: xr_data["latitude"].values,
+        BatchKey.pv_longitude: xr_data["longitude"].values,
+    }
+
+    return example
+
+
 @functional_datapipe("convert_pv_to_numpy_batch")
 class ConvertPVToNumpyBatchIterDataPipe(IterDataPipe):
     """Convert PV Xarray to NumpyBatch"""
@@ -25,20 +42,6 @@ def __init__(self, source_datapipe: IterDataPipe):
         self.source_datapipe = source_datapipe
 
     def __iter__(self) -> NumpyBatch:
-        """Iterate and convert PV Xarray to NumpyBatch"""
+        """Convert PV Xarray to NumpyBatch"""
         for xr_data in self.source_datapipe:
-            logger.debug("Converting PV xarray to numpy example")
-
-            example: NumpyBatch = {
-                BatchKey.pv: xr_data.values,
-                BatchKey.pv_t0_idx: xr_data.attrs["t0_idx"],
-                BatchKey.pv_ml_id: xr_data["ml_id"].values,
-                BatchKey.pv_id: xr_data["pv_system_id"].values.astype(np.float32),
-                BatchKey.pv_observed_capacity_wp: (xr_data["observed_capacity_wp"].values),
-                BatchKey.pv_nominal_capacity_wp: (xr_data["nominal_capacity_wp"].values),
-                BatchKey.pv_time_utc: datetime64_to_float(xr_data["time_utc"].values),
-                BatchKey.pv_latitude: xr_data["latitude"].values,
-                BatchKey.pv_longitude: xr_data["longitude"].values,
-            }
-
-            yield example
+            yield convert_pv_to_numpy_batch(xr_data)

ocf_datapipes/convert/numpy_batch/satellite.py

@@ -5,6 +5,49 @@
 from ocf_datapipes.utils.utils import datetime64_to_float
 
 
+def _convert_satellite_to_numpy_batch(xr_data):
+    example: NumpyBatch = {
+        BatchKey.satellite_actual: xr_data.values,
+        BatchKey.satellite_t0_idx: xr_data.attrs["t0_idx"],
+        BatchKey.satellite_time_utc: datetime64_to_float(xr_data["time_utc"].values),
+    }
+
+    for batch_key, dataset_key in (
+        (BatchKey.satellite_y_geostationary, "y_geostationary"),
+        (BatchKey.satellite_x_geostationary, "x_geostationary"),
+    ):
+        # HRVSatellite coords are already float32.
+        example[batch_key] = xr_data[dataset_key].values
+
+    return example
+
+
+def _convert_hrvsatellite_to_numpy_batch(xr_data):
+    example: NumpyBatch = {
+        BatchKey.hrvsatellite_actual: xr_data.values,
+        BatchKey.hrvsatellite_t0_idx: xr_data.attrs["t0_idx"],
+        BatchKey.hrvsatellite_time_utc: datetime64_to_float(xr_data["time_utc"].values),
+    }
+
+    for batch_key, dataset_key in (
+        (BatchKey.hrvsatellite_y_geostationary, "y_geostationary"),
+        (BatchKey.hrvsatellite_x_geostationary, "x_geostationary"),
+    ):
+        # Satellite coords are already float32.
+        example[batch_key] = xr_data[dataset_key].values
+
+    return example
+
+
+def convert_satellite_to_numpy_batch(xr_data, is_hrv=False):
+    """Converts Xarray Satellite to NumpyBatch object"""
+    if is_hrv:
+        example = _convert_hrvsatellite_to_numpy_batch(xr_data)
+    else:
+        example = _convert_satellite_to_numpy_batch(xr_data)
+    return example
+
+
 @functional_datapipe("convert_satellite_to_numpy_batch")
 class ConvertSatelliteToNumpyBatchIterDataPipe(IterDataPipe):
     """Converts Xarray Satellite to NumpyBatch object"""
@@ -24,31 +67,4 @@ def __init__(self, source_datapipe: IterDataPipe, is_hrv: bool = False):
     def __iter__(self) -> NumpyBatch:
         """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,
-                    BatchKey.hrvsatellite_t0_idx: xr_data.attrs["t0_idx"],
-                    BatchKey.hrvsatellite_time_utc: datetime64_to_float(xr_data["time_utc"].values),
-                }
-
-                for batch_key, dataset_key in (
-                    (BatchKey.hrvsatellite_y_geostationary, "y_geostationary"),
-                    (BatchKey.hrvsatellite_x_geostationary, "x_geostationary"),
-                ):
-                    # HRVSatellite coords are already float32.
-                    example[batch_key] = xr_data[dataset_key].values
-            else:
-                example: NumpyBatch = {
-                    BatchKey.satellite_actual: xr_data.values,
-                    BatchKey.satellite_t0_idx: xr_data.attrs["t0_idx"],
-                    BatchKey.satellite_time_utc: datetime64_to_float(xr_data["time_utc"].values),
-                }
-
-                for batch_key, dataset_key in (
-                    (BatchKey.satellite_y_geostationary, "y_geostationary"),
-                    (BatchKey.satellite_x_geostationary, "x_geostationary"),
-                ):
-                    # HRVSatellite coords are already float32.
-                    example[batch_key] = xr_data[dataset_key].values
-
-            yield example
+            yield convert_satellite_to_numpy_batch(xr_data, self.is_hrv)

ocf_datapipes/convert/numpy_batch/sensor.py

@@ -10,34 +10,37 @@
 logger = logging.getLogger(__name__)
 
 
+def convert_sensor_to_numpy_batch(xr_data):
+    """Convert Sensor Xarray to NumpyBatch"""
+
+    example: NumpyBatch = {
+        BatchKey.sensor: xr_data.values,
+        BatchKey.sensor_t0_idx: xr_data.attrs["t0_idx"],
+        BatchKey.sensor_id: xr_data["station_id"].values.astype(np.float32),
+        # BatchKey.sensor_observed_capacity_wp: (xr_data["observed_capacity_wp"].values),
+        # BatchKey.sensor_nominal_capacity_wp: (xr_data["nominal_capacity_wp"].values),
+        BatchKey.sensor_time_utc: datetime64_to_float(xr_data["time_utc"].values),
+        BatchKey.sensor_latitude: xr_data["latitude"].values,
+        BatchKey.sensor_longitude: xr_data["longitude"].values,
+    }
+    return example
+
+
 @functional_datapipe("convert_sensor_to_numpy_batch")
 class ConvertSensorToNumpyBatchIterDataPipe(IterDataPipe):
     """Convert Sensor Xarray to NumpyBatch"""
 
     def __init__(self, source_datapipe: IterDataPipe):
         """
-        Convert PV Xarray objects to NumpyBatch objects
+        Convert sensor Xarray objects to NumpyBatch objects
 
         Args:
-            source_datapipe: Datapipe emitting PV Xarray objects
+            source_datapipe: Datapipe emitting sensor Xarray objects
         """
         super().__init__()
         self.source_datapipe = source_datapipe
 
     def __iter__(self) -> NumpyBatch:
-        """Iterate and convert PV Xarray to NumpyBatch"""
+        """Iterate and convert sensor Xarray to NumpyBatch"""
         for xr_data in self.source_datapipe:
-            logger.debug("Converting Sensor xarray to numpy example")
-
-            example: NumpyBatch = {
-                BatchKey.sensor: xr_data.values,
-                BatchKey.sensor_t0_idx: xr_data.attrs["t0_idx"],
-                BatchKey.sensor_id: xr_data["station_id"].values.astype(np.float32),
-                # BatchKey.sensor_observed_capacity_wp: (xr_data["observed_capacity_wp"].values),
-                # BatchKey.sensor_nominal_capacity_wp: (xr_data["nominal_capacity_wp"].values),
-                BatchKey.sensor_time_utc: datetime64_to_float(xr_data["time_utc"].values),
-                BatchKey.sensor_latitude: xr_data["latitude"].values,
-                BatchKey.sensor_longitude: xr_data["longitude"].values,
-            }
-
-            yield example
+            yield convert_sensor_to_numpy_batch(xr_data)

ocf_datapipes/convert/numpy_batch/wind.py

@@ -10,6 +10,24 @@
 logger = logging.getLogger(__name__)
 
 
+def convert_wind_to_numpy_batch(xr_data):
+    """Convert Wind Xarray to NumpyBatch"""
+
+    example: NumpyBatch = {
+        BatchKey.wind: xr_data.values,
+        BatchKey.wind_t0_idx: xr_data.attrs["t0_idx"],
+        BatchKey.wind_ml_id: xr_data["ml_id"].values,
+        BatchKey.wind_id: xr_data["wind_system_id"].values.astype(np.float32),
+        BatchKey.wind_observed_capacity_mwp: (xr_data["observed_capacity_mwp"].values),
+        BatchKey.wind_nominal_capacity_mwp: (xr_data["nominal_capacity_mwp"].values),
+        BatchKey.wind_time_utc: datetime64_to_float(xr_data["time_utc"].values),
+        BatchKey.wind_latitude: xr_data["latitude"].values,
+        BatchKey.wind_longitude: xr_data["longitude"].values,
+    }
+
+    return example
+
+
 @functional_datapipe("convert_wind_to_numpy_batch")
 class ConvertWindToNumpyBatchIterDataPipe(IterDataPipe):
     """Convert Wind Xarray to NumpyBatch"""
@@ -27,18 +45,4 @@ def __init__(self, source_datapipe: IterDataPipe):
     def __iter__(self) -> NumpyBatch:
         """Iterate and convert PV Xarray to NumpyBatch"""
         for xr_data in self.source_datapipe:
-            logger.debug("Converting Wind xarray to numpy example")
-
-            example: NumpyBatch = {
-                BatchKey.wind: xr_data.values,
-                BatchKey.wind_t0_idx: xr_data.attrs["t0_idx"],
-                BatchKey.wind_ml_id: xr_data["ml_id"].values,
-                BatchKey.wind_id: xr_data["wind_system_id"].values.astype(np.float32),
-                BatchKey.wind_observed_capacity_mwp: (xr_data["observed_capacity_mwp"].values),
-                BatchKey.wind_nominal_capacity_mwp: (xr_data["nominal_capacity_mwp"].values),
-                BatchKey.wind_time_utc: datetime64_to_float(xr_data["time_utc"].values),
-                BatchKey.wind_latitude: xr_data["latitude"].values,
-                BatchKey.wind_longitude: xr_data["longitude"].values,
-            }
-
-            yield example
+            yield convert_wind_to_numpy_batch(xr_data)