Update PVRollingPowerWindow and docstrings

Author: jacobbieker

ocf_datapipes/transform/xarray/add_t0idx_and_sample_period_duration.py

@@ -3,17 +3,27 @@
 import xarray as xr
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterDataPipe
-
+from datetime import timedelta
 
 @functional_datapipe("add_t0_idx_and_sample_period_duration")
 class AddT0IdxAndSamplePeriodDurationIterDataPipe(IterDataPipe):
-    def __init__(self, source_datapipe: IterDataPipe, sample_period_duration, history_duration):
+    """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):
+        """
+        Adds two attributes, t0_idx, and sample_period_duration for downstream datapipes to use
+
+        Args:
+            source_datapipe: Datapipe emitting a Xarray DataSet or DataArray
+            sample_period_duration: Time between samples
+            history_duration: Amount of history in each example
+        """
         self.source_datapipe = source_datapipe
         self.sample_period_duration = sample_period_duration
         self.history_duration = history_duration
         self.t0_idx = int(self.history_duration / self.sample_period_duration)
 
     def __iter__(self) -> Union[xr.DataArray, xr.Dataset]:
+        """Adds the two attributes to the xarray objects and returns them"""
         for xr_data in self.source_datapipe:
             xr_data.attrs["t0_idx"] = self.t0_idx
             xr_data.attrs["sample_period_duration"] = self.sample_period_duration

ocf_datapipes/transform/xarray/pv_power_rolling_window.py

@@ -1,6 +1,5 @@
 from typing import Optional, Union
 
-import numpy as np
 import pandas as pd
 import xarray as xr
 from torchdata.datapipes import functional_datapipe
@@ -9,24 +8,60 @@
 
 @functional_datapipe("pv_power_rolling_window")
 class PVPowerRollingWindowIterDataPipe(IterDataPipe):
+    """Compute rolling mean of PV power."""
     def __init__(
         self,
-        source_dp: IterDataPipe,
+        source_datapipe: IterDataPipe,
         window: Union[int, pd.tseries.offsets.DateOffset, pd.core.indexers.objects.BaseIndexer] = 3,
         min_periods: Optional[int] = 2,
         center: bool = True,
         win_type: Optional[str] = None,
         expect_dataset: bool = True,
     ):
-        self.source_dp = source_dp
+        """
+        Compute the rolling mean of PV power data
+
+        Args:
+            source_datapipe: Datapipe emitting PV Xarray object
+
+            window: Size of the moving window.
+            If an integer, the fixed number of observations used for each window.
+
+            If an offset, the time period of each window. Each window will be a variable sized
+            based on the observations included in the time-period. This is only valid for
+            datetimelike indexes. To learn more about the offsets & frequency strings, please see:
+            https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases
+
+            If a BaseIndexer subclass, the window boundaries based on the defined
+            `get_window_bounds` method. Additional rolling keyword arguments,
+            namely `min_periods` and `center` will be passed to `get_window_bounds`.
+
+            min_periods: Minimum number of observations in window required to have a value;
+            otherwise, result is `np.nan`.
+
+            To avoid NaNs at the start and end of the timeseries, this should be <= ceil(window/2).
+
+            For a window that is specified by an offset, `min_periods` will default to 1.
+
+            For a window that is specified by an integer, `min_periods` will default to the size of
+            the window.
+
+            center: If False, set the window labels as the right edge of the window index.
+            If True, set the window labels as the center of the window index.
+
+            win_type: Window type
+            expect_dataset: Whether to expect a dataset or DataArray
+        """
+        self.source_datapipe = source_datapipe
         self.window = window
         self.min_periods = min_periods
         self.center = center
         self.win_type = win_type
         self.expect_dataset = expect_dataset
 
-    def __iter__(self):
-        for xr_data in self.source_dp:
+    def __iter__(self) -> Union[xr.DataArray, xr.Dataset]:
+        """Compute rolling mean of PV power"""
+        for xr_data in self.source_datapipe:
             if self.expect_dataset:
                 data_to_resample = xr_data["power_w"]
             else:
@@ -47,22 +82,3 @@ def __iter__(self):
                 resampled.attrs[attr_name] = xr_data.attrs[attr_name]
 
             yield resampled
-
-
-def set_new_sample_period_and_t0_idx_attrs(xr_data, new_sample_period) -> xr.DataArray:
-    orig_sample_period = xr_data.attrs["sample_period_duration"]
-    orig_t0_idx = xr_data.attrs["t0_idx"]
-    new_sample_period = pd.Timedelta(new_sample_period)
-    assert new_sample_period >= orig_sample_period
-    new_t0_idx = orig_t0_idx / (new_sample_period / orig_sample_period)
-    np.testing.assert_almost_equal(
-        int(new_t0_idx),
-        new_t0_idx,
-        err_msg=(
-            "The original t0_idx must be exactly divisible by"
-            " (new_sample_period / orig_sample_period)"
-        ),
-    )
-    xr_data.attrs["sample_period_duration"] = new_sample_period
-    xr_data.attrs["t0_idx"] = int(new_t0_idx)
-    return xr_data

tests/transform/xarray/test_pv_power_rolling_window.py

@@ -9,6 +9,7 @@ def test_pv_power_rolling_window_passiv(passiv_dp):
         history_duration=timedelta(minutes=60),
         sample_period_duration=timedelta(minutes=5),
     )
+    data_before = next(iter(passiv_dp))
     passiv_dp = PVPowerRollingWindow(passiv_dp, expect_dataset=False)
     data = next(iter(passiv_dp))
-    assert data is not None
+    assert len(data.time_utc) == len(data_before.time_utc)