ezmsg-sigproc 1.8.2__py3-none-any.whl → 2.1.0__py3-none-any.whl

ezmsg/sigproc/affinetransform.py

@@ -1,112 +1,112 @@
 import os
 from pathlib import Path
-import typing
 
 import numpy as np
 import numpy.typing as npt
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray, AxisBase
 from ezmsg.util.messages.util import replace
-from ezmsg.util.generator import consumer
 
-from .base import GenAxisArray
+from .base import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    BaseTransformer,
+    processor_state,
+)
 
 
-@consumer
-def affine_transform(
-    weights: np.ndarray | str | Path,
-    axis: str | None = None,
-    right_multiply: bool = True,
-) -> typing.Generator[AxisArray, AxisArray, None]:
+class AffineTransformSettings(ez.Settings):
+    """
+    Settings for :obj:`AffineTransform`.
+    See :obj:`affine_transform` for argument details.
     """
-    Perform affine transformations on streaming data.
 
-    Args:
-        weights: An array of weights or a path to a file with weights compatible with np.loadtxt.
-        axis: The name of the axis to apply the transformation to. Defaults to the leading (0th) axis in the array.
-        right_multiply: Set False to transpose the weights before applying.
+    weights: np.ndarray | str | Path
+    """An array of weights or a path to a file with weights compatible with np.loadtxt."""
 
-    Returns:
-        A primed generator object that yields an :obj:`AxisArray` object for every
-        :obj:`AxisArray` it receives via `send`.
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
+    axis: str | None = None
+    """The name of the axis to apply the transformation to. Defaults to the leading (0th) axis in the array."""
 
-    # Check parameters
-    if isinstance(weights, str):
-        if weights == "passthrough":
-            weights = None
-        else:
-            weights = Path(os.path.abspath(os.path.expanduser(weights)))
-    if isinstance(weights, Path):
-        weights = np.loadtxt(weights, delimiter=",")
-    if not right_multiply:
-        weights = weights.T
-    if weights is not None:
-        weights = np.ascontiguousarray(weights)
-
-    # State variables
-    # New axis with transformed labels, if required
+    right_multiply: bool = True
+    """Set False to transpose the weights before applying."""
+
+
+@processor_state
+class AffineTransformState:
+    weights: npt.NDArray | None = None
     new_axis: AxisBase | None = None
 
-    # Reset if any of these change.
-    check_input = {"key": None}
-    # We assume key change catches labels change; we don't want to check labels every message
-    # We don't need to check if input size has changed because weights multiplication will fail if so.
-
-    while True:
-        msg_in: AxisArray = yield msg_out
-
-        if weights is None:
-            msg_out = msg_in
-            continue
-
-        axis = axis or msg_in.dims[-1]  # Note: Most nodes default do dim[0]
-        axis_idx = msg_in.get_axis_idx(axis)
-
-        b_reset = msg_in.key != check_input["key"]
-        if b_reset:
-            # First sample or key has changed. Reset the state.
-            check_input["key"] = msg_in.key
-            # Determine if we need to modify the transformed axis.
-            if (
-                axis in msg_in.axes
-                and hasattr(msg_in.axes[axis], "data")
-                and weights.shape[0] != weights.shape[1]
-            ):
-                in_labels = msg_in.axes[axis].data
-                new_labels = []
-                n_in, n_out = weights.shape
-                if len(in_labels) != n_in:
-                    # Something upstream did something it wasn't supposed to. We will drop the labels.
-                    ez.logger.warning(
-                        f"Received {len(in_labels)} for {n_in} inputs. Check upstream labels."
-                    )
-                else:
-                    b_filled_outputs = np.any(weights, axis=0)
-                    b_used_inputs = np.any(weights, axis=1)
-                    if np.all(b_used_inputs) and np.all(b_filled_outputs):
-                        # All inputs are used and all outputs are used, but n_in != n_out.
-                        # Mapping cannot be determined.
-                        new_labels = []
-                    elif np.all(b_used_inputs):
-                        # Strange scenario: New outputs are filled with empty data.
-                        in_ix = 0
-                        new_labels = []
-                        for out_ix in range(n_out):
-                            if b_filled_outputs[out_ix]:
-                                new_labels.append(in_labels[in_ix])
-                                in_ix += 1
-                            else:
-                                new_labels.append("")
-                    elif np.all(b_filled_outputs):
-                        # Transform is dropping some of the inputs.
-                        new_labels = np.array(in_labels)[b_used_inputs]
-                new_axis = replace(msg_in.axes[axis], data=np.array(new_labels))
-
-        data = msg_in.data
-
-        if data.shape[axis_idx] == (weights.shape[0] - 1):
+
+class AffineTransformTransformer(
+    BaseStatefulTransformer[
+        AffineTransformSettings, AxisArray, AxisArray, AffineTransformState
+    ]
+):
+    def __call__(self, message: AxisArray) -> AxisArray:
+        # Override __call__ so we can shortcut if weights are None.
+        if self.settings.weights is None or (
+            isinstance(self.settings.weights, str)
+            and self.settings.weights == "passthrough"
+        ):
+            return message
+        return super().__call__(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        return hash(message.key)
+
+    def _reset_state(self, message: AxisArray) -> None:
+        weights = self.settings.weights
+        if isinstance(weights, str):
+            weights = Path(os.path.abspath(os.path.expanduser(weights)))
+        if isinstance(weights, Path):
+            weights = np.loadtxt(weights, delimiter=",")
+        if not self.settings.right_multiply:
+            weights = weights.T
+        if weights is not None:
+            weights = np.ascontiguousarray(weights)
+
+        self._state.weights = weights
+
+        axis = self.settings.axis or message.dims[-1]
+        if (
+            axis in message.axes
+            and hasattr(message.axes[axis], "data")
+            and weights.shape[0] != weights.shape[1]
+        ):
+            in_labels = message.axes[axis].data
+            new_labels = []
+            n_in, n_out = weights.shape
+            if len(in_labels) != n_in:
+                ez.logger.warning(
+                    f"Received {len(in_labels)} for {n_in} inputs. Check upstream labels."
+                )
+            else:
+                b_filled_outputs = np.any(weights, axis=0)
+                b_used_inputs = np.any(weights, axis=1)
+                if np.all(b_used_inputs) and np.all(b_filled_outputs):
+                    new_labels = []
+                elif np.all(b_used_inputs):
+                    in_ix = 0
+                    new_labels = []
+                    for out_ix in range(n_out):
+                        if b_filled_outputs[out_ix]:
+                            new_labels.append(in_labels[in_ix])
+                            in_ix += 1
+                        else:
+                            new_labels.append("")
+                elif np.all(b_filled_outputs):
+                    new_labels = np.array(in_labels)[b_used_inputs]
+
+            self._state.new_axis = replace(
+                message.axes[axis], data=np.array(new_labels)
+            )
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+        data = message.data
+
+        if data.shape[axis_idx] == (self._state.weights.shape[0] - 1):
             # The weights are stacked A|B where A is the transform and B is a single row
             #  in the equation y = Ax + B. This supports NeuroKey's weights matrices.
             sample_shape = data.shape[:axis_idx] + (1,) + data.shape[axis_idx + 1 :]
@@ -114,82 +114,87 @@ def affine_transform(
                 (data, np.ones(sample_shape).astype(data.dtype)), axis=axis_idx
             )
 
-        if axis_idx in [-1, len(msg_in.dims) - 1]:
-            data = np.matmul(data, weights)
+        if axis_idx in [-1, len(message.dims) - 1]:
+            data = np.matmul(data, self._state.weights)
         else:
             data = np.moveaxis(data, axis_idx, -1)
-            data = np.matmul(data, weights)
+            data = np.matmul(data, self._state.weights)
             data = np.moveaxis(data, -1, axis_idx)
 
         replace_kwargs = {"data": data}
-        if new_axis is not None:
-            replace_kwargs["axes"] = {**msg_in.axes, axis: new_axis}
-        msg_out = replace(msg_in, **replace_kwargs)
+        if self._state.new_axis is not None:
+            replace_kwargs["axes"] = {**message.axes, axis: self._state.new_axis}
 
+        return replace(message, **replace_kwargs)
 
-class AffineTransformSettings(ez.Settings):
-    """
-    Settings for :obj:`AffineTransform`.
-    See :obj:`affine_transform` for argument details.
-    """
 
-    weights: np.ndarray | str | Path
-    axis: str | None = None
-    right_multiply: bool = True
+class AffineTransform(
+    BaseTransformerUnit[
+        AffineTransformSettings, AxisArray, AxisArray, AffineTransformTransformer
+    ]
+):
+    SETTINGS = AffineTransformSettings
 
 
-class AffineTransform(GenAxisArray):
-    """:obj:`Unit` for :obj:`affine_transform`"""
+def affine_transform(
+    weights: np.ndarray | str | Path,
+    axis: str | None = None,
+    right_multiply: bool = True,
+) -> AffineTransformTransformer:
+    """
+    Perform affine transformations on streaming data.
 
-    SETTINGS = AffineTransformSettings
+    Args:
+        weights: An array of weights or a path to a file with weights compatible with np.loadtxt.
+        axis: The name of the axis to apply the transformation to. Defaults to the leading (0th) axis in the array.
+        right_multiply: Set False to transpose the weights before applying.
 
-    def construct_generator(self):
-        self.STATE.gen = affine_transform(
-            weights=self.SETTINGS.weights,
-            axis=self.SETTINGS.axis,
-            right_multiply=self.SETTINGS.right_multiply,
+    Returns:
+        :obj:`AffineTransformTransformer`.
+    """
+    return AffineTransformTransformer(
+        AffineTransformSettings(
+            weights=weights, axis=axis, right_multiply=right_multiply
         )
+    )
 
 
 def zeros_for_noop(data: npt.NDArray, **ignore_kwargs) -> npt.NDArray:
     return np.zeros_like(data)
 
 
-@consumer
-def common_rereference(
-    mode: str = "mean", axis: str | None = None, include_current: bool = True
-) -> typing.Generator[AxisArray, AxisArray, None]:
+class CommonRereferenceSettings(ez.Settings):
+    """
+    Settings for :obj:`CommonRereference`
     """
-    Perform common average referencing (CAR) on streaming data.
 
-    Args:
-        mode: The statistical mode to apply -- either "mean" or "median"
-        axis: The name of hte axis to apply the transformation to.
-        include_current: Set False to exclude each channel from participating in the calculation of its reference.
+    mode: str = "mean"
+    """The statistical mode to apply -- either "mean" or "median"."""
 
-    Returns:
-        A primed generator object that yields an :obj:`AxisArray` object
-        for every :obj:`AxisArray` it receives via `send`.
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
+    axis: str | None = None
+    """The name of the axis to apply the transformation to."""
 
-    if mode == "passthrough":
-        include_current = True
+    include_current: bool = True
+    """Set False to exclude each channel from participating in the calculation of its reference."""
 
-    func = {"mean": np.mean, "median": np.median, "passthrough": zeros_for_noop}[mode]
 
-    while True:
-        msg_in: AxisArray = yield msg_out
+class CommonRereferenceTransformer(
+    BaseTransformer[CommonRereferenceSettings, AxisArray, AxisArray]
+):
+    def _process(self, message: AxisArray) -> AxisArray:
+        if self.settings.mode == "passthrough":
+            return message
 
-        if axis is None:
-            axis = msg_in.dims[-1]
-            axis_idx = -1
-        else:
-            axis_idx = msg_in.get_axis_idx(axis)
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+
+        func = {"mean": np.mean, "median": np.median, "passthrough": zeros_for_noop}[
+            self.settings.mode
+        ]
 
-        ref_data = func(msg_in.data, axis=axis_idx, keepdims=True)
+        ref_data = func(message.data, axis=axis_idx, keepdims=True)
 
-        if not include_current:
+        if not self.settings.include_current:
             # Typical `CAR = x[0]/N + x[1]/N + ... x[i-1]/N + x[i]/N + x[i+1]/N + ... + x[N-1]/N`
             # and is the same for all i, so it is calculated only once in `ref_data`.
             # However, if we had excluded the current channel,
@@ -200,34 +205,35 @@ def common_rereference(
             # from the current channel (i.e., `x[i] / (N-1)`)
             #  i.e., `CAR[i] = (N / (N-1)) * common_CAR - x[i]/(N-1)`
             # We can use broadcasting subtraction instead of looping over channels.
-            N = msg_in.data.shape[axis_idx]
-            ref_data = (N / (N - 1)) * ref_data - msg_in.data / (N - 1)
-            # Side note: I profiled using affine_transform and it's about 30x slower than this implementation.
+            N = message.data.shape[axis_idx]
+            ref_data = (N / (N - 1)) * ref_data - message.data / (N - 1)
+            # Note: I profiled using AffineTransformTransformer; it's ~30x slower than this implementation.
 
-        msg_out = replace(msg_in, data=msg_in.data - ref_data)
+        return replace(message, data=message.data - ref_data)
 
 
-class CommonRereferenceSettings(ez.Settings):
-    """
-    Settings for :obj:`CommonRereference`
-    See :obj:`common_rereference` for argument details.
-    """
-
-    mode: str = "mean"
-    axis: str | None = None
-    include_current: bool = True
+class CommonRereference(
+    BaseTransformerUnit[
+        CommonRereferenceSettings, AxisArray, AxisArray, CommonRereferenceTransformer
+    ]
+):
+    SETTINGS = CommonRereferenceSettings
 
 
-class CommonRereference(GenAxisArray):
-    """
-    :obj:`Unit` for :obj:`common_rereference`.
+def common_rereference(
+    mode: str = "mean", axis: str | None = None, include_current: bool = True
+) -> CommonRereferenceTransformer:
     """
+    Perform common average referencing (CAR) on streaming data.
 
-    SETTINGS = CommonRereferenceSettings
+    Args:
+        mode: The statistical mode to apply -- either "mean" or "median"
+        axis: The name of hte axis to apply the transformation to.
+        include_current: Set False to exclude each channel from participating in the calculation of its reference.
 
-    def construct_generator(self):
-        self.STATE.gen = common_rereference(
-            mode=self.SETTINGS.mode,
-            axis=self.SETTINGS.axis,
-            include_current=self.SETTINGS.include_current,
-        )
+    Returns:
+        :obj:`CommonRereferenceTransformer`
+    """
+    return CommonRereferenceTransformer(
+        CommonRereferenceSettings(mode=mode, axis=axis, include_current=include_current)
+    )