Add option to impute missing values in window method for moving window (#669)

The `fill_value` option imputes missing values `NaN` or `None` in the
ring buffer and the moving window with the corresponding value. If
`None` is passed, no imputation is done. By default, missing values are
imputed with `NaN`. This prevents returning outdated return values
accidentally.

If missing values should be filled, the force_copy argument has to be
true to avoid overwriting the underlying data.

Author: cwasicki

RELEASE_NOTES.md

@@ -28,6 +28,7 @@
 - In `OrderedRingBuffer` and `MovingWindow`:
   - Support for integer indices is added.
   - Add `count_covered` method to count the number of elements covered by the used time range.
+  - Add `fill_value` option to window method to impute missing values. By default missing values are imputed with `NaN`.
 - Add `at` method to `MovingWindow` to access a single element and use it in `__getitem__` magic to fully support single element access.
 
 

src/frequenz/sdk/timeseries/_moving_window.py

@@ -289,6 +289,7 @@ def window(
         end: datetime | int | None,
         *,
         force_copy: bool = True,
+        fill_value: float | None = np.nan,
     ) -> ArrayLike:
         """
         Return an array containing the samples in the given time interval.
@@ -305,11 +306,16 @@ def window(
             force_copy: If `True`, the returned array is a copy of the underlying
                 data. Otherwise, if possible, a view of the underlying data is
                 returned.
+            fill_value: If not None, will use this value to fill missing values.
+                If missing values should be set, force_copy must be True.
+                Defaults to NaN to avoid returning outdated data unexpectedly.
 
         Returns:
             An array containing the samples in the given time interval.
         """
-        return self._buffer.window(start, end, force_copy=force_copy)
+        return self._buffer.window(
+            start, end, force_copy=force_copy, fill_value=fill_value
+        )
 
     async def _run_impl(self) -> None:
         """Awaits samples from the receiver and updates the underlying ring buffer.

src/frequenz/sdk/timeseries/_ringbuffer/buffer.py

@@ -304,6 +304,7 @@ def window(
         end: datetime | int | None,
         *,
         force_copy: bool = True,
+        fill_value: float | None = np.nan,
     ) -> FloatArray:
         """Request a copy or view on the data between start timestamp and end timestamp.
 
@@ -326,13 +327,21 @@ def window(
             end: end time of the window.
             force_copy: optional, default True. If True, will always create a
                 copy of the data.
+            fill_value: If not None, will use this value to fill missing values.
+                If missing values should be filled, force_copy must be True.
+                Defaults to NaN to avoid returning outdated data unexpectedly.
 
         Raises:
             IndexError: When start and end are not both datetime or index.
+            ValueError: When fill_value is not None and force_copy is False.
 
         Returns:
             The requested window
         """
+        # We don't want to modify the original buffer
+        if fill_value is not None and not force_copy:
+            raise ValueError("fill_value only supported for force_copy=True")
+
         if self.count_covered() == 0:
             return np.array([]) if isinstance(self._buffer, np.ndarray) else []
 
@@ -359,7 +368,48 @@ def window(
         start_pos = self.to_internal_index(start)
         end_pos = self.to_internal_index(end)
 
-        return self._wrapped_buffer_window(self._buffer, start_pos, end_pos, force_copy)
+        window = self._wrapped_buffer_window(
+            self._buffer, start_pos, end_pos, force_copy
+        )
+
+        if fill_value is not None:
+            window = self._fill_gaps(window, fill_value, start, self.gaps)
+        return window
+
+    def _fill_gaps(
+        self,
+        data: FloatArray,
+        fill_value: float,
+        oldest_timestamp: datetime,
+        gaps: list[Gap],
+    ) -> FloatArray:
+        """Fill the gaps in the data with the given fill_value.
+
+        Args:
+            data: The data to fill.
+            fill_value: The value to fill the gaps with.
+            oldest_timestamp: The oldest timestamp in the data.
+            gaps: List of gaps to fill.
+
+        Returns:
+            The filled data.
+        """
+        assert isinstance(
+            data, (np.ndarray, list)
+        ), f"Unsupported data type {type(data)}"
+        for gap in gaps:
+            start_index = (gap.start - oldest_timestamp) // self._sampling_period
+            end_index = (gap.end - oldest_timestamp) // self._sampling_period
+            start_index = max(start_index, 0)
+            end_index = min(end_index, len(data))
+            if start_index < end_index:
+                if isinstance(data, np.ndarray):
+                    data[start_index:end_index] = fill_value
+                elif isinstance(data, list):
+                    data[start_index:end_index] = [fill_value] * (
+                        end_index - start_index
+                    )
+        return data
 
     @staticmethod
     def _wrapped_buffer_window(

tests/timeseries/test_moving_window.py

@@ -163,10 +163,16 @@ def test_eq(expected: list[float], start: int | None, end: int | None) -> None:
             sender, [3.0], start_ts=UNIX_EPOCH + timedelta(seconds=3)
         )
         test_eq([0.0, 1.0], 0, 2)
-        # gap fill not supported yet:
-        # test_eq([0.0, 1.0, np.nan, 3.0], 0, None)
-        # test_eq([0.0, 1.0, np.nan, 3.0], -9, None)
-        # test_eq([np.nan, 3.0], -2, None)
+        # test gaps to be NaN
+        test_eq([0.0, 1.0, np.nan, 3.0], 0, None)
+        test_eq([np.nan, 3.0], -2, None)
+
+        # Test fill_value
+        assert np.allclose(
+            np.array([0.0, 1.0, 2.0, 3.0]),
+            window.window(0, None, fill_value=2.0),
+            equal_nan=True,
+        )
 
         # Complete window
         await push_logical_meter_data(sender, [0.0, 1.0, 2.0, 3.0, 4.0])

tests/timeseries/test_ringbuffer.py

@@ -526,12 +526,12 @@ def get_orb(data: FloatArray) -> OrderedRingBuffer[FloatArray]:
 def test_window_datetime() -> None:
     """Test the window function with datetime."""
     buffer = get_orb(np.array([0, None, 2, 3, 4]))
-    win = buffer.window(dt(0), dt(3), force_copy=False)
+    win = buffer.window(dt(0), dt(3), force_copy=False, fill_value=None)
     assert [0, np.nan, 2] == list(win)
     buffer._buffer[1] = 1  # pylint: disable=protected-access
     # Test whether the window is a view or a copy
     assert [0, 1, 2] == list(win)
-    win = buffer.window(dt(0), dt(3), force_copy=False)
+    win = buffer.window(dt(0), dt(3), force_copy=False, fill_value=None)
     assert [0, 1, 2] == list(win)
     # Empty array
     assert 0 == buffer.window(dt(1), dt(1)).size
@@ -561,8 +561,45 @@ def test_window_index() -> None:
     assert [] == buffer.window(-3, 0)
 
 
+def test_window_index_fill_value() -> None:
+    """Test fill_value functionality of window function."""
+    # Init with dummy data of size 3 and fill with [0, nan, 2]
+    buffer = OrderedRingBuffer([4711.0] * 3, ONE_SECOND)
+    buffer.update(Sample(dt(0), Quantity(0)))
+    buffer.update(Sample(dt(1), None))
+    buffer.update(Sample(dt(2), Quantity(2)))
+
+    # Test fill_value for explicitly set gaps
+    assert [0.0, np.nan, 2.0] == buffer.window(0, None)
+    assert [0.0, np.nan, 2.0] == buffer.window(0, None, fill_value=None)
+    assert [0.0, 1.0, 2.0] == buffer.window(0, None, fill_value=1)
+
+    # initial nan is ignored (optional implementation decision)
+    buffer.update(Sample(dt(3), Quantity(3)))  # -> [nan, 2, 3]
+    assert [2.0, 3.0] == buffer.window(0, None)
+    assert [2.0, 3.0] == buffer.window(0, None, fill_value=1)
+
+    # Test fill_value for gaps implicitly created gaps by time jumps
+    buffer.update(Sample(dt(4), Quantity(4)))
+    buffer.update(Sample(dt(6), Quantity(6)))  # -> [4, ?, 6]
+    # Default is fill missing values with NaNs
+    assert [4.0, np.nan, 6.0] == buffer.window(0, None)
+    assert [4.0, np.nan, 6.0] == buffer.window(0, None, fill_value=np.nan)
+    assert [4.0, 5.0, 6.0] == buffer.window(0, None, fill_value=5)
+    # If missing values not filled, outdated values can be returned unexpectedly
+    assert [4.0, 2, 6.0] == buffer.window(0, None, fill_value=None)
+
+    # Some edge cases
+    buffer.update(Sample(dt(7), Quantity(7)))  # -> [?, 6, 7]
+    assert [6.0, 7.0] == buffer.window(0, None)
+    buffer.update(Sample(dt(8), Quantity(np.nan)))  # -> [6, 7, nan]
+    assert [6.0, 7.0, np.nan] == buffer.window(0, None)
+    buffer.update(Sample(dt(9), None))  # -> [7, nan, nan]
+    assert [7.0, np.nan, np.nan] == buffer.window(0, None)
+
+
 def test_window_fail() -> None:
-    """Test the window function with invalid indices."""
+    """Test the window function with invalid arguments."""
     buffer = get_orb([0.0, 1.0, 2.0, 3.0, 4.0])
     # Go crazy with the indices
     with pytest.raises(IndexError):
@@ -573,6 +610,9 @@ def test_window_fail() -> None:
         buffer.window(None, dt(2))
     with pytest.raises(IndexError):
         buffer.window(dt(2), None)
+    # Invalid argument combination
+    with pytest.raises(ValueError):
+        buffer.window(0, 1, force_copy=False, fill_value=0)
 
 
 def test_wrapped_buffer_window() -> None: