ezmsg-sigproc 1.4.2__py3-none-any.whl → 1.6.0__py3-none-any.whl

ezmsg/sigproc/scaler.py

@@ -1,10 +1,12 @@
-from dataclasses import replace
+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
@@ -28,9 +30,139 @@ def _alpha_from_tau(tau: float, dt: float) -> float:
     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
+
+
 @consumer
 def scaler(
-    time_constant: float = 1.0, axis: typing.Optional[str] = None
+    time_constant: float = 1.0, axis: str | None = None
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Apply the adaptive standard scaler from https://riverml.xyz/latest/api/preprocessing/AdaptiveStandardScaler/
@@ -78,7 +210,7 @@ def scaler(
 
 @consumer
 def scaler_np(
-    time_constant: float = 1.0, axis: typing.Optional[str] = None
+    time_constant: float = 1.0, axis: str | None = None
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Create a generator function that applies an adaptive standard scaler.
@@ -87,6 +219,7 @@ def scaler_np(
     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)`
@@ -95,10 +228,8 @@ def scaler_np(
     msg_out = AxisArray(np.array([]), dims=[""])
 
     # State variables
-    alpha: float = 0.0
-    means: typing.Optional[npt.NDArray] = None
-    vars_means: typing.Optional[npt.NDArray] = None
-    vars_sq_means: typing.Optional[npt.NDArray] = None
+    samps_ewma: EWMA | None = None
+    vars_sq_ewma: EWMA | None = None
 
     # Reset if input changes
     check_input = {
@@ -107,45 +238,32 @@ def scaler_np(
         "key": None,  # Key change implies buffered means/vars are invalid.
     }
 
-    def _ew_update(arr, prev, _alpha):
-        if np.all(prev == 0):
-            return arr
-        # return _alpha * arr + (1 - _alpha) * prev
-        # Micro-optimization: sub, mult, add (below) is faster than sub, mult, mult, add (above)
-        return prev + _alpha * (arr - prev)
-
     while True:
         msg_in: AxisArray = yield msg_out
 
         axis = axis or msg_in.dims[0]
         axis_idx = msg_in.get_axis_idx(axis)
 
-        if msg_in.axes[axis].gain != check_input["gain"]:
-            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
-            check_input["gain"] = msg_in.axes[axis].gain
-
         data: npt.NDArray = np.moveaxis(msg_in.data, axis_idx, 0)
         b_reset = data.shape[1:] != check_input["shape"]
-        b_reset |= msg_in.key != check_input["key"]
+        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
-            vars_sq_means = np.zeros_like(data[0], dtype=float)
-            vars_means = np.zeros_like(data[0], dtype=float)
-            means = np.zeros_like(data[0], dtype=float)
-
-        result = np.zeros_like(data)
-        for sample_ix in range(data.shape[0]):
-            sample = data[sample_ix]
-            # Update step
-            vars_means = _ew_update(sample, vars_means, alpha)
-            vars_sq_means = _ew_update(sample**2, vars_sq_means, alpha)
-            means = _ew_update(sample, means, alpha)
-            # Get step
-            varis = vars_sq_means - vars_means**2
-            y = (sample - means) / (varis**0.5)
-            result[sample_ix] = y
+            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
+            samps_ewma = EWMA(alpha=alpha)
+            vars_sq_ewma = EWMA(alpha=alpha)
+
+        # Update step
+        means = samps_ewma.compute(data)
+        vars_sq_means = vars_sq_ewma.compute(data**2)
 
+        # Get step
+        varis = vars_sq_means - means**2
+        with np.errstate(divide="ignore", invalid="ignore"):
+            result = (data - means) / (varis**0.5)
         result[np.isnan(result)] = 0.0
         result = np.moveaxis(result, 0, axis_idx)
         msg_out = replace(msg_in, data=result)
@@ -158,7 +276,7 @@ class AdaptiveStandardScalerSettings(ez.Settings):
     """
 
     time_constant: float = 1.0
-    axis: typing.Optional[str] = None
+    axis: str | None = None
 
 
 class AdaptiveStandardScaler(GenAxisArray):

ezmsg/sigproc/signalinjector.py

@@ -1,22 +1,22 @@
-from dataclasses import replace
 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
 
 
 class SignalInjectorSettings(ez.Settings):
     time_dim: str = "time"  # Input signal needs a time dimension with units in sec.
-    frequency: typing.Optional[float] = None  # Hz
+    frequency: float | None = None  # Hz
     amplitude: float = 1.0
-    mixing_seed: typing.Optional[int] = None
+    mixing_seed: int | None = None
 
 
 class SignalInjectorState(ez.State):
-    cur_shape: typing.Optional[typing.Tuple[int, ...]] = None
-    cur_frequency: typing.Optional[float] = None
+    cur_shape: tuple[int, ...] | None = None
+    cur_frequency: float | None = None
     cur_amplitude: float
     mixing: npt.NDArray
 
@@ -30,7 +30,7 @@ class SignalInjector(ez.Unit):
     SETTINGS = SignalInjectorSettings
     STATE = SignalInjectorState
 
-    INPUT_FREQUENCY = ez.InputStream(typing.Optional[float])
+    INPUT_FREQUENCY = ez.InputStream(float | None)
     INPUT_AMPLITUDE = ez.InputStream(float)
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
@@ -41,7 +41,7 @@ class SignalInjector(ez.Unit):
         self.STATE.mixing = np.array([])
 
     @ez.subscriber(INPUT_FREQUENCY)
-    async def on_frequency(self, msg: typing.Optional[float]) -> None:
+    async def on_frequency(self, msg: float | None) -> None:
         self.STATE.cur_frequency = msg
 
     @ez.subscriber(INPUT_AMPLITUDE)

ezmsg/sigproc/slicer.py

@@ -1,10 +1,14 @@
-from dataclasses import replace
 import typing
 
 import numpy as np
 import numpy.typing as npt
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    slice_along_axis,
+    AxisBase,
+    replace,
+)
 from ezmsg.util.generator import consumer
 
 from .base import GenAxisArray
@@ -15,7 +19,10 @@ Slicer:Select a subset of data along a particular axis.
 """
 
 
-def parse_slice(s: str) -> typing.Tuple[typing.Union[slice, int], ...]:
+def parse_slice(
+    s: str,
+    axinfo: AxisArray.CoordinateAxis | None = None,
+) -> tuple[slice | int, ...]:
     """
     Parses a string representation of a slice and returns a tuple of slice objects.
 
@@ -26,9 +33,13 @@ def parse_slice(s: str) -> typing.Tuple[typing.Union[slice, int], ...]:
     - "5" (or any integer) -> (5,). Take only that item.
         applying this to a ndarray or AxisArray will drop the dimension.
     - A comma-separated list of the above -> a tuple of slices | ints
+    - A comma-separated list of values and axinfo is provided and is a CoordinateAxis -> a tuple of ints
 
     Args:
         s: The string representation of the slice.
+        axinfo: (Optional) If provided, and of type CoordinateAxis,
+          and `s` is a comma-separated list of values, then the values
+          in s will be checked against the values in axinfo.data.
 
     Returns:
         A tuple of slice objects and/or ints.
@@ -38,15 +49,21 @@ def parse_slice(s: str) -> typing.Tuple[typing.Union[slice, int], ...]:
     if "," not in s:
         parts = [part.strip() for part in s.split(":")]
         if len(parts) == 1:
+            if (
+                axinfo is not None
+                and hasattr(axinfo, "data")
+                and parts[0] in axinfo.data
+            ):
+                return tuple(np.where(axinfo.data == parts[0])[0])
             return (int(parts[0]),)
         return (slice(*(int(part.strip()) if part else None for part in parts)),)
-    suplist = [parse_slice(_) for _ in s.split(",")]
+    suplist = [parse_slice(_, axinfo=axinfo) for _ in s.split(",")]
     return tuple([item for sublist in suplist for item in sublist])
 
 
 @consumer
 def slicer(
-    selection: str = "", axis: typing.Optional[str] = None
+    selection: str = "", axis: str | None = None
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Slice along a particular axis.
@@ -63,8 +80,8 @@ def slicer(
     msg_out = AxisArray(np.array([]), dims=[""])
 
     # State variables
-    _slice: typing.Optional[typing.Union[slice, npt.NDArray]] = None
-    new_axis: typing.Optional[AxisArray.Axis] = None
+    _slice: slice | 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
@@ -92,7 +109,7 @@ def slicer(
             b_change_dims = False
 
             # Calculate the slice
-            _slices = parse_slice(selection)
+            _slices = parse_slice(selection, msg_in.axes.get(axis, None))
             if len(_slices) == 1:
                 _slice = _slices[0]
                 # Do we drop the sliced dimension?
@@ -107,12 +124,15 @@ def slicer(
             # Create the output axis.
             if (
                 axis in msg_in.axes
-                and hasattr(msg_in.axes[axis], "labels")
-                and len(msg_in.axes[axis].labels) > 0
+                and hasattr(msg_in.axes[axis], "data")
+                and len(msg_in.axes[axis].data) > 0
             ):
-                in_labels = np.array(msg_in.axes[axis].labels)
-                new_labels = in_labels[_slice].tolist()
-                new_axis = replace(msg_in.axes[axis], labels=new_labels)
+                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)
 
         replace_kwargs = {}
         if b_change_dims:
@@ -134,7 +154,7 @@ def slicer(
 
 class SlicerSettings(ez.Settings):
     selection: str = ""
-    axis: typing.Optional[str] = None
+    axis: str | None = None
 
 
 class Slicer(GenAxisArray):

ezmsg/sigproc/spectrogram.py

@@ -12,12 +12,12 @@ from .base import GenAxisArray
 
 @consumer
 def spectrogram(
-    window_dur: typing.Optional[float] = None,
-    window_shift: typing.Optional[float] = None,
+    window_dur: float | None = None,
+    window_shift: float | None = None,
     window: WindowFunction = WindowFunction.HANNING,
     transform: SpectralTransform = SpectralTransform.REL_DB,
     output: SpectralOutput = SpectralOutput.POSITIVE,
-) -> typing.Generator[typing.Optional[AxisArray], AxisArray, None]:
+) -> typing.Generator[AxisArray | None, AxisArray, None]:
     """
     Calculate a spectrogram on streaming data.
 
@@ -50,7 +50,7 @@ def spectrogram(
     )
 
     # State variables
-    msg_out: typing.Optional[AxisArray] = None
+    msg_out: AxisArray | None = None
 
     while True:
         msg_in: AxisArray = yield msg_out
@@ -63,8 +63,8 @@ class SpectrogramSettings(ez.Settings):
     See :obj:`spectrogram` for a description of the parameters.
     """
 
-    window_dur: typing.Optional[float] = None  # window duration in seconds
-    window_shift: typing.Optional[float] = None
+    window_dur: float | None = None  # window duration in seconds
+    window_shift: float | None = None
     """"window step in seconds. If None, window_shift == window_dur"""
 
     # See SpectrumSettings for details of following settings:

ezmsg/sigproc/spectrum.py

@@ -1,11 +1,14 @@
-from dataclasses import replace
 import enum
 from functools import partial
 import typing
 
 import numpy as np
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    slice_along_axis,
+    replace,
+)
 from ezmsg.util.generator import consumer
 
 from .base import GenAxisArray
@@ -65,20 +68,21 @@ class SpectralOutput(OptionsEnum):
 
 @consumer
 def spectrum(
-    axis: typing.Optional[str] = None,
-    out_axis: typing.Optional[str] = "freq",
+    axis: str | None = None,
+    out_axis: str | None = "freq",
     window: WindowFunction = WindowFunction.HANNING,
     transform: SpectralTransform = SpectralTransform.REL_DB,
     output: SpectralOutput = SpectralOutput.POSITIVE,
-    norm: typing.Optional[str] = "forward",
+    norm: str | None = "forward",
     do_fftshift: bool = True,
-    nfft: typing.Optional[int] = None,
+    nfft: int | None = None,
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Calculate a spectrum on a data slice.
 
     Args:
         axis: The name of the axis on which to calculate the spectrum.
+            Note: The axis must have an .axes entry of type LinearAxis, not CoordinateAxis.
         out_axis: The name of the new axis. Defaults to "freq".
         window: The :obj:`WindowFunction` to apply to the data slice prior to calculating the spectrum.
         transform: The :obj:`SpectralTransform` to apply to the spectral magnitude.
@@ -101,10 +105,10 @@ def spectrum(
     apply_window = window != WindowFunction.NONE
     do_fftshift &= output == SpectralOutput.FULL
     f_sl = slice(None)
-    freq_axis: typing.Optional[AxisArray.Axis] = None
-    fftfun: typing.Optional[typing.Callable] = None
-    f_transform: typing.Optional[typing.Callable] = None
-    new_dims: typing.Optional[typing.List[str]] = None
+    freq_axis: AxisArray.LinearAxis | None = None
+    fftfun: typing.Callable | None = None
+    f_transform: typing.Callable | None = None
+    new_dims: list[str] | None = None
 
     # Reset if input changes substantially
     check_input = {
@@ -174,7 +178,7 @@ def spectrum(
                     freqs = np.fft.fftshift(freqs, axes=-1)
                 freqs = freqs[f_sl]
             freqs = freqs.tolist()  # To please type checking
-            freq_axis = AxisArray.Axis(
+            freq_axis = AxisArray.LinearAxis(
                 unit="Hz", gain=freqs[1] - freqs[0], offset=freqs[0]
             )
             if out_axis is None:
@@ -234,9 +238,9 @@ class SpectrumSettings(ez.Settings):
     See :obj:`spectrum` for a description of the parameters.
     """
 
-    axis: typing.Optional[str] = None
-    # n: typing.Optional[int] = None # n parameter for fft
-    out_axis: typing.Optional[str] = "freq"  # If none; don't change dim name
+    axis: str | None = None
+    # n: int | None = None # n parameter for fft
+    out_axis: str | None = "freq"  # If none; don't change dim name
     window: WindowFunction = WindowFunction.HAMMING
     transform: SpectralTransform = SpectralTransform.REL_DB
     output: SpectralOutput = SpectralOutput.POSITIVE