ezmsg-sigproc 1.8.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

ezmsg/sigproc/scaler.py

@@ -1,163 +1,21 @@
-import functools
 import typing
 
 import numpy as np
-import numpy.typing as npt
-import scipy.signal
-import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
-from .base import GenAxisArray
+from .base import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from .ewma import EWMATransformer, EWMASettings, _alpha_from_tau
 
-
-def _tau_from_alpha(alpha: float, dt: float) -> float:
-    """
-    Inverse of _alpha_from_tau. See that function for explanation.
-    """
-    return -dt / np.log(1 - alpha)
-
-
-def _alpha_from_tau(tau: float, dt: float) -> float:
-    """
-    # https://en.wikipedia.org/wiki/Exponential_smoothing#Time_constant
-    :param tau: The amount of time for the smoothed response of a unit step function to reach
-        1 - 1/e approx-eq 63.2%.
-    :param dt: sampling period, or 1 / sampling_rate.
-    :return: alpha, the "fading factor" in exponential smoothing.
-    """
-    return 1 - np.exp(-dt / tau)
-
-
-def ewma_step(
-    sample: npt.NDArray, zi: npt.NDArray, alpha: float, beta: float | None = None
-):
-    """
-    Do an exponentially weighted moving average step.
-
-    Args:
-        sample: The new sample.
-        zi: The output of the previous step.
-        alpha: Fading factor.
-        beta: Persisting factor. If None, it is calculated as 1-alpha.
-
-    Returns:
-        alpha * sample + beta * zi
-
-    """
-    # Potential micro-optimization:
-    #  Current: scalar-arr multiplication, scalar-arr multiplication, arr-arr addition
-    #  Alternative: arr-arr subtraction, arr-arr multiplication, arr-arr addition
-    # return zi + alpha * (new_sample - zi)
-    beta = beta or (1 - alpha)
-    return alpha * sample + beta * zi
-
-
-class EWMA:
-    def __init__(self, alpha: float):
-        self.beta = 1 - alpha
-        self._filt_func = functools.partial(
-            scipy.signal.lfilter, [alpha], [1.0, alpha - 1.0], axis=0
-        )
-        self.prev = None
-
-    def compute(self, arr: npt.NDArray) -> npt.NDArray:
-        if self.prev is None:
-            self.prev = self.beta * arr[:1]
-        expected, self.prev = self._filt_func(arr, zi=self.prev)
-        return expected
-
-
-class EWMA_Deprecated:
-    """
-    Grabbed these methods from https://stackoverflow.com/a/70998068 and other answers in that topic,
-    but they ended up being slower than the scipy.signal.lfilter method.
-    Additionally, `compute` and `compute2` suffer from potential errors as the vector length increases
-    and beta**n approaches zero.
-    """
-
-    def __init__(self, alpha: float, max_len: int):
-        self.alpha = alpha
-        self.beta = 1 - alpha
-        self.prev: npt.NDArray | None = None
-        self.weights = np.empty((max_len + 1,), float)
-        self._precalc_weights(max_len)
-        self._step_func = functools.partial(ewma_step, alpha=self.alpha, beta=self.beta)
-
-    def _precalc_weights(self, n: int):
-        #   (1-α)^0, (1-α)^1, (1-α)^2, ..., (1-α)^n
-        np.power(self.beta, np.arange(n + 1), out=self.weights)
-
-    def compute(self, arr: npt.NDArray, out: npt.NDArray | None = None) -> npt.NDArray:
-        if out is None:
-            out = np.empty(arr.shape, arr.dtype)
-
-        n = arr.shape[0]
-        weights = self.weights[:n]
-        weights = np.expand_dims(weights, list(range(1, arr.ndim)))
-
-        #   α*P0, α*P1, α*P2, ..., α*Pn
-        np.multiply(self.alpha, arr, out)
-
-        #   α*P0/(1-α)^0, α*P1/(1-α)^1, α*P2/(1-α)^2, ..., α*Pn/(1-α)^n
-        np.divide(out, weights, out)
-
-        #   α*P0/(1-α)^0, α*P0/(1-α)^0 + α*P1/(1-α)^1, ...
-        np.cumsum(out, axis=0, out=out)
-
-        #   (α*P0/(1-α)^0)*(1-α)^0, (α*P0/(1-α)^0 + α*P1/(1-α)^1)*(1-α)^1, ...
-        np.multiply(out, weights, out)
-
-        # Add the previous output
-        if self.prev is None:
-            self.prev = arr[:1]
-
-        out += self.prev * np.expand_dims(
-            self.weights[1 : n + 1], list(range(1, arr.ndim))
-        )
-
-        self.prev = out[-1:]
-
-        return out
-
-    def compute2(self, arr: npt.NDArray) -> npt.NDArray:
-        """
-        Compute the Exponentially Weighted Moving Average (EWMA) of the input array.
-
-        Args:
-            arr: The input array to be smoothed.
-
-        Returns:
-            The smoothed array.
-        """
-        n = arr.shape[0]
-        if n > len(self.weights):
-            self._precalc_weights(n)
-        weights = self.weights[:n][::-1]
-        weights = np.expand_dims(weights, list(range(1, arr.ndim)))
-
-        result = np.cumsum(self.alpha * weights * arr, axis=0)
-        result = result / weights
-
-        # Handle the first call when prev is unset
-        if self.prev is None:
-            self.prev = arr[:1]
-
-        result += self.prev * np.expand_dims(
-            self.weights[1 : n + 1], list(range(1, arr.ndim))
-        )
-
-        # Store the result back into prev
-        self.prev = result[-1]
-
-        return result
-
-    def compute_sample(self, new_sample: npt.NDArray) -> npt.NDArray:
-        if self.prev is None:
-            self.prev = new_sample
-        self.prev = self._step_func(new_sample, self.prev)
-        return self.prev
+# Imports for backwards compatibility with previous module location
+from .ewma import EWMA_Deprecated as EWMA_Deprecated
+from .ewma import ewma_step as ewma_step
+from .ewma import _tau_from_alpha as _tau_from_alpha
 
 
 @consumer
@@ -208,83 +66,62 @@ def scaler(
         msg_out = replace(msg_in, data=result)
 
 
-@consumer
-def scaler_np(
-    time_constant: float = 1.0, axis: str | None = None
-) -> typing.Generator[AxisArray, AxisArray, None]:
-    """
-    Create a generator function that applies an adaptive standard scaler.
-    This is faster than :obj:`scaler` for multichannel data.
-
-    Args:
-        time_constant: Decay constant `tau` in seconds.
-        axis: The name of the axis to accumulate statistics over.
-            Note: The axis must be in the msg.axes and be of type AxisArray.LinearAxis.
-
-    Returns:
-        A primed generator object that expects to be sent a :obj:`AxisArray` via `.send(axis_array)`
-         and yields an :obj:`AxisArray` with its data being a standardized, or "Z-scored" version of the input data.
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
-
-    # State variables
-    samps_ewma: EWMA | None = None
-    vars_sq_ewma: EWMA | None = None
+class AdaptiveStandardScalerSettings(EWMASettings): ...
 
-    # Reset if input changes
-    check_input = {
-        "gain": None,  # Resets alpha
-        "shape": None,
-        "key": None,  # Key change implies buffered means/vars are invalid.
-    }
 
-    while True:
-        msg_in: AxisArray = yield msg_out
+@processor_state
+class AdaptiveStandardScalerState:
+    samps_ewma: EWMATransformer | None = None
+    vars_sq_ewma: EWMATransformer | None = None
+    alpha: float | None = None
 
-        axis = axis or msg_in.dims[0]
-        axis_idx = msg_in.get_axis_idx(axis)
 
-        data: npt.NDArray = np.moveaxis(msg_in.data, axis_idx, 0)
-        b_reset = data.shape[1:] != check_input["shape"]
-        b_reset = b_reset or msg_in.axes[axis].gain != check_input["gain"]
-        b_reset = b_reset or msg_in.key != check_input["key"]
-        if b_reset:
-            check_input["shape"] = data.shape[1:]
-            check_input["gain"] = msg_in.axes[axis].gain
-            check_input["key"] = msg_in.key
-            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
-            samps_ewma = EWMA(alpha=alpha)
-            vars_sq_ewma = EWMA(alpha=alpha)
+class AdaptiveStandardScalerTransformer(
+    BaseStatefulTransformer[
+        AdaptiveStandardScalerSettings,
+        AxisArray,
+        AxisArray,
+        AdaptiveStandardScalerState,
+    ]
+):
+    def _reset_state(self, message: AxisArray) -> None:
+        self._state.samps_ewma = EWMATransformer(
+            time_constant=self.settings.time_constant, axis=self.settings.axis
+        )
+        self._state.vars_sq_ewma = EWMATransformer(
+            time_constant=self.settings.time_constant, axis=self.settings.axis
+        )
 
+    def _process(self, message: AxisArray) -> AxisArray:
         # Update step
-        means = samps_ewma.compute(data)
-        vars_sq_means = vars_sq_ewma.compute(data**2)
+        mean_message = self._state.samps_ewma(message)
+        var_sq_message = self._state.vars_sq_ewma(
+            replace(message, data=message.data**2)
+        )
 
         # Get step
-        varis = vars_sq_means - means**2
+        varis = var_sq_message.data - mean_message.data**2
         with np.errstate(divide="ignore", invalid="ignore"):
-            result = (data - means) / (varis**0.5)
+            result = (message.data - mean_message.data) / (varis**0.5)
         result[np.isnan(result)] = 0.0
-        result = np.moveaxis(result, 0, axis_idx)
-        msg_out = replace(msg_in, data=result)
+        return replace(message, data=result)
 
 
-class AdaptiveStandardScalerSettings(ez.Settings):
-    """
-    Settings for :obj:`AdaptiveStandardScaler`.
-    See :obj:`scaler_np` for a description of the parameters.
-    """
-
-    time_constant: float = 1.0
-    axis: str | None = None
-
-
-class AdaptiveStandardScaler(GenAxisArray):
-    """Unit for :obj:`scaler_np`"""
-
+class AdaptiveStandardScaler(
+    BaseTransformerUnit[
+        AdaptiveStandardScalerSettings,
+        AxisArray,
+        AxisArray,
+        AdaptiveStandardScalerTransformer,
+    ]
+):
     SETTINGS = AdaptiveStandardScalerSettings
 
-    def construct_generator(self):
-        self.STATE.gen = scaler_np(
-            time_constant=self.SETTINGS.time_constant, axis=self.SETTINGS.axis
-        )
+
+# Backwards compatibility...
+def scaler_np(
+    time_constant: float = 1.0, axis: str | None = None
+) -> AdaptiveStandardScalerTransformer:
+    return AdaptiveStandardScalerTransformer(
+        settings=AdaptiveStandardScalerSettings(time_constant=time_constant, axis=axis)
+    )

ezmsg/sigproc/signalinjector.py

@@ -1,12 +1,14 @@
-import typing
-
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 import numpy as np
 import numpy.typing as npt
 
-from .util.profile import profile_subpub
+from .base import (
+    BaseAsyncTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 
 
 class SignalInjectorSettings(ez.Settings):
@@ -16,57 +18,64 @@ class SignalInjectorSettings(ez.Settings):
     mixing_seed: int | None = None
 
 
-class SignalInjectorState(ez.State):
+@processor_state
+class SignalInjectorState:
     cur_shape: tuple[int, ...] | None = None
     cur_frequency: float | None = None
-    cur_amplitude: float
-    mixing: npt.NDArray
+    cur_amplitude: float | None = None
+    mixing: npt.NDArray | None = None
 
 
-class SignalInjector(ez.Unit):
-    """
-    Add a sinusoidal signal to the input signal. Each feature gets a different amplitude of the sinusoid.
-    All features get the same frequency sinusoid. The frequency and base amplitude can be changed while running.
-    """
+class SignalInjectorTransformer(
+    BaseAsyncTransformer[
+        SignalInjectorSettings, AxisArray, AxisArray, SignalInjectorState
+    ]
+):
+    def _hash_message(self, message: AxisArray) -> int:
+        time_ax_idx = message.get_axis_idx(self.settings.time_dim)
+        sample_shape = (
+            message.data.shape[:time_ax_idx] + message.data.shape[time_ax_idx + 1 :]
+        )
+        return hash((message.key,) + sample_shape)
+
+    def _reset_state(self, message: AxisArray) -> None:
+        if self._state.cur_frequency is None:
+            self._state.cur_frequency = self.settings.frequency
+        if self._state.cur_amplitude is None:
+            self._state.cur_amplitude = self.settings.amplitude
+        time_ax_idx = message.get_axis_idx(self.settings.time_dim)
+        self._state.cur_shape = (
+            message.data.shape[:time_ax_idx] + message.data.shape[time_ax_idx + 1 :]
+        )
+        rng = np.random.default_rng(self.settings.mixing_seed)
+        self._state.mixing = rng.random((1, message.shape2d(self.settings.time_dim)[1]))
+        self._state.mixing = (self._state.mixing * 2.0) - 1.0
+
+    async def _aprocess(self, message: AxisArray) -> AxisArray:
+        if self._state.cur_frequency is None:
+            return message
+        out_msg = replace(message, data=message.data.copy())
+        t = out_msg.ax(self.settings.time_dim).values[..., np.newaxis]
+        signal = np.sin(2 * np.pi * self._state.cur_frequency * t)
+        mixed_signal = signal * self._state.mixing * self._state.cur_amplitude
+        with out_msg.view2d(self.settings.time_dim) as view:
+            view[...] = view + mixed_signal.astype(view.dtype)
+        return out_msg
 
-    SETTINGS = SignalInjectorSettings
-    STATE = SignalInjectorState
 
+class SignalInjector(
+    BaseTransformerUnit[
+        SignalInjectorSettings, AxisArray, AxisArray, SignalInjectorTransformer
+    ]
+):
+    SETTINGS = SignalInjectorSettings
     INPUT_FREQUENCY = ez.InputStream(float | None)
     INPUT_AMPLITUDE = ez.InputStream(float)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
-
-    async def initialize(self) -> None:
-        self.STATE.cur_frequency = self.SETTINGS.frequency
-        self.STATE.cur_amplitude = self.SETTINGS.amplitude
-        self.STATE.mixing = np.array([])
 
     @ez.subscriber(INPUT_FREQUENCY)
     async def on_frequency(self, msg: float | None) -> None:
-        self.STATE.cur_frequency = msg
+        self.processor.state.cur_frequency = msg
 
     @ez.subscriber(INPUT_AMPLITUDE)
     async def on_amplitude(self, msg: float) -> None:
-        self.STATE.cur_amplitude = msg
-
-    @ez.subscriber(INPUT_SIGNAL)
-    @ez.publisher(OUTPUT_SIGNAL)
-    @profile_subpub(trace_oldest=False)
-    async def inject(self, msg: AxisArray) -> typing.AsyncGenerator:
-        if self.STATE.cur_shape != msg.shape:
-            self.STATE.cur_shape = msg.shape
-            rng = np.random.default_rng(self.SETTINGS.mixing_seed)
-            self.STATE.mixing = rng.random((1, msg.shape2d(self.SETTINGS.time_dim)[1]))
-            self.STATE.mixing = (self.STATE.mixing * 2.0) - 1.0
-
-        if self.STATE.cur_frequency is None:
-            yield self.OUTPUT_SIGNAL, msg
-        else:
-            out_msg = replace(msg, data=msg.data.copy())
-            t = out_msg.ax(self.SETTINGS.time_dim).values[..., np.newaxis]
-            signal = np.sin(2 * np.pi * self.STATE.cur_frequency * t)
-            mixed_signal = signal * self.STATE.mixing * self.STATE.cur_amplitude
-            with out_msg.view2d(self.SETTINGS.time_dim) as view:
-                view[...] = view + mixed_signal.astype(view.dtype)
-            yield self.OUTPUT_SIGNAL, out_msg
+        self.processor.state.cur_amplitude = msg

ezmsg/sigproc/slicer.py

@@ -1,5 +1,3 @@
-import typing
-
 import numpy as np
 import numpy.typing as npt
 import ezmsg.core as ez
@@ -9,10 +7,12 @@ from ezmsg.util.messages.axisarray import (
     AxisBase,
     replace,
 )
-from ezmsg.util.generator import consumer
-
-from .base import GenAxisArray
 
+from .base import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 
 """
 Slicer:Select a subset of data along a particular axis.
@@ -61,106 +61,98 @@ def parse_slice(
     return tuple([item for sublist in suplist for item in sublist])
 
 
-@consumer
-def slicer(
-    selection: str = "", axis: str | None = None
-) -> typing.Generator[AxisArray, AxisArray, None]:
-    """
-    Slice along a particular axis.
-
-    Args:
-        selection: See :obj:`ezmsg.sigproc.slicer.parse_slice` for details.
-        axis: The name of the axis to slice along. If None, the last axis is used.
+class SlicerSettings(ez.Settings):
+    selection: str = ""
+    """selection: See :obj:`ezmsg.sigproc.slicer.parse_slice` for details."""
 
-    Returns:
-        A primed generator object ready to yield an :obj:`AxisArray` for each .send(axis_array)
-        with the data payload containing a sliced view of the input data.
+    axis: str | None = None
+    """The name of the axis to slice along. If None, the last axis is used."""
 
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
 
-    # State variables
-    _slice: slice | npt.NDArray | None = None
+@processor_state
+class SlicerState:
+    slice_: slice | int | npt.NDArray | None = None
     new_axis: AxisBase | None = None
-    b_change_dims: bool = False  # If number of dimensions changes when slicing
-
-    # Reset if input changes
-    check_input = {
-        "key": None,  # key change used as proxy for label change, which we don't check explicitly
-        "len": None,
-    }
-
-    while True:
-        msg_in: AxisArray = yield msg_out
-
-        axis = axis or msg_in.dims[-1]
-        axis_idx = msg_in.get_axis_idx(axis)
-
-        b_reset = _slice is None  # or new_axis is None
-        b_reset = b_reset or msg_in.key != check_input["key"]
-        b_reset = b_reset or (
-            (msg_in.data.shape[axis_idx] != check_input["len"])
-            and (type(_slice) is np.ndarray)
-        )
-        if b_reset:
-            check_input["key"] = msg_in.key
-            check_input["len"] = msg_in.data.shape[axis_idx]
-            new_axis = None  # Will hold updated metadata
-            b_change_dims = False
-
-            # Calculate the slice
-            _slices = parse_slice(selection, msg_in.axes.get(axis, None))
-            if len(_slices) == 1:
-                _slice = _slices[0]
-                # Do we drop the sliced dimension?
-                b_change_dims = isinstance(_slice, int)
+    b_change_dims: bool = False
+
+
+class SlicerTransformer(
+    BaseStatefulTransformer[SlicerSettings, AxisArray, AxisArray, SlicerState]
+):
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+        return hash((message.key, message.data.shape[axis_idx]))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+        self._state.new_axis = None
+        self._state.b_change_dims = False
+
+        # Calculate the slice
+        _slices = parse_slice(self.settings.selection, message.axes.get(axis, None))
+        if len(_slices) == 1:
+            self._state.slice_ = _slices[0]
+            self._state.b_change_dims = isinstance(self._state.slice_, int)
+        else:
+            indices = np.arange(message.data.shape[axis_idx])
+            indices = np.hstack([indices[_] for _ in _slices])
+            self._state.slice_ = np.s_[indices]
+
+        # Create the output axis
+        if (
+            axis in message.axes
+            and hasattr(message.axes[axis], "data")
+            and len(message.axes[axis].data) > 0
+        ):
+            in_data = np.array(message.axes[axis].data)
+            if self._state.b_change_dims:
+                out_data = in_data[self._state.slice_ : self._state.slice_ + 1]
             else:
-                # Multiple slices, but this cannot be done in a single step, so we convert the slices
-                #  to a discontinuous set of integer indexes.
-                indices = np.arange(msg_in.data.shape[axis_idx])
-                indices = np.hstack([indices[_] for _ in _slices])
-                _slice = np.s_[indices]  # Integer scalar array
+                out_data = in_data[self._state.slice_]
+            self._state.new_axis = replace(message.axes[axis], data=out_data)
 
-            # Create the output axis.
-            if (
-                axis in msg_in.axes
-                and hasattr(msg_in.axes[axis], "data")
-                and len(msg_in.axes[axis].data) > 0
-            ):
-                in_data = np.array(msg_in.axes[axis].data)
-                if b_change_dims:
-                    out_data = in_data[_slice : _slice + 1]
-                else:
-                    out_data = in_data[_slice]
-                new_axis = replace(msg_in.axes[axis], data=out_data)
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
 
         replace_kwargs = {}
-        if b_change_dims:
-            # Dropping the target axis
+        if self._state.b_change_dims:
             replace_kwargs["dims"] = [
-                _ for dim_ix, _ in enumerate(msg_in.dims) if dim_ix != axis_idx
+                _ for dim_ix, _ in enumerate(message.dims) if dim_ix != axis_idx
             ]
-            replace_kwargs["axes"] = {k: v for k, v in msg_in.axes.items() if k != axis}
-        elif new_axis is not None:
             replace_kwargs["axes"] = {
-                k: (v if k != axis else new_axis) for k, v in msg_in.axes.items()
+                k: v for k, v in message.axes.items() if k != axis
             }
-        msg_out = replace(
-            msg_in,
-            data=slice_along_axis(msg_in.data, _slice, axis_idx),
+        elif self._state.new_axis is not None:
+            replace_kwargs["axes"] = {
+                k: (v if k != axis else self._state.new_axis)
+                for k, v in message.axes.items()
+            }
+
+        return replace(
+            message,
+            data=slice_along_axis(message.data, self._state.slice_, axis_idx),
             **replace_kwargs,
         )
 
 
-class SlicerSettings(ez.Settings):
-    selection: str = ""
-    axis: str | None = None
+class Slicer(
+    BaseTransformerUnit[SlicerSettings, AxisArray, AxisArray, SlicerTransformer]
+):
+    SETTINGS = SlicerSettings
 
 
-class Slicer(GenAxisArray):
-    SETTINGS = SlicerSettings
+def slicer(selection: str = "", axis: str | None = None) -> SlicerTransformer:
+    """
+    Slice along a particular axis.
 
-    def construct_generator(self):
-        self.STATE.gen = slicer(
-            selection=self.SETTINGS.selection, axis=self.SETTINGS.axis
-        )
+    Args:
+        selection: See :obj:`ezmsg.sigproc.slicer.parse_slice` for details.
+        axis: The name of the axis to slice along. If None, the last axis is used.
+
+    Returns:
+        :obj:`SlicerTransformer`
+    """
+    return SlicerTransformer(SlicerSettings(selection=selection, axis=axis))