Add second stage resampling to the example

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

examples/resampling.py

@@ -4,9 +4,10 @@
 """Frequenz Python SDK resampling example."""
 
 import asyncio
+from datetime import datetime, timezone
 
 from frequenz.channels import Broadcast
-from frequenz.channels.util import MergeNamed
+from frequenz.channels.util import Merge
 
 from frequenz.sdk import microgrid
 from frequenz.sdk.actor import (
@@ -16,11 +17,29 @@
     DataSourcingActor,
 )
 from frequenz.sdk.microgrid.component import ComponentCategory, ComponentMetricId
+from frequenz.sdk.timeseries import Sample, Sink, Source
+from frequenz.sdk.timeseries.resampling import Resampler
 
 HOST = "microgrid.sandbox.api.frequenz.io"
 PORT = 61060
 
 
+async def _calculate_average(source: Source, sink: Sink) -> None:
+    avg: float = 0.0
+    count: int = 0
+    async for sample in source:
+        print(f"Received sample to average at {sample.timestamp}: {sample.value}")
+        count += 1
+        if sample.value is None:
+            continue
+        avg = avg * (count - 1) / count + sample.value / count
+        await sink(Sample(datetime.now(timezone.utc), avg))
+
+
+async def _print_sample(sample: Sample) -> None:
+    print(f"\nResampled average at {sample.timestamp}: {sample.value}\n")
+
+
 async def run() -> None:  # pylint: disable=too-many-locals
     """Run main functions that initializes and creates everything."""
     await microgrid.initialize(HOST, PORT)
@@ -46,7 +65,7 @@ async def run() -> None:  # pylint: disable=too-many-locals
         channel_registry=channel_registry,
         data_sourcing_request_sender=data_source_request_sender,
         resampling_request_receiver=resampling_request_receiver,
-        resampling_period_s=1.0,
+        resampling_period_s=1,
     )
 
     components = await microgrid.get().api_client.components()
@@ -56,7 +75,9 @@ async def run() -> None:  # pylint: disable=too-many-locals
         if comp.category == ComponentCategory.BATTERY
     ]
 
-    # Create subscription requests for each time series id
+    print(f"Found {len(battery_ids)} batteries: {battery_ids}")
+
+    # Create subscription requests for each battery's SoC
     subscription_requests = [
         ComponentMetricRequest(
             namespace="resampling",
@@ -73,17 +94,37 @@ async def run() -> None:  # pylint: disable=too-many-locals
     )
 
     # Merge sample receivers for each subscription into one receiver
-    merged_receiver = MergeNamed(
-        **{
-            req.get_channel_name(): channel_registry.new_receiver(
-                req.get_channel_name()
-            )
+    merged_receiver = Merge(
+        *[
+            channel_registry.new_receiver(req.get_channel_name())
             for req in subscription_requests
-        }
+        ]
     )
 
-    async for channel_name, msg in merged_receiver:
-        print(f"{channel_name}: {msg}")
+    # Create a channel to calculate an average for all the data
+    average_chan = Broadcast[Sample]("average")
+
+    second_stage_resampler = Resampler(resampling_period_s=3.0)
+    second_stage_resampler.add_timeseries(average_chan.new_receiver(), _print_sample)
+
+    average_sender = average_chan.new_sender()
+    # Needed until channels Senders raises exceptions on errors
+    async def sink_adapter(sample: Sample) -> None:
+        assert await average_sender.send(sample)
+
+    print("Starting...")
+
+    try:
+        # This will run until it is interrupted (with Ctrl-C for example)
+        await asyncio.gather(
+            _calculate_average(merged_receiver, sink_adapter),
+            second_stage_resampler.resample(),
+        )
+    finally:
+        await second_stage_resampler.stop()
 
 
-asyncio.run(run())
+try:
+    asyncio.run(run())
+except KeyboardInterrupt:
+    print("Bye!")