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

ezmsg/sigproc/downsample.py

@@ -1,82 +1,86 @@
-import typing
-
 import numpy as np
 from ezmsg.util.messages.axisarray import (
     AxisArray,
     slice_along_axis,
     replace,
 )
-from ezmsg.util.generator import consumer
 import ezmsg.core as ez
 
-from .base import GenAxisArray
+from .base import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 
 
-@consumer
-def downsample(
-    axis: str | None = None, target_rate: float | None = None
-) -> typing.Generator[AxisArray, AxisArray, None]:
+class DownsampleSettings(ez.Settings):
     """
-    Construct a generator that yields a downsampled version of the data .send() to it.
-    Downsampled data simply comprise every `factor`th sample.
-    This should only be used following appropriate lowpass filtering.
-    If your pipeline does not already have lowpass filtering then consider
-    using the :obj:`Decimate` collection instead.
+    Settings for :obj:`Downsample` node.
+    """
+
+    axis: str = "time"
+    """The name of the axis along which to downsample."""
 
-    Args:
-        axis: The name of the axis along which to downsample.
-            Note: The axis must exist in the message .axes and be of type AxisArray.LinearAxis.
-        target_rate: Desired rate after downsampling. The actual rate will be the nearest integer factor of the
-            input rate that is the same or higher than the target rate.
+    target_rate: float | None = None
+    """Desired rate after downsampling. The actual rate will be the nearest integer factor of the
+            input rate that is the same or higher than the target rate."""
 
-    Returns:
-        A primed generator object ready to receive an :obj:`AxisArray` via `.send(axis_array)`
-        and yields an :obj:`AxisArray` with its data downsampled.
-        Note that if a send chunk does not have sufficient samples to reach the
-        next downsample interval then an :obj:`AxisArray` with size-zero data is yielded.
+    factor: int | None = None
+    """Explicitly specify downsample factor.  If specified, target_rate is ignored."""
 
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
 
-    # state variables
-    factor: int = 0  # The integer downsampling factor. It will be determined based on the target rate.
-    s_idx: int = 0  # Index of the next msg's first sample into the virtual rotating ds_factor counter.
+@processor_state
+class DownsampleState:
+    q: int = 0
+    """The integer downsampling factor. It will be determined based on the target rate."""
 
-    check_input = {"gain": None, "key": None}
+    s_idx: int = 0
+    """Index of the next msg's first sample into the virtual rotating ds_factor counter."""
+
+
+class DownsampleTransformer(
+    BaseStatefulTransformer[DownsampleSettings, AxisArray, AxisArray, DownsampleState]
+):
+    """
+    Downsampled data simply comprise every `factor`th sample.
+    This should only be used following appropriate lowpass filtering.
+    If your pipeline does not already have lowpass filtering then consider
+    using the :obj:`Decimate` collection instead.
+    """
 
-    while True:
-        msg_in: AxisArray = yield msg_out
+    def _hash_message(self, message: AxisArray) -> int:
+        return hash((message.axes[self.settings.axis].gain, message.key))
 
-        if axis is None:
-            axis = msg_in.dims[0]
-        axis_info = msg_in.get_axis(axis)
-        axis_idx = msg_in.get_axis_idx(axis)
+    def _reset_state(self, message: AxisArray) -> None:
+        axis_info = message.get_axis(self.settings.axis)
 
-        b_reset = (
-            msg_in.axes[axis].gain != check_input["gain"]
-            or msg_in.key != check_input["key"]
+        if self.settings.factor is not None:
+            q = self.settings.factor
+        elif self.settings.target_rate is None:
+            q = 1
+        else:
+            q = int(1 / (axis_info.gain * self.settings.target_rate))
+        if q < 1:
+            ez.logger.warning(
+                f"Target rate {self.settings.target_rate} cannot be achieved with input rate of {1 / axis_info.gain}."
+                "Setting factor to 1."
+            )
+            q = 1
+        self._state.q = q
+        self._state.s_idx = 0
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis
+        axis_info = message.get_axis(axis)
+        axis_idx = message.get_axis_idx(axis)
+
+        n_samples = message.data.shape[axis_idx]
+        samples = (
+            np.arange(self.state.s_idx, self.state.s_idx + n_samples) % self._state.q
         )
-        if b_reset:
-            check_input["gain"] = axis_info.gain
-            check_input["key"] = msg_in.key
-            # Reset state variables
-            s_idx = 0
-            if target_rate is None:
-                factor = 1
-            else:
-                factor = int(1 / (axis_info.gain * target_rate))
-            if factor < 1:
-                ez.logger.warning(
-                    f"Target rate {target_rate} cannot be achieved with input rate of {1/axis_info.gain}."
-                    "Setting factor to 1."
-                )
-                factor = 1
-
-        n_samples = msg_in.data.shape[axis_idx]
-        samples = np.arange(s_idx, s_idx + n_samples) % factor
         if n_samples > 0:
             # Update state for next iteration.
-            s_idx = samples[-1] + 1
+            self._state.s_idx = samples[-1] + 1
 
         pub_samples = np.where(samples == 0)[0]
         if len(pub_samples) > 0:
@@ -86,35 +90,31 @@ def downsample(
             n_step = 0
             data_slice = slice(None, 0, None)
         msg_out = replace(
-            msg_in,
-            data=slice_along_axis(msg_in.data, data_slice, axis=axis_idx),
+            message,
+            data=slice_along_axis(message.data, data_slice, axis=axis_idx),
             axes={
-                **msg_in.axes,
+                **message.axes,
                 axis: replace(
                     axis_info,
-                    gain=axis_info.gain * factor,
+                    gain=axis_info.gain * self._state.q,
                     offset=axis_info.offset + axis_info.gain * n_step,
                 ),
             },
         )
+        return msg_out
 
 
-class DownsampleSettings(ez.Settings):
-    """
-    Settings for :obj:`Downsample` node.
-    See :obj:`downsample` documentation for a description of the parameters.
-    """
-
-    axis: str | None = None
-    target_rate: float | None = None
-
-
-class Downsample(GenAxisArray):
-    """:obj:`Unit` for :obj:`bandpower`."""
-
+class Downsample(
+    BaseTransformerUnit[DownsampleSettings, AxisArray, AxisArray, DownsampleTransformer]
+):
     SETTINGS = DownsampleSettings
 
-    def construct_generator(self):
-        self.STATE.gen = downsample(
-            axis=self.SETTINGS.axis, target_rate=self.SETTINGS.target_rate
-        )
+
+def downsample(
+    axis: str = "time",
+    target_rate: float | None = None,
+    factor: int | None = None,
+) -> DownsampleTransformer:
+    return DownsampleTransformer(
+        DownsampleSettings(axis=axis, target_rate=target_rate, factor=factor)
+    )

ezmsg/sigproc/ewma.py

@@ -0,0 +1,197 @@
+from dataclasses import field
+import functools
+
+import numpy as np
+import numpy.typing as npt
+import scipy.signal as sps
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.util import replace
+
+from .base import BaseStatefulTransformer, processor_state, BaseTransformerUnit
+
+
+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_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
+
+
+class EWMASettings(ez.Settings):
+    time_constant: float = 1.0
+    axis: str | None = None
+
+
+@processor_state
+class EWMAState:
+    alpha: float = field(default_factory=lambda: _alpha_from_tau(1.0, 1000.0))
+    zi: npt.NDArray | None = None
+
+
+class EWMATransformer(
+    BaseStatefulTransformer[EWMASettings, AxisArray, AxisArray, EWMAState]
+):
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        sample_shape = (
+            message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        )
+        return hash((sample_shape, message.axes[axis].gain, message.key))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis = self.settings.axis or message.dims[0]
+        self._state.alpha = _alpha_from_tau(
+            self.settings.time_constant, message.axes[axis].gain
+        )
+        sub_dat = slice_along_axis(
+            message.data, slice(None, 1, None), axis=message.get_axis_idx(axis)
+        )
+        self._state.zi = (1 - self._state.alpha) * sub_dat
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        if np.prod(message.data.shape) == 0:
+            return message
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        expected, self._state.zi = sps.lfilter(
+            [self._state.alpha],
+            [1.0, self._state.alpha - 1.0],
+            message.data,
+            axis=axis_idx,
+            zi=self._state.zi,
+        )
+        return replace(message, data=expected)
+
+
+class EWMAUnit(
+    BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, EWMATransformer]
+):
+    SETTINGS = EWMASettings

ezmsg/sigproc/extract_axis.py

@@ -0,0 +1,41 @@
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.sigproc.base import BaseTransformer, BaseTransformerUnit
+
+
+class ExtractAxisSettings(ez.Settings):
+    axis: str = "freq"
+    reference: str = "time"
+
+
+class ExtractAxisData(BaseTransformer[ExtractAxisSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        targ_ax = message.axes[self.settings.axis]
+        if hasattr(targ_ax, "data"):
+            # Extracted axis is of type CoordinateAxis
+            return replace(
+                message,
+                data=targ_ax.data,
+                dims=targ_ax.dims,
+                axes={k: v for k, v in message.axes.items() if k in targ_ax.dims},
+            )
+            # Note: So far we don't have any transformers where the coordinate axis has its own axes,
+            # but if that happens in the future, we'd need to consider how to handle that.
+
+        else:
+            # Extracted axis is of type LinearAxis
+            #  LinearAxis can only yield a 1d array data which simplifies dims and axes.
+            n = message.data.shape[message.get_axis_idx(self.settings.reference)]
+            return replace(
+                message,
+                data=targ_ax.value(np.arange(n)),
+                dims=[self.settings.reference],
+                axes={self.settings.reference: message.axes[self.settings.reference]},
+            )
+
+
+class ExtractAxisDataUnit(
+    BaseTransformerUnit[ExtractAxisSettings, AxisArray, AxisArray, ExtractAxisData]
+):
+    SETTINGS = ExtractAxisSettings