[dirty-commit] Add resampler to the MovingWindow

Signed-off-by: Daniel Zullo <[email protected]>

Author: daniel-zullo-frequenz

src/frequenz/sdk/timeseries/_moving_window.py

@@ -13,11 +13,12 @@
 from typing import SupportsIndex, overload
 
 import numpy as np
-from frequenz.channels import Receiver
+from frequenz.channels import Broadcast, Receiver, Sender
 from numpy.typing import ArrayLike
 
 from .._internal.asyncio import cancel_and_await
 from . import Sample
+from ._resampling import Resampler, ResamplerConfig
 from ._ringbuffer import OrderedRingBuffer
 
 log = logging.getLogger(__name__)
@@ -91,11 +92,12 @@ class MovingWindow:
     ```
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         size: timedelta,
         resampled_data_recv: Receiver[Sample],
         sampling_period: timedelta,
+        resampling_period: timedelta | None = None,
         window_alignment: datetime = datetime(1, 1, 1),
     ) -> None:
         """
@@ -110,6 +112,7 @@ def __init__(
             resampled_data_recv: A receiver that delivers samples with a
                 given sampling period.
             sampling_period: The sampling period.
+            resampling_period: The resampling period.
             window_alignment: A datetime object that defines a point in time to which
                 the window is aligned to modulo window size.
                 (default is midnight 01.01.01)
@@ -125,20 +128,39 @@ def __init__(
             sampling_period <= size
         ), "The sampling period should be equal to or lower than the window size."
 
+        sampling = sampling_period
+        self._resampler: Resampler | None = None
+        self._avg_sender: Sender[Sample] | None = None
+        self._resampler_task: asyncio.Task[None] | None = None
+        if resampling_period is not None:
+            assert (
+                resampling_period >= sampling_period
+            ), "Resampling must be equal to or higher than sampling, only down-sampling supported."
+
+            if resampling_period > sampling_period:
+                self._resampler = Resampler(
+                    ResamplerConfig(
+                        resampling_period_s=resampling_period.total_seconds()
+                    )
+                )
+                sampling = resampling_period
+
         # Sampling period might not fit perfectly into the window size.
-        num_samples = math.ceil(size / sampling_period)
+        num_samples = math.ceil(size.total_seconds() / sampling.total_seconds())
 
         self._resampled_data_recv = resampled_data_recv
         self._buffer = OrderedRingBuffer(
             np.empty(shape=num_samples, dtype=float),
-            sampling_period=sampling_period,
+            sampling_period=sampling,
             time_index_alignment=window_alignment,
         )
 
+        if self._resampler:
+            self._configure_resampler()
+
         self._update_window_task: asyncio.Task[None] = asyncio.create_task(
             self._run_impl()
         )
-        log.debug("Cancelling MovingWindow task: %s", __name__)
 
     async def _run_impl(self) -> None:
         """Awaits samples from the receiver and updates the underlying ringbuffer.
@@ -149,16 +171,34 @@ async def _run_impl(self) -> None:
         try:
             async for sample in self._resampled_data_recv:
                 log.debug("Received new sample: %s", sample)
-                self._buffer.update(sample)
+                if self._resampler:
+                    await self._avg_sender.send(sample)  # type: ignore
+                else:
+                    self._buffer.update(sample)
+
         except asyncio.CancelledError:
             log.info("MovingWindow task has been cancelled.")
             raise
 
         log.error("Channel has been closed")
 
     async def stop(self) -> None:
-        """Cancel the running task and stop the MovingWindow."""
+        """Cancel the running tasks and stop the MovingWindow."""
         await cancel_and_await(self._update_window_task)
+        if self._resampler_task:
+            await cancel_and_await(self._resampler_task)
+
+    def _configure_resampler(self) -> None:
+        """Configure the components needed to run the resampler."""
+        assert isinstance(self._resampler, Resampler)
+
+        async def sink_buffer(sample: Sample) -> None:
+            self._buffer.update(sample)
+
+        avg_channel = Broadcast[Sample]("average")
+        self._avg_sender = avg_channel.new_sender()
+        self._resampler.add_timeseries("avg", avg_channel.new_receiver(), sink_buffer)
+        self._resampler_task = asyncio.create_task(self._resampler.resample())
 
     def __len__(self) -> int:
         """