Update SelectLiveT0Time docstring and typing and test

Author: jacobbieker

ocf_datapipes/load/configuration.py

@@ -16,11 +16,11 @@ def __init__(self, configuration_filename: str):
         self.configuration_filename = configuration_filename
 
     def __iter__(self):
-        logger.debug(f'Going to open {self.configuration_filename}')
+        logger.debug(f"Going to open {self.configuration_filename}")
         with fsspec.open(self.configuration_filename, mode="r") as stream:
             configuration = parse_config(data=stream)
 
-        logger.debug(f'Converting to Configuration ({configuration})')
+        logger.debug(f"Converting to Configuration ({configuration})")
         configuration = Configuration(**configuration)
 
         while True:

ocf_datapipes/load/pv.py

@@ -214,10 +214,10 @@ def _load_pv_metadata(filename: str) -> pd.DataFrame:
             latitude, longitude, system_id, x_osgb, y_osgb
     """
     _log.info(f"Loading PV metadata from {filename}")
-    if 'passiv' in str(filename):
-        index_col = 'ss_id'
+    if "passiv" in str(filename):
+        index_col = "ss_id"
     else:
-        index_col = 'system_id'
+        index_col = "system_id"
     pv_metadata = pd.read_csv(filename, index_col=index_col)
 
     if "Unnamed: 0" in pv_metadata.columns:

ocf_datapipes/load/satellite.py

@@ -45,12 +45,24 @@ def open_sat_data(
     # Note that `rename` renames *both* the coordinates and dimensions, and keeps
     # the connection between the dims and coordinates, so we don't have to manually
     # use `data_array.set_index()`.
-    dataset = dataset.rename({"time": "time_utc",})
+    dataset = dataset.rename(
+        {
+            "time": "time_utc",
+        }
+    )
     if "y" in dataset.coords.keys():
-        dataset = dataset.rename({"y": "y_geostationary",})
+        dataset = dataset.rename(
+            {
+                "y": "y_geostationary",
+            }
+        )
 
     if "x" in dataset.coords.keys():
-        dataset = dataset.rename({"x": "x_geostationary",})
+        dataset = dataset.rename(
+            {
+                "x": "x_geostationary",
+            }
+        )
 
     # Flip coordinates to top-left first
     if dataset.y_geostationary[0] < dataset.y_geostationary[-1]:

ocf_datapipes/select/select_live_t0_time.py

@@ -8,10 +8,18 @@ class SelectLiveT0TimeIterDataPipe(IterDataPipe):
     """Select the history for the live data"""
 
     def __init__(self, source_datapipe: IterDataPipe, dim_name: str = "time_utc"):
+        """
+        Select history for the Xarray object
+
+        Args:
+            source_datapipe: Datapipe emitting Xarray objects
+            dim_name: The time dimension name to use
+        """
         self.source_datapipe = source_datapipe
         self.dim_name = dim_name
 
-    def __iter__(self):
+    def __iter__(self) -> pd.Timestamp:
+        """Get the latest timestamp and return it"""
         for xr_data in self.source_datapipe:
             # Get most recent time in data
             # Select the history that goes back that far

ocf_datapipes/transform/xarray/add_t0idx_and_sample_period_duration.py

@@ -5,10 +5,17 @@
 from torchdata.datapipes.iter import IterDataPipe
 from datetime import timedelta
 
+
 @functional_datapipe("add_t0_idx_and_sample_period_duration")
 class AddT0IdxAndSamplePeriodDurationIterDataPipe(IterDataPipe):
     """Add t0_idx and sample_period_duration attributes to datasets for downstream tasks"""
-    def __init__(self, source_datapipe: IterDataPipe, sample_period_duration: timedelta, history_duration: timedelta):
+
+    def __init__(
+        self,
+        source_datapipe: IterDataPipe,
+        sample_period_duration: timedelta,
+        history_duration: timedelta,
+    ):
         """
         Adds two attributes, t0_idx, and sample_period_duration for downstream datapipes to use
 

ocf_datapipes/transform/xarray/convert_to_nwp_target_times.py

@@ -10,6 +10,7 @@
 @functional_datapipe("convert_to_nwp_target_time")
 class ConvertToNWPTargetTimeIterDataPipe(IterDataPipe):
     """Converts NWP Xarray to use the target time"""
+
     def __init__(
         self,
         source_datapipe: IterDataPipe,

ocf_datapipes/transform/xarray/downsample.py

@@ -5,7 +5,15 @@
 @functional_datapipe("downsample")
 class DownsampleIterDataPipe(IterDataPipe):
     """Downsample Xarray dataset with coarsen"""
-    def __init__(self, source_datapipe: IterDataPipe, y_coarsen: int, x_coarsen: int, x_dim_name: str = "x_osgb", y_dim_name: str = "y_osgb"):
+
+    def __init__(
+        self,
+        source_datapipe: IterDataPipe,
+        y_coarsen: int,
+        x_coarsen: int,
+        x_dim_name: str = "x_osgb",
+        y_dim_name: str = "y_osgb",
+    ):
         """
         Downsample xarray dataset/dataarrays with coarsen
 
@@ -25,8 +33,10 @@ def __init__(self, source_datapipe: IterDataPipe, y_coarsen: int, x_coarsen: int
     def __iter__(self):
         """Coarsen the data on the specified dimensions"""
         for xr_data in self.source_datapipe:
-            yield xr_data.coarsen({
-                self.y_dim_name: self.y_coarsen,
-                self.x_dim_name: self.x_coarsen,},
+            yield xr_data.coarsen(
+                {
+                    self.y_dim_name: self.y_coarsen,
+                    self.x_dim_name: self.x_coarsen,
+                },
                 boundary="trim",
             ).mean()

ocf_datapipes/transform/xarray/get_contiguous_time_periods.py

@@ -1,24 +1,39 @@
-import datetime
-
 import pandas as pd
-import xarray as xr
+import numpy as np
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterDataPipe
+from datetime import timedelta
 
 
 @functional_datapipe("add_contiguous_time_periods")
 class GetContiguousT0TimePeriodsIterDataPipe(IterDataPipe):
+    """Get contiguous time periods for training"""
+
     def __init__(
-        self, source_dp: IterDataPipe, history_duration, forecast_duration, sample_period_duration
+        self,
+        source_datapipe: IterDataPipe,
+        history_duration: timedelta,
+        forecast_duration: timedelta,
+        sample_period_duration: timedelta,
     ):
-        self.source_dp = source_dp
+        """
+        Get contiguous time periods for use in determing t0 times for training
+
+        Args:
+            source_datapipe: Datapipe emitting a Xarray dataset
+            history_duration: Amount of time for the history of an example
+            forecast_duration: Amount of time for the forecast of an example
+            sample_period_duration: The sampling period of the data source
+        """
+        self.source_datapipe = source_datapipe
         self.history_duration = history_duration
         self.forecast_duration = forecast_duration
         self.total_duration = history_duration + forecast_duration
         self.sample_period_duration = sample_period_duration
 
     def __iter__(self) -> pd.DataFrame:
-        for xr_data in self.source_dp:
+        """Calculate contiguous time periods and return a dataframe containing them"""
+        for xr_data in self.source_datapipe:
             contiguous_time_periods = get_contiguous_time_periods(
                 datetimes=pd.DatetimeIndex(xr_data["time_utc"]),
                 min_seq_length=int(self.total_duration / self.sample_period_duration) + 1,
@@ -33,7 +48,7 @@ def __iter__(self) -> pd.DataFrame:
 
 
 def get_contiguous_t0_time_periods(
-    contiguous_time_periods, history_duration, forecast_duration
+    contiguous_time_periods: pd.DataFrame, history_duration: timedelta, forecast_duration: timedelta
 ) -> pd.DataFrame:
     """Get all time periods which contain valid t0 datetimes.
 
@@ -52,7 +67,7 @@ def get_contiguous_t0_time_periods(
 def get_contiguous_time_periods(
     datetimes: pd.DatetimeIndex,
     min_seq_length: int,
-    max_gap_duration: datetime.timedelta,
+    max_gap_duration: timedelta,
 ) -> pd.DataFrame:
     """Return a pd.DataFrame where each row records the boundary of a contiguous time period.
 

ocf_datapipes/transform/xarray/normalize.py

@@ -8,12 +8,13 @@
 @functional_datapipe("normalize")
 class NormalizeIterDataPipe(IterDataPipe):
     """Normalize the data in various methods"""
+
     def __init__(
         self,
         source_datapipe: IterDataPipe,
-        mean: Optional[Union[xr.Dataset, xr.DataArray, np.ndarray]]=None,
-        std: Optional[Union[xr.Dataset, xr.DataArray, np.ndarray]]=None,
-        max_value: Optional[Union[int, float]]=None,
+        mean: Optional[Union[xr.Dataset, xr.DataArray, np.ndarray]] = None,
+        std: Optional[Union[xr.Dataset, xr.DataArray, np.ndarray]] = None,
+        max_value: Optional[Union[int, float]] = None,
         calculate_mean_std_from_example: bool = False,
         normalize_fn: Optional[Callable] = None,
     ):

ocf_datapipes/transform/xarray/pv_power_rolling_window.py

@@ -9,6 +9,7 @@
 @functional_datapipe("pv_power_rolling_window")
 class PVPowerRollingWindowIterDataPipe(IterDataPipe):
     """Compute rolling mean of PV power."""
+
     def __init__(
         self,
         source_datapipe: IterDataPipe,

tests/convert/test_gsp.py

@@ -4,7 +4,6 @@
 from datetime import timedelta
 
 
-
 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)

tests/convert/test_nwp.py

@@ -9,7 +9,7 @@ def test_convert_nwp_to_numpy_batch(nwp_dp):
     nwp_dp = AddT0IdxAndSamplePeriodDuration(
         nwp_dp, sample_period_duration=timedelta(minutes=60), history_duration=timedelta(minutes=60)
     )
-    t0_dp = SelectLiveT0Time(nwp_dp, dim_name='init_time_utc')
+    t0_dp = SelectLiveT0Time(nwp_dp, dim_name="init_time_utc")
 
     nwp_dp = ConvertToNWPTargetTime(
         nwp_dp,

tests/convert/test_pv.py

@@ -6,7 +6,9 @@
 
 def test_convert_passiv_to_numpy_batch(passiv_dp):
     passiv_dp = AddT0IdxAndSamplePeriodDuration(
-        passiv_dp, sample_period_duration=timedelta(minutes=5), history_duration=timedelta(minutes=60)
+        passiv_dp,
+        sample_period_duration=timedelta(minutes=5),
+        history_duration=timedelta(minutes=60),
     )
     passiv_dp = ConvertPVToNumpyBatch(passiv_dp)
     data = next(iter(passiv_dp))
@@ -15,7 +17,9 @@ def test_convert_passiv_to_numpy_batch(passiv_dp):
 
 def test_convert_pvoutput_to_numpy_batch(pvoutput_dp):
     pvoutput_dp = AddT0IdxAndSamplePeriodDuration(
-        pvoutput_dp, sample_period_duration=timedelta(minutes=5), history_duration=timedelta(minutes=60)
+        pvoutput_dp,
+        sample_period_duration=timedelta(minutes=5),
+        history_duration=timedelta(minutes=60),
     )
 
     pvoutput_dp = ConvertPVToNumpyBatch(pvoutput_dp)

tests/load/test_load_config.py

@@ -1,9 +1,7 @@
-
 from ocf_datapipes.load import OpenConfiguration
 
 
 def test_open_config():
     config_dp = OpenConfiguration("tests/config/test.yaml")
     configuration = next(iter(config_dp))
     print(configuration)
-

tests/production/test_pp_production.py

@@ -6,7 +6,7 @@
 import pytest
 
 
-@pytest.mark.skip('Need to set up laod PV from database first')
+@pytest.mark.skip("Need to set up laod PV from database first")
 def test_pp_production_datapipe():
 
     filename = os.path.join(os.path.dirname(ocf_datapipes.__file__), "../tests/config/test.yaml")

tests/select/test_select_live_t0_time_slice.py

@@ -58,7 +58,7 @@ def test_select_passiv(passiv_dp):
 
 def test_select_pvoutput(pvoutput_dp):
     time_len = len(next(iter(pvoutput_dp)).time_utc.values)
-    t0_dp = SelectLiveT0Time(pvoutput_dp, dim_name='time_utc')
+    t0_dp = SelectLiveT0Time(pvoutput_dp, dim_name="time_utc")
     pvoutput_dp = SelectLiveTimeSlice(
         pvoutput_dp,
         t0_datapipe=t0_dp,

tests/select/test_select_live_time_slice.py

@@ -1,7 +1,7 @@
 from datetime import timedelta
 
 from ocf_datapipes.select import SelectLiveT0Time, SelectLiveTimeSlice
-# from ocf_datapipes.transform.xarray import AddNWPTargetTime
+from ocf_datapipes.transform.xarray import ConvertToNWPTargetTime
 
 
 def test_select_hrv(sat_hrv_dp):
@@ -22,28 +22,25 @@ def test_select_gsp(gsp_dp):
     assert len(data.time_utc.values) < time_len
 
 
-# def test_select_nwp(nwp_dp):
-#     t0_dp = SelectLiveT0Time(nwp_dp, dim_name="init_time_utc")
-#     nwp_dp = AddNWPTargetTime(
-#         nwp_dp,
-#         t0_dp,
-#         sample_period_duration=timedelta(hours=1),
-#         history_duration=timedelta(hours=2),
-#         forecast_duration=timedelta(hours=4),
-#     )
-#     data = next(iter(nwp_dp))
-#     print(data)
-#     print(data.init_time_utc.values)
-#     time_len = len(next(iter(nwp_dp)).target_time_utc.values)
-#     nwp_dp = SelectLiveTimeSlice(
-#         nwp_dp,
-#         t0_datapipe=t0_dp,
-#         history_duration=timedelta(minutes=120),
-#         dim_name="target_time_utc",
-#     )
-#     data = next(iter(nwp_dp))
-#     assert len(data.time_utc.values) == 3
-#     assert len(data.time_utc.values) < time_len
+def test_select_nwp(nwp_dp):
+    t0_dp = SelectLiveT0Time(nwp_dp, dim_name="init_time_utc")
+    nwp_dp = ConvertToNWPTargetTime(
+        nwp_dp,
+        t0_dp,
+        sample_period_duration=timedelta(hours=1),
+        history_duration=timedelta(hours=2),
+        forecast_duration=timedelta(hours=4),
+    )
+    time_len = len(next(iter(nwp_dp)).target_time_utc.values)
+    nwp_dp = SelectLiveTimeSlice(
+         nwp_dp,
+         t0_datapipe=t0_dp,
+         history_duration=timedelta(minutes=120),
+         dim_name="target_time_utc",
+    )
+    data = next(iter(nwp_dp))
+    assert len(data.target_time_utc.values) == 3
+    assert len(data.target_time_utc.values) < time_len
 
 
 def test_select_passiv(passiv_dp):

tests/transform/xarray/test_add_nwp_target_time.py

@@ -2,10 +2,16 @@
 from ocf_datapipes.select import SelectLiveT0Time
 from datetime import timedelta
 
+
 def test_add_nwp_target_time(nwp_dp):
     t0_dp = SelectLiveT0Time(nwp_dp, dim_name="init_time_utc")
-    nwp_dp = ConvertToNWPTargetTime(nwp_dp, t0_dp, sample_period_duration=timedelta(minutes=60), history_duration=timedelta(hours=2), forecast_duration=timedelta(hours=3))
+    nwp_dp = ConvertToNWPTargetTime(
+        nwp_dp,
+        t0_dp,
+        sample_period_duration=timedelta(minutes=60),
+        history_duration=timedelta(hours=2),
+        forecast_duration=timedelta(hours=3),
+    )
     data = next(iter(nwp_dp))
     assert "target_time_utc" in data.coords
     assert len(data.coords["target_time_utc"]) == 6
-

tests/transform/xarray/test_downsample.py

@@ -10,7 +10,13 @@ def test_nwp_downsample(nwp_dp):
 
 
 def test_sat_downsample(sat_dp):
-    sat_dp = Downsample(sat_dp, y_coarsen=16, x_coarsen=16, y_dim_name="y_geostationary", x_dim_name="x_geostationary")
+    sat_dp = Downsample(
+        sat_dp,
+        y_coarsen=16,
+        x_coarsen=16,
+        y_dim_name="y_geostationary",
+        x_dim_name="x_geostationary",
+    )
     data = next(iter(sat_dp))
     assert data.shape[-1] == 38
     assert data.shape[-2] == 18