Merge pull request #92 from openclimatefix/location_picker_x_y_names

Location picker x y names

Author: peterdudfield

ocf_datapipes/config/model.py

@@ -41,6 +41,9 @@
 
 logger = logging.getLogger(__name__)
 
+# add SV to list of providers
+providers.append("SV")
+
 
 class Base(BaseModel):
     """Pydantic Base model where no extras can be added"""

ocf_datapipes/select/location_picker.py

@@ -1,5 +1,6 @@
 """Pick locations from a dataset"""
 import logging
+from typing import Optional
 
 import numpy as np
 from torchdata.datapipes import functional_datapipe
@@ -14,18 +15,28 @@
 class LocationPickerIterDataPipe(IterDataPipe):
     """Picks locations from a dataset and returns them"""
 
-    def __init__(self, source_datapipe: IterDataPipe, return_all_locations: bool = False):
+    def __init__(
+        self,
+        source_datapipe: IterDataPipe,
+        return_all_locations: bool = False,
+        x_dim_name: Optional[str] = "x_osgb",
+        y_dim_name: Optional[str] = "y_osgb",
+    ):
         """
         Picks locations from a dataset and returns them
 
         Args:
             source_datapipe: Datapipe emitting Xarray Dataset
             return_all_locations: Whether to return all locations,
-            if True, also returns them in order
+                if True, also returns them in order
+            x_dim_name: x dimension name, defaulted to 'x_osgb'
+            y_dim_name: y dimension name, defaulted to 'y_osgb'
         """
         super().__init__()
         self.source_datapipe = source_datapipe
         self.return_all_locations = return_all_locations
+        self.x_dim_name = x_dim_name
+        self.y_dim_name = y_dim_name
 
     def __iter__(self) -> Location:
         """Returns locations from the inputs datapipe"""
@@ -35,10 +46,10 @@ def __iter__(self) -> Location:
 
             if self.return_all_locations:
                 # Iterate through all locations in dataset
-                for location_idx in range(len(xr_dataset["x_osgb"])):
+                for location_idx in range(len(xr_dataset[self.x_dim_name])):
                     location = Location(
-                        x=xr_dataset["x_osgb"][location_idx].values,
-                        y=xr_dataset["y_osgb"][location_idx].values,
+                        x=xr_dataset[self.x_dim_name][location_idx].values,
+                        y=xr_dataset[self.y_dim_name][location_idx].values,
                     )
                     if "pv_system_id" in xr_dataset.coords.keys():
 
@@ -48,12 +59,13 @@ def __iter__(self) -> Location:
             else:
                 # Assumes all datasets have osgb coordinates for selecting locations
                 # Pick 1 random location from the input dataset
-                location_idx = np.random.randint(0, len(xr_dataset["x_osgb"]))
+                location_idx = np.random.randint(0, len(xr_dataset[self.x_dim_name]))
                 location = Location(
-                    x=xr_dataset["x_osgb"][location_idx].values,
-                    y=xr_dataset["y_osgb"][location_idx].values,
+                    x=xr_dataset[self.x_dim_name][location_idx].values,
+                    y=xr_dataset[self.y_dim_name][location_idx].values,
                 )
                 if "pv_system_id" in xr_dataset.coords.keys():
                     location.id = int(xr_dataset["pv_system_id"][location_idx].values)
                     logger.debug(f"Have selected location.id {location.id}")
-                yield location
+                for i in range(0, 10):
+                    yield location

ocf_datapipes/training/simple_pv.py

@@ -15,6 +15,9 @@
 logger = logging.getLogger(__name__)
 xarray.set_options(keep_attrs=True)
 
+# default is set to 1000
+BUFFERSIZE = -1
+
 
 def simple_pv_datapipe(
     configuration_filename: Union[Path, str], tag: Optional[str] = "train"
@@ -39,7 +42,9 @@ def simple_pv_datapipe(
 
     logger.debug("Opening Datasets")
     pv_datapipe, pv_location_datapipe = (
-        OpenPVFromNetCDF(pv=configuration.input_data.pv).pv_fill_night_nans().fork(2)
+        OpenPVFromNetCDF(pv=configuration.input_data.pv)
+        .pv_fill_night_nans()
+        .fork(2, buffer_size=BUFFERSIZE)
     )
 
     logger.debug("Add t0 idx")
@@ -71,7 +76,7 @@ def simple_pv_datapipe(
         )
         .ensure_n_pv_systems_per_example(n_pv_systems_per_example=1)
         .remove_nans()
-        .fork(3)
+        .fork(3, buffer_size=BUFFERSIZE)
     )
 
     # get contiguous time periods