Switch to the rust component graph

src/frequenz/sdk/timeseries/formulas/__init__.py

@@ -0,0 +1,4 @@
+# License: MIT
+# Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+"""Formulas on telemetry streams."""

src/frequenz/sdk/timeseries/formulas/_ast.py

@@ -0,0 +1,205 @@
+# License: MIT
+# Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+"""Formula AST nodes and evaluation logic."""
+
+from __future__ import annotations
+
+import abc
+import logging
+import math
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Generic
+
+from typing_extensions import override
+
+from ..._internal._math import is_close_to_zero
+from .._base_types import QuantityT, Sample
+from ._functions import Function
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Node(abc.ABC):
+    """An abstract syntax tree node representing a formula expression."""
+
+    span: tuple[int, int] | None
+
+    @abc.abstractmethod
+    def evaluate(self) -> float | None:
+        """Evaluate the expression and return its numerical value."""
+
+    @abc.abstractmethod
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the node."""
+
+    @override
+    def __str__(self) -> str:
+        """Return the string representation of the node."""
+        return self.format()
+
+
+@dataclass
+class TelemetryStream(Node, Generic[QuantityT]):
+    """A AST node that retrieves values from a component's telemetry stream."""
+
+    source: str
+    stream: AsyncIterator[Sample[QuantityT]]
+    _latest_sample: Sample[QuantityT] | None = None
+
+    @property
+    def latest_sample(self) -> Sample[QuantityT] | None:
+        """Return the latest fetched sample for this component."""
+        return self._latest_sample
+
+    @override
+    def evaluate(self) -> float | None:
+        """Return the base value of the latest sample for this component."""
+        if self._latest_sample is None:
+            raise ValueError("Next value has not been fetched yet.")
+        if self._latest_sample.value is None:
+            return None
+        return self._latest_sample.value.base_value
+
+    @override
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the telemetry stream node."""
+        return f"{self.source}"
+
+    async def fetch_next(self) -> None:
+        """Fetch the next value for this component and store it internally."""
+        self._latest_sample = await anext(self.stream)
+
+
+@dataclass
+class FunCall(Node):
+    """A function call in the formula."""
+
+    function: Function
+    args: list[Node]
+
+    @override
+    def evaluate(self) -> float | None:
+        """Evaluate the function call with its arguments."""
+        return self.function(arg.evaluate() for arg in self.args)
+
+    @override
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the function call node."""
+        args_str = ", ".join(str(arg) for arg in self.args)
+        return f"{self.function.name}({args_str})"
+
+
+@dataclass
+class Constant(Node):
+    """A constant numerical value in the formula."""
+
+    value: float
+
+    @override
+    def evaluate(self) -> float | None:
+        """Return the constant value."""
+        return self.value
+
+    @override
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the constant node."""
+        return str(self.value)
+
+
+@dataclass
+class Add(Node):
+    """Addition operation node."""
+
+    left: Node
+    right: Node
+
+    @override
+    def evaluate(self) -> float | None:
+        """Evaluate the addition of the left and right nodes."""
+        left = self.left.evaluate()
+        right = self.right.evaluate()
+        if left is None or right is None:
+            return None
+        return left + right
+
+    @override
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the addition node."""
+        expr = f"{self.left} + {self.right}"
+        if wrap:
+            expr = f"({expr})"
+        return expr
+
+
+@dataclass
+class Sub(Node):
+    """Subtraction operation node."""
+
+    left: Node
+    right: Node
+
+    @override
+    def evaluate(self) -> float | None:
+        """Evaluate the subtraction of the right node from the left node."""
+        left = self.left.evaluate()
+        right = self.right.evaluate()
+        if left is None or right is None:
+            return None
+        return left - right
+
+    @override
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the subtraction node."""
+        expr = f"{self.left} - {self.right.format(True)}"
+        if wrap:
+            expr = f"({expr})"
+        return expr
+
+
+@dataclass
+class Mul(Node):
+    """Multiplication operation node."""
+
+    left: Node
+    right: Node
+
+    @override
+    def evaluate(self) -> float | None:
+        """Evaluate the multiplication of the left and right nodes."""
+        left = self.left.evaluate()
+        right = self.right.evaluate()
+        if left is None or right is None:
+            return None
+        return left * right
+
+    @override
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the multiplication node."""
+        return f"{self.left.format(True)} * {self.right.format(True)}"
+
+
+@dataclass
+class Div(Node):
+    """Division operation node."""
+
+    left: Node
+    right: Node
+
+    @override
+    def evaluate(self) -> float | None:
+        """Evaluate the division of the left node by the right node."""
+        left = self.left.evaluate()
+        right = self.right.evaluate()
+        if left is None or right is None:
+            return None
+        if is_close_to_zero(right):
+            return math.nan
+        return left / right
+
+    @override
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the division node."""
+        return f"{self.left.format(True)} / {self.right.format(True)}"

src/frequenz/sdk/timeseries/formulas/_formula_evaluator.py

@@ -0,0 +1,129 @@
+# License: MIT
+# Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+"""An evaluator for a formula represented as an AST."""
+
+
+import asyncio
+import logging
+from collections.abc import Callable
+from datetime import datetime
+from typing import Generic
+
+from frequenz.channels import Broadcast, ReceiverStoppedError, Sender
+from typing_extensions import override
+
+from ...actor import Actor
+from .._base_types import QuantityT, Sample
+from . import _ast
+from ._resampled_stream_fetcher import ResampledStreamFetcher
+
+_logger = logging.getLogger(__name__)
+
+
+class FormulaEvaluatingActor(Generic[QuantityT], Actor):
+    """An evaluator for a formula represented as an AST."""
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        *,
+        root: _ast.Node,
+        components: list[_ast.TelemetryStream[QuantityT]],
+        create_method: Callable[[float], QuantityT],
+        output_channel: Broadcast[Sample[QuantityT]],
+        metric_fetcher: ResampledStreamFetcher | None = None,
+    ) -> None:
+        """Create a `FormulaEvaluatingActor` instance.
+
+        Args:
+            root: The root node of the formula AST.
+            components: The telemetry streams that the formula depends on.
+            create_method: A method to generate the output values with.  If the
+                formula is for generating power values, this would be
+                `Power.from_watts`, for example.
+            output_channel: The channel to send evaluated samples to.
+            metric_fetcher: An optional metric fetcher that needs to be started
+                before the formula can be evaluated.
+        """
+        super().__init__()
+
+        self._root: _ast.Node = root
+        self._components: list[_ast.TelemetryStream[QuantityT]] = components
+        self._create_method: Callable[[float], QuantityT] = create_method
+        self._metric_fetcher: ResampledStreamFetcher | None = metric_fetcher
+        self._output_channel: Broadcast[Sample[QuantityT]] = output_channel
+
+        self._output_sender: Sender[Sample[QuantityT]] = output_channel.new_sender()
+
+    @override
+    async def _run(self) -> None:
+        """Run the formula evaluator actor."""
+        if self._metric_fetcher is not None:
+            await self._metric_fetcher.subscribe()
+        await synchronize_receivers(self._components)
+
+        while True:
+            try:
+                timestamp = next(
+                    comp.latest_sample.timestamp
+                    for comp in self._components
+                    if comp.latest_sample is not None
+                )
+
+                res = self._root.evaluate()
+                next_sample = Sample(
+                    timestamp, None if res is None else self._create_method(res)
+                )
+                await self._output_sender.send(next_sample)
+            except (StopAsyncIteration, StopIteration):
+                _logger.debug(
+                    "No more input samples available; stopping formula evaluator. (%s)",
+                    self._root,
+                )
+                await self._output_channel.close()
+                return
+            except Exception as e:  # pylint: disable=broad-except
+                _logger.error(
+                    "Error evaluating formula %s: %s", self._root, e, exc_info=True
+                )
+                await self._output_channel.close()
+                return
+
+            fetch_results = await asyncio.gather(
+                *(comp.fetch_next() for comp in self._components),
+                return_exceptions=True,
+            )
+            if ex := next((e for e in fetch_results if isinstance(e, Exception)), None):
+                if isinstance(ex, (StopAsyncIteration, ReceiverStoppedError)):
+                    _logger.debug(
+                        "input streams closed; stopping formula evaluator. (%s)",
+                        self._root,
+                    )
+                    await self._output_channel.close()
+                    return
+                raise ex
+
+
+async def synchronize_receivers(
+    components: list[_ast.TelemetryStream[QuantityT]],
+) -> None:
+    """Synchronize the given telemetry stream receivers."""
+    _ = await asyncio.gather(
+        *(comp.fetch_next() for comp in components),
+    )
+    latest_ts: datetime | None = None
+    for comp in components:
+        if comp.latest_sample is not None and (
+            latest_ts is None or comp.latest_sample.timestamp > latest_ts
+        ):
+            latest_ts = comp.latest_sample.timestamp
+    if latest_ts is None:
+        _logger.debug("No samples available to synchronize receivers.")
+        return
+    for comp in components:
+        if comp.latest_sample is None:
+            raise RuntimeError("Can't synchronize receivers.")
+        ctr = 0
+        while ctr < 10 and comp.latest_sample.timestamp < latest_ts:
+            await comp.fetch_next()
+            ctr += 1