Add Power Perceiver Production DataPipes

#minor

Author: jacobbieker

.gitignore

@@ -6,6 +6,10 @@ __pycache__/
 # C extensions
 *.so
 
+.idea/*
+.idea
+
+
 # Distribution / packaging
 .Python
 build/

README.md

@@ -1,2 +1,27 @@
 # ocf_datapipes
 OCF's DataPipe based dataloader for training and inference
+
+
+## Adding a new DataPipe
+A general outline for a new DataPipe should go something
+like this:
+
+```python
+from torchdata.datapipes.iter import IterDataPipe
+from torchdata.datapipes import functional_datapipe
+
+@functional_datapipe("<pipelet_name>")
+class <PipeletName>IterDataPipe(IterDataPipe):
+    def __init__(self):
+        pass
+
+    def __iter__(self):
+        pass
+```
+
+### Experimental DataPipes
+
+For new datapipes being developed for new models or input modalities, to somewhat separate the more experimental and in
+development datapipes from the ones better tested for production purposes, there is an `ocf_datapipes.experimental` namespace for
+developing these more research-y datapipes. These datapipes might not, and probably are not, tested.
+Once the model(s) using them are in production, they should be upgraded to one of the other namespaces and have tests added.

environment.yml

@@ -15,3 +15,6 @@ dependencies:
 pip:
   - einops
   - pathy
+  - git+https://github.com/SheffieldSolar/PV_Live-API
+  - pyaml_env
+  - nowcasting_datamodel

ocf_datapipes/batch/__init__.py

@@ -0,0 +1,4 @@
+from .merge_numpy_examples_to_batch import (
+    MergeNumpyExamplesToBatchIterDataPipe as MergeNumpyExamplesToBatch,
+)
+from .merge_numpy_modalities import MergeNumpyModalitiesIterDataPipe as MergeNumpyModalities

ocf_datapipes/batch/merge_numpy_examples_to_batch.py

@@ -0,0 +1,20 @@
+from torchdata.datapipes import functional_datapipe
+from torchdata.datapipes.iter import IterDataPipe
+
+from ocf_datapipes.utils.consts import NumpyBatch
+from ocf_datapipes.utils.utils import stack_np_examples_into_batch
+
+
+@functional_datapipe("merge_numpy_examples_to_batch")
+class MergeNumpyExamplesToBatchIterDataPipe(IterDataPipe):
+    def __init__(self, source_datapipe: IterDataPipe, n_examples_per_batch: int):
+        self.source_datapipe = source_datapipe
+        self.n_examples_per_batch = n_examples_per_batch
+
+    def __iter__(self) -> NumpyBatch:
+        np_examples = []
+        for np_batch in self.source_datapipe:
+            np_examples.append(np_batch)
+            if len(np_examples) == self.n_examples_per_batch:
+                yield stack_np_examples_into_batch(np_examples)
+                np_examples = []

ocf_datapipes/batch/merge_numpy_modalities.py

@@ -0,0 +1,17 @@
+from torchdata.datapipes import functional_datapipe
+from torchdata.datapipes.iter import IterDataPipe, Zipper
+
+from ocf_datapipes.utils.consts import BatchKey, NumpyBatch
+
+
+@functional_datapipe("merge_numpy_modalities")
+class MergeNumpyModalitiesIterDataPipe(IterDataPipe):
+    def __init__(self, source_datapipes: [IterDataPipe]):
+        self.source_datapipes = source_datapipes
+
+    def __iter__(self) -> NumpyBatch:
+        for np_batches in Zipper(*self.source_datapipes):
+            example: NumpyBatch = {}
+            for np_batch in np_batches:
+                example.update(np_batch)
+            yield example

ocf_datapipes/config/model.py

@@ -22,15 +22,12 @@
 from nowcasting_datamodel.models.pv import providers, pv_output, solar_sheffield_passiv
 
 # nowcasting_dataset imports
-from nowcasting_dataset.consts import (
+from ocf_datapipes.utils.consts import (
     DEFAULT_N_GSP_PER_EXAMPLE,
     DEFAULT_N_PV_SYSTEMS_PER_EXAMPLE,
     NWP_VARIABLE_NAMES,
     SAT_VARIABLE_NAMES,
 )
-from nowcasting_dataset.dataset.split import split
-from pathy import Pathy
-from pydantic import BaseModel, Field, root_validator, validator
 
 IMAGE_SIZE_PIXELS = 64
 IMAGE_SIZE_PIXELS_FIELD = Field(
@@ -622,14 +619,6 @@ class Process(Base):
             " return data at 11:30, 12:00, 12:30, and 13:00."
         ),
     )
-    split_method: split.SplitMethod = Field(
-        split.SplitMethod.DAY_RANDOM_TEST_DATE,
-        description=(
-            "The method used to split the t0 datetimes into train, validation and test sets."
-            " If the split method produces no t0 datetimes for any split_name, then"
-            " n_<split_name>_batches must also be set to 0."
-        ),
-    )
 
     train_test_validation_split: List[int] = Field(
         [10, 1, 1],

ocf_datapipes/convert/__init__.py

@@ -0,0 +1,4 @@
+from .gsp import ConvertGSPToNumpyBatchIterDataPipe as ConvertGSPToNumpyBatch
+from .nwp import ConvertNWPToNumpyBatchIterDataPipe as ConvertNWPToNumpyBatch
+from .pv import ConvertPVToNumpyBatchIterDataPipe as ConvertPVToNumpyBatch
+from .satellite import ConvertSatelliteToNumpyBatchIterDataPipe as ConvertSatelliteToNumpyBatch

ocf_datapipes/convert/gsp.py

@@ -0,0 +1,33 @@
+from torchdata.datapipes import functional_datapipe
+from torchdata.datapipes.iter import IterDataPipe
+
+from ocf_datapipes.utils.consts import BatchKey, NumpyBatch
+from ocf_datapipes.utils.utils import datetime64_to_float
+
+
+@functional_datapipe("convert_gsp_to_numpy_batch")
+class ConvertGSPToNumpyBatchIterDataPipe(IterDataPipe):
+    def __init__(self, source_dp: IterDataPipe):
+        super().__init__()
+        self.source_dp = source_dp
+
+    def __iter__(self) -> NumpyBatch:
+        for xr_data in self.source_dp:
+            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_capacity_mwp: xr_data.isel(time_utc=0)["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"),
+            ):
+                values = xr_data[dataset_key].values
+                # Expand dims so EncodeSpaceTime works!
+                example[batch_key] = values  # np.expand_dims(values, axis=1)
+
+            yield example

ocf_datapipes/convert/nwp.py

@@ -0,0 +1,36 @@
+import numpy as np
+from torchdata.datapipes import functional_datapipe
+from torchdata.datapipes.iter import IterDataPipe
+
+from ocf_datapipes.utils.consts import BatchKey, NumpyBatch
+from ocf_datapipes.utils.utils import datetime64_to_float
+
+
+@functional_datapipe("convert_nwp_to_numpy_batch")
+class ConvertNWPToNumpyBatchIterDataPipe(IterDataPipe):
+    def __init__(self, source_dp: IterDataPipe):
+        super().__init__()
+        self.source_dp = source_dp
+
+    def __iter__(self) -> NumpyBatch:
+        for xr_data in self.source_dp:
+            example: NumpyBatch = {
+                BatchKey.nwp: xr_data.values,
+                BatchKey.nwp_t0_idx: xr_data.attrs["t0_idx"],
+            }
+
+            target_time = xr_data.target_time_utc.values
+            example[BatchKey.nwp_target_time_utc] = datetime64_to_float(target_time)
+            example[BatchKey.nwp_channel_names] = xr_data.channel.values
+            example[BatchKey.nwp_step] = (xr_data.step.values / np.timedelta64(1, "h")).astype(
+                np.int64
+            )
+            example[BatchKey.nwp_init_time_utc] = datetime64_to_float(xr_data.init_time_utc.values)
+
+            for batch_key, dataset_key in (
+                (BatchKey.nwp_y_osgb, "y_osgb"),
+                (BatchKey.nwp_x_osgb, "x_osgb"),
+            ):
+                example[batch_key] = xr_data[dataset_key].values
+
+            yield example

ocf_datapipes/convert/pv.py

@@ -0,0 +1,28 @@
+import numpy as np
+from torchdata.datapipes import functional_datapipe
+from torchdata.datapipes.iter import IterDataPipe
+
+from ocf_datapipes.utils.consts import BatchKey, NumpyBatch
+from ocf_datapipes.utils.utils import datetime64_to_float
+
+
+@functional_datapipe("convert_pv_to_numpy_batch")
+class ConvertPVToNumpyBatchIterDataPipe(IterDataPipe):
+    def __init__(self, source_dp: IterDataPipe):
+        super().__init__()
+        self.source_dp = source_dp
+
+    def __iter__(self) -> NumpyBatch:
+        for xr_data in self.source_dp:
+            example: NumpyBatch = {
+                BatchKey.pv: xr_data.values,
+                BatchKey.pv_t0_idx: xr_data.attrs["t0_idx"],
+                BatchKey.pv_system_row_number: xr_data["pv_system_row_number"].values,
+                BatchKey.pv_id: xr_data["pv_system_id"].values.astype(np.float32),
+                BatchKey.pv_capacity_wp: xr_data["capacity_wp"].values,
+                BatchKey.pv_time_utc: datetime64_to_float(xr_data["time_utc"].values),
+                BatchKey.pv_x_osgb: xr_data["x_osgb"].values,
+                BatchKey.pv_y_osgb: xr_data["y_osgb"].values,
+            }
+
+            yield example

ocf_datapipes/convert/satellite.py

@@ -0,0 +1,48 @@
+from torchdata.datapipes import functional_datapipe
+from torchdata.datapipes.iter import IterDataPipe
+
+from ocf_datapipes.utils.consts import BatchKey, NumpyBatch
+from ocf_datapipes.utils.utils import datetime64_to_float
+
+
+@functional_datapipe("convert_satellite_to_numpy_batch")
+class ConvertSatelliteToNumpyBatchIterDataPipe(IterDataPipe):
+    def __init__(self, source_dp: IterDataPipe, is_hrv: bool = False):
+        super().__init__()
+        self.source_dp = source_dp
+        self.is_hrv = is_hrv
+
+    def __iter__(self) -> NumpyBatch:
+        for xr_data in self.source_dp:
+            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_osgb, "y_osgb"),
+                    (BatchKey.hrvsatellite_x_osgb, "x_osgb"),
+                    (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_osgb, "y_osgb"),
+                    (BatchKey.satellite_x_osgb, "x_osgb"),
+                    (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

ocf_datapipes/experimental/__init__.py

@@ -0,0 +1,2 @@
+"""Collection of experimental datapipes, explicitly not ready for production,
+ but can be upgraded later, these might not be tested, and are to make quick experiments easier"""

ocf_datapipes/load/__init__.py

@@ -1,8 +1,9 @@
+# from .configuration import OpenConfigurationIterDataPipe as OpenConfiguration
 from .gsp import OpenGSPIterDataPipe as OpenGSP
 from .nwp import OpenNWPIterDataPipe as OpenNWP
 from .pv import OpenPVFromDBIterDataPipe as OpenPVFromDB
 from .pv import OpenPVFromNetCDFIterDataPipe as OpenPVFromNetCDF
-from .satellite import OpenSatelliteDataPipe as OpenSatellite
+from .satellite import OpenSatelliteIterDataPipe as OpenSatellite
 
 try:
     import rioxarray  # Rioxarray is sometimes a pain to install, so only load this if its installed

ocf_datapipes/load/configuration.py

@@ -4,6 +4,8 @@
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterDataPipe
 
+from ocf_datapipes.config.model import Configuration
+
 
 @functional_datapipe("open_config")
 class OpenConfigurationIterDataPipe(IterDataPipe):
@@ -14,7 +16,7 @@ def __iter__(self):
         with fsspec.open(self.configuration_filename, mode="r") as stream:
             configuration = parse_config(data=stream)
 
-        # TODO Load into Pydantic Configuration class
+        configuration = Configuration(**configuration)
 
         while True:
             yield configuration

ocf_datapipes/load/gsp.py

@@ -1,4 +1,6 @@
 import datetime
+from pathlib import Path
+from typing import Optional, Union
 
 import geopandas as gpd
 import numpy as np
@@ -7,20 +9,38 @@
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterDataPipe
 
+from ocf_datapipes.utils.eso import get_gsp_metadata_from_eso, get_gsp_shape_from_eso
+
+try:
+    from ocf_datapipes.utils.eso import get_gsp_metadata_from_eso, get_gsp_shape_from_eso
+
+    _has_pvlive = True
+except ImportError:
+    print("Unable to import PVLive utils, please provide filenames with OpenGSP")
+    _has_pvlive = False
+
 
 @functional_datapipe("open_gsp")
 class OpenGSPIterDataPipe(IterDataPipe):
     def __init__(
         self,
-        gsp_pv_power_zarr_path: str,
-        gsp_id_to_region_id_filename: str,
-        sheffield_solar_region_path: str,
+        gsp_pv_power_zarr_path: Union[str, Path],
+        gsp_id_to_region_id_filename: Optional[str] = None,
+        sheffield_solar_region_path: Optional[str] = None,
         threshold_mw: int = 0,
         sample_period_duration: datetime.timedelta = datetime.timedelta(minutes=30),
     ):
         self.gsp_pv_power_zarr_path = gsp_pv_power_zarr_path
-        self.gsp_id_to_region_id_filename = gsp_id_to_region_id_filename
-        self.sheffield_solar_region_path = sheffield_solar_region_path
+        if (
+            gsp_id_to_region_id_filename is None
+            or sheffield_solar_region_path is None
+            and _has_pvlive
+        ):
+            self.gsp_id_to_region_id_filename = get_gsp_metadata_from_eso()
+            self.sheffield_solar_region_path = get_gsp_shape_from_eso()
+        else:
+            self.gsp_id_to_region_id_filename = gsp_id_to_region_id_filename
+            self.sheffield_solar_region_path = sheffield_solar_region_path
         self.threshold_mw = threshold_mw
         self.sample_period_duration = sample_period_duration
 
@@ -33,6 +53,12 @@ def __iter__(self):
         # Load GSP generation xr.Dataset:
         gsp_pv_power_mw_ds = xr.open_dataset(self.gsp_pv_power_zarr_path, engine="zarr")
 
+        # Have to remove ID 0 (National one) for rest to work
+        # TODO Do filtering later, deal with national here for now
+        gsp_pv_power_mw_ds = gsp_pv_power_mw_ds.isel(
+            gsp_id=slice(1, len(gsp_pv_power_mw_ds.gsp_id))
+        )
+
         # Ensure the centroids have the same GSP ID index as the GSP PV power:
         gsp_id_to_shape = gsp_id_to_shape.loc[gsp_pv_power_mw_ds.gsp_id]
 
@@ -44,7 +70,6 @@ def __iter__(self):
             x_osgb=gsp_id_to_shape.geometry.centroid.x.astype(np.float32),
             y_osgb=gsp_id_to_shape.geometry.centroid.y.astype(np.float32),
             capacity_mwp=gsp_pv_power_mw_ds.installedcapacity_mwp.data.astype(np.float32),
-            t0_idx=self.t0_idx,
         )
 
         del gsp_id_to_shape, gsp_pv_power_mw_ds
@@ -86,7 +111,6 @@ def _put_gsp_data_into_an_xr_dataarray(
     x_osgb: np.ndarray,
     y_osgb: np.ndarray,
     capacity_mwp: np.ndarray,
-    t0_idx: int,
 ) -> xr.DataArray:
     # Convert to xr.DataArray:
     data_array = xr.DataArray(
@@ -99,5 +123,4 @@ def _put_gsp_data_into_an_xr_dataarray(
         y_osgb=("gsp_id", y_osgb),
         capacity_mwp=(("time_utc", "gsp_id"), capacity_mwp),
     )
-    data_array.attrs["t0_idx"] = t0_idx
     return data_array