Use Sample[Energy] to stream BatteryPool.capacity values

Because the capacity bounds are no longer streamed, and only the
usable capacity is streamed, a custom type is no longer necessary.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/sdk/timeseries/battery_pool/_metric_calculator.py

@@ -12,9 +12,9 @@
 
 from ...microgrid import connection_manager
 from ...microgrid.component import ComponentCategory, ComponentMetricId, InverterType
-from ...timeseries import Quantity, Sample
+from ...timeseries import Energy, Quantity, Sample
 from ._component_metrics import ComponentMetricsData
-from ._result_types import Bound, CapacityMetrics, PowerMetrics
+from ._result_types import Bound, PowerMetrics
 
 _logger = logging.getLogger(__name__)
 _MIN_TIMESTAMP = datetime.min.replace(tzinfo=timezone.utc)
@@ -59,7 +59,7 @@ def battery_inverter_mapping(batteries: Iterable[int]) -> dict[int, int]:
 
 # Formula output types class have no common interface
 # Print all possible types here.
-T = TypeVar("T", Sample[Quantity], CapacityMetrics, PowerMetrics)
+T = TypeVar("T", Sample[Quantity], Sample[Energy], PowerMetrics)
 
 
 class MetricCalculator(ABC, Generic[T]):
@@ -142,7 +142,7 @@ def calculate(
         """
 
 
-class CapacityCalculator(MetricCalculator[CapacityMetrics]):
+class CapacityCalculator(MetricCalculator[Sample[Energy]]):
     """Define how to calculate Capacity metrics."""
 
     def __init__(self, batteries: Set[int]) -> None:
@@ -190,7 +190,7 @@ def calculate(
         self,
         metrics_data: dict[int, ComponentMetricsData],
         working_batteries: set[int],
-    ) -> CapacityMetrics | None:
+    ) -> Sample[Energy] | None:
         """Aggregate the metrics_data and calculate high level metric.
 
         Missing components will be ignored. Formula will be calculated for all
@@ -207,7 +207,8 @@ def calculate(
             High level metric calculated from the given metrics.
             Return None if there are no component metrics.
         """
-        result = CapacityMetrics(timestamp=_MIN_TIMESTAMP, total_capacity=0)
+        timestamp = _MIN_TIMESTAMP
+        total_capacity = 0.0
 
         for battery_id in working_batteries:
             if battery_id not in metrics_data:
@@ -223,10 +224,14 @@ def calculate(
             if capacity is None or soc_lower_bound is None or soc_upper_bound is None:
                 continue
             usable_capacity = capacity * (soc_upper_bound - soc_lower_bound) / 100
-            result.timestamp = max(result.timestamp, metrics.timestamp)
-            result.total_capacity += usable_capacity
+            timestamp = max(timestamp, metrics.timestamp)
+            total_capacity += usable_capacity
 
-        return None if result.timestamp == _MIN_TIMESTAMP else result
+        return (
+            None
+            if timestamp == _MIN_TIMESTAMP
+            else Sample[Energy](timestamp, Energy.from_watt_hours(total_capacity))
+        )
 
 
 class SoCCalculator(MetricCalculator[Sample[Quantity]]):

src/frequenz/sdk/timeseries/battery_pool/battery_pool.py

@@ -27,10 +27,10 @@
     FormulaGeneratorConfig,
     FormulaType,
 )
-from .._quantities import Power
+from .._quantities import Energy, Power
 from ._methods import MetricAggregator, SendOnUpdate
 from ._metric_calculator import CapacityCalculator, PowerBoundsCalculator, SoCCalculator
-from ._result_types import CapacityMetrics, PowerMetrics
+from ._result_types import PowerMetrics
 
 
 class BatteryPool:
@@ -357,7 +357,7 @@ def soc(self) -> MetricAggregator[Sample[Quantity]]:
         return self._active_methods[method_name]
 
     @property
-    def capacity(self) -> MetricAggregator[CapacityMetrics]:
+    def capacity(self) -> MetricAggregator[Sample[Energy]]:
         """Get receiver to receive new capacity metrics when they change.
 
         Capacity formulas are described in the receiver return type.  None will be send

tests/timeseries/_battery_pool/test_battery_pool.py

@@ -24,14 +24,8 @@
 from frequenz.sdk.actor import ResamplerConfig
 from frequenz.sdk.actor.power_distributing import BatteryStatus
 from frequenz.sdk.microgrid.component import ComponentCategory
-from frequenz.sdk.timeseries import Quantity, Sample
-from frequenz.sdk.timeseries._quantities import Power
-from frequenz.sdk.timeseries.battery_pool import (
-    BatteryPool,
-    Bound,
-    CapacityMetrics,
-    PowerMetrics,
-)
+from frequenz.sdk.timeseries import Energy, Power, Quantity, Sample
+from frequenz.sdk.timeseries.battery_pool import BatteryPool, Bound, PowerMetrics
 from frequenz.sdk.timeseries.battery_pool._metric_calculator import (
     battery_inverter_mapping,
 )
@@ -508,93 +502,115 @@ async def run_capacity_test(setup_args: SetupArgs) -> None:
         capacity_receiver.receive(), timeout=WAIT_FOR_COMPONENT_DATA_SEC + 0.2
     )
     now = datetime.now(tz=timezone.utc)
-    expected = CapacityMetrics(
+    expected = Sample[Energy](
         timestamp=now,
-        total_capacity=50.0,  # 50% of 50 kWh + 50% of 50 kWh = 25 + 25 = 50 kWh
+        value=Energy.from_watt_hours(
+            50.0
+        ),  # 50% of 50 kWh + 50% of 50 kWh = 25 + 25 = 50 kWh
     )
     compare_messages(msg, expected, WAIT_FOR_COMPONENT_DATA_SEC + 0.2)
 
     batteries_in_pool = list(battery_pool.battery_ids)
-    scenarios: list[Scenario[CapacityMetrics]] = [
+    scenarios: list[Scenario[Sample[Energy]]] = [
         Scenario(
             batteries_in_pool[0],
             {"capacity": 90.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                70.0,  # 50% of 90 kWh + 50% of 50 kWh = 45 + 25 = 70 kWh
+                Energy.from_watt_hours(
+                    70.0
+                ),  # 50% of 90 kWh + 50% of 50 kWh = 45 + 25 = 70 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[1],
             {"soc_lower_bound": 0.0, "soc_upper_bound": 90.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                90.0,  # 50% of 90 kWh + 90% of 50 kWh = 45 + 45 = 90 kWh
+                Energy.from_watt_hours(
+                    90.0
+                ),  # 50% of 90 kWh + 90% of 50 kWh = 45 + 45 = 90 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[0],
             {"capacity": 0.0, "soc_lower_bound": 0.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                45.0,  # 75% of 0 kWh + 90% of 50 kWh = 0 + 45 = 45 kWh
+                Energy.from_watt_hours(
+                    45.0
+                ),  # 75% of 0 kWh + 90% of 50 kWh = 0 + 45 = 45 kWh
             ),
         ),
         # Test zero division error
         Scenario(
             batteries_in_pool[1],
             {"capacity": 0.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                0.0,  # 75% of 0 kWh + 90% of 0 kWh = 0 + 0 = 0 kWh
+                Energy.from_watt_hours(
+                    0.0
+                ),  # 75% of 0 kWh + 90% of 0 kWh = 0 + 0 = 0 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[1],
             {"capacity": 50.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                45.0,  # 75% of 0 kWh + 90% of 50 kWh = 0 + 45 = 45 kWh
+                Energy.from_watt_hours(
+                    45.0
+                ),  # 75% of 0 kWh + 90% of 50 kWh = 0 + 45 = 45 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[1],
             {"soc_upper_bound": float("NaN")},
-            CapacityMetrics(
+            Sample(
                 now,
-                0.0,  # 75% of 0 kWh + 90% of 0 kWh = 0 + 0 = 0 kWh
+                Energy.from_watt_hours(
+                    0.0
+                ),  # 75% of 0 kWh + 90% of 0 kWh = 0 + 0 = 0 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[0],
             {"capacity": 30.0, "soc_lower_bound": 20.0, "soc_upper_bound": 90.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                21.0,  # 70% of 30 kWh + 90% of 0 kWh = 21 + 0 = 21 kWh
+                Energy.from_watt_hours(
+                    21.0
+                ),  # 70% of 30 kWh + 90% of 0 kWh = 21 + 0 = 21 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[1],
             {"capacity": 200.0, "soc_lower_bound": 20.0, "soc_upper_bound": 90.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                161.0,  # 70% of 30 kWh + 70% of 200 kWh = 21 + 140 = 161 kWh
+                Energy.from_watt_hours(
+                    161.0
+                ),  # 70% of 30 kWh + 70% of 200 kWh = 21 + 140 = 161 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[1],
             {"capacity": float("NaN")},
-            CapacityMetrics(
+            Sample(
                 now,
-                21.0,  # 70% of 30 kWh + 70% of 0 kWh = 21 + 0 = 21 kWh
+                Energy.from_watt_hours(
+                    21.0
+                ),  # 70% of 30 kWh + 70% of 0 kWh = 21 + 0 = 21 kWh
             ),
         ),
         Scenario(
             batteries_in_pool[1],
             {"capacity": 200.0},
-            CapacityMetrics(
+            Sample(
                 now,
-                161.0,  # 70% of 30 kWh + 70% of 200 kWh = 21 + 140 = 161 kWh
+                Energy.from_watt_hours(
+                    161.0
+                ),  # 70% of 30 kWh + 70% of 200 kWh = 21 + 140 = 161 kWh
             ),
         ),
     ]
@@ -608,15 +624,15 @@ async def run_capacity_test(setup_args: SetupArgs) -> None:
         all_batteries=all_batteries,
         batteries_in_pool=batteries_in_pool,
         waiting_time_sec=waiting_time_sec,
-        all_pool_result=CapacityMetrics(now, 161.0),
-        only_first_battery_result=CapacityMetrics(now, 21.0),
+        all_pool_result=Sample(now, Energy.from_watt_hours(161.0)),
+        only_first_battery_result=Sample(now, Energy.from_watt_hours(21.0)),
     )
 
     # One battery stopped sending data.
     await streamer.stop_streaming(batteries_in_pool[1])
     await asyncio.sleep(MAX_BATTERY_DATA_AGE_SEC + 0.2)
     msg = await asyncio.wait_for(capacity_receiver.receive(), timeout=waiting_time_sec)
-    compare_messages(msg, CapacityMetrics(now, 21.0), 0.2)
+    compare_messages(msg, Sample(now, Energy.from_watt_hours(21.0)), 0.2)
 
     # All batteries stopped sending data.
     await streamer.stop_streaming(batteries_in_pool[0])
@@ -628,7 +644,7 @@ async def run_capacity_test(setup_args: SetupArgs) -> None:
     latest_data = streamer.get_current_component_data(batteries_in_pool[0])
     streamer.start_streaming(latest_data, sampling_rate=0.1)
     msg = await asyncio.wait_for(capacity_receiver.receive(), timeout=waiting_time_sec)
-    compare_messages(msg, CapacityMetrics(now, 21.0), 0.2)
+    compare_messages(msg, Sample(now, Energy.from_watt_hours(21.0)), 0.2)
 
 
 async def run_soc_test(setup_args: SetupArgs) -> None: