ezmsg-sigproc 1.2.2__py3-none-any.whl → 2.10.0__py3-none-any.whl

ezmsg/sigproc/scaler.py

@@ -0,0 +1,165 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.generator import consumer
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+# Imports for backwards compatibility with previous module location
+from .ewma import EWMA_Deprecated as EWMA_Deprecated
+from .ewma import EWMASettings, EWMATransformer, _alpha_from_tau
+from .ewma import _tau_from_alpha as _tau_from_alpha
+from .ewma import ewma_step as ewma_step
+
+
+@consumer
+def scaler(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/
+    This is faster than :obj:`scaler_np` for single-channel data.
+
+    Args:
+        time_constant: Decay constant `tau` in seconds.
+        axis: The name of the axis to accumulate statistics over.
+
+    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.
+    """
+    from river import preprocessing
+
+    msg_out = AxisArray(np.array([]), dims=[""])
+    _scaler = None
+    while True:
+        msg_in: AxisArray = yield msg_out
+        data = msg_in.data
+        if axis is None:
+            axis = msg_in.dims[0]
+            axis_idx = 0
+        else:
+            axis_idx = msg_in.get_axis_idx(axis)
+            if axis_idx != 0:
+                data = np.moveaxis(data, axis_idx, 0)
+
+        if _scaler is None:
+            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
+            _scaler = preprocessing.AdaptiveStandardScaler(fading_factor=alpha)
+
+        result = []
+        for sample in data:
+            x = {k: v for k, v in enumerate(sample.flatten().tolist())}
+            _scaler.learn_one(x)
+            y = _scaler.transform_one(x)
+            k = sorted(y.keys())
+            result.append(np.array([y[_] for _ in k]).reshape(sample.shape))
+
+        result = np.stack(result)
+        result = np.moveaxis(result, 0, axis_idx)
+        msg_out = replace(msg_in, data=result)
+
+
+class AdaptiveStandardScalerSettings(EWMASettings): ...
+
+
+@processor_state
+class AdaptiveStandardScalerState:
+    samps_ewma: EWMATransformer | None = None
+    vars_sq_ewma: EWMATransformer | None = None
+    alpha: float | None = None
+
+
+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,
+            accumulate=self.settings.accumulate,
+        )
+        self._state.vars_sq_ewma = EWMATransformer(
+            time_constant=self.settings.time_constant,
+            axis=self.settings.axis,
+            accumulate=self.settings.accumulate,
+        )
+
+    @property
+    def accumulate(self) -> bool:
+        """Whether to accumulate statistics from incoming samples."""
+        return self.settings.accumulate
+
+    @accumulate.setter
+    def accumulate(self, value: bool) -> None:
+        """
+        Set the accumulate mode and propagate to child EWMA transformers.
+
+        Args:
+            value: If True, update statistics with each sample.
+                   If False, only apply current statistics without updating.
+        """
+        if self._state.samps_ewma is not None:
+            self._state.samps_ewma.settings = replace(self._state.samps_ewma.settings, accumulate=value)
+        if self._state.vars_sq_ewma is not None:
+            self._state.vars_sq_ewma.settings = replace(self._state.vars_sq_ewma.settings, accumulate=value)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        # Update step (respects accumulate setting via child EWMAs)
+        mean_message = self._state.samps_ewma(message)
+        var_sq_message = self._state.vars_sq_ewma(replace(message, data=message.data**2))
+
+        # Get step
+        varis = var_sq_message.data - mean_message.data**2
+        with np.errstate(divide="ignore", invalid="ignore"):
+            result = (message.data - mean_message.data) / (varis**0.5)
+        result[np.isnan(result)] = 0.0
+        return replace(message, data=result)
+
+
+class AdaptiveStandardScaler(
+    BaseTransformerUnit[
+        AdaptiveStandardScalerSettings,
+        AxisArray,
+        AxisArray,
+        AdaptiveStandardScalerTransformer,
+    ]
+):
+    SETTINGS = AdaptiveStandardScalerSettings
+
+    @ez.subscriber(BaseTransformerUnit.INPUT_SETTINGS)
+    async def on_settings(self, msg: AdaptiveStandardScalerSettings) -> None:
+        """
+        Handle settings updates with smart reset behavior.
+
+        Only resets state if `axis` changes (structural change).
+        Changes to `time_constant` or `accumulate` are applied without
+        resetting accumulated statistics.
+        """
+        old_axis = self.SETTINGS.axis
+        self.apply_settings(msg)
+
+        if msg.axis != old_axis:
+            # Axis changed - need full reset
+            self.create_processor()
+        else:
+            # Update accumulate on processor (propagates to child EWMAs)
+            self.processor.accumulate = msg.accumulate
+            # Also update own settings reference
+            self.processor.settings = msg
+
+
+# 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

@@ -0,0 +1,70 @@
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from ezmsg.baseproc import (
+    BaseAsyncTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class SignalInjectorSettings(ez.Settings):
+    time_dim: str = "time"  # Input signal needs a time dimension with units in sec.
+    frequency: float | None = None  # Hz
+    amplitude: float = 1.0
+    mixing_seed: int | None = None
+
+
+@processor_state
+class SignalInjectorState:
+    cur_shape: tuple[int, ...] | None = None
+    cur_frequency: float | None = None
+    cur_amplitude: float | None = None
+    mixing: npt.NDArray | None = None
+
+
+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
+
+
+class SignalInjector(BaseTransformerUnit[SignalInjectorSettings, AxisArray, AxisArray, SignalInjectorTransformer]):
+    SETTINGS = SignalInjectorSettings
+    INPUT_FREQUENCY = ez.InputStream(float | None)
+    INPUT_AMPLITUDE = ez.InputStream(float)
+
+    @ez.subscriber(INPUT_FREQUENCY)
+    async def on_frequency(self, msg: float | None) -> None:
+        self.processor.state.cur_frequency = msg
+
+    @ez.subscriber(INPUT_AMPLITUDE)
+    async def on_amplitude(self, msg: float) -> None:
+        self.processor.state.cur_amplitude = msg

ezmsg/sigproc/slicer.py

@@ -0,0 +1,138 @@
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    AxisBase,
+    replace,
+    slice_along_axis,
+)
+
+"""
+Slicer:Select a subset of data along a particular axis.
+"""
+
+
+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.
+
+    - "" -> slice(None, None, None)  (take all)
+    - ":" -> slice(None, None, None)
+    - '"none"` (case-insensitive) -> slice(None, None, None)
+    - "{start}:{stop}" or {start}:{stop}:{step} -> slice(start, stop, step)
+    - "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.
+    """
+    if s.lower() in ["", ":", "none"]:
+        return (slice(None),)
+    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(_, axinfo=axinfo) for _ in s.split(",")]
+    return tuple([item for sublist in suplist for item in sublist])
+
+
+class SlicerSettings(ez.Settings):
+    selection: str = ""
+    """selection: See :obj:`ezmsg.sigproc.slicer.parse_slice` for details."""
+
+    axis: str | None = None
+    """The name of the axis to slice along. If None, the last axis is used."""
+
+
+@processor_state
+class SlicerState:
+    slice_: slice | int | npt.NDArray | None = None
+    new_axis: AxisBase | None = None
+    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:
+                out_data = in_data[self._state.slice_]
+            self._state.new_axis = replace(message.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 self._state.b_change_dims:
+            replace_kwargs["dims"] = [_ for dim_ix, _ in enumerate(message.dims) if dim_ix != axis_idx]
+            replace_kwargs["axes"] = {k: v for k, v in message.axes.items() if k != axis}
+        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 Slicer(BaseTransformerUnit[SlicerSettings, AxisArray, AxisArray, SlicerTransformer]):
+    SETTINGS = SlicerSettings
+
+
+def slicer(selection: str = "", axis: str | None = None) -> SlicerTransformer:
+    """
+    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.
+
+    Returns:
+        :obj:`SlicerTransformer`
+    """
+    return SlicerTransformer(SlicerSettings(selection=selection, axis=axis))

ezmsg/sigproc/spectral.py

@@ -1,132 +1,6 @@
-import enum
-
-from dataclasses import replace
-
-import numpy as np
-import ezmsg.core as ez
-
-from ezmsg.util.messages.axisarray import AxisArray
-
-from typing import Optional, AsyncGenerator
-
-
-class OptionsEnum(enum.Enum):
-    @classmethod
-    def options(cls):
-        return list(map(lambda c: c.value, cls))
-
-
-class WindowFunction(OptionsEnum):
-    NONE = "None (Rectangular)"
-    HAMMING = "Hamming"
-    HANNING = "Hanning"
-    BARTLETT = "Bartlett"
-    BLACKMAN = "Blackman"
-
-
-WINDOWS = {
-    WindowFunction.NONE: np.ones,
-    WindowFunction.HAMMING: np.hamming,
-    WindowFunction.HANNING: np.hanning,
-    WindowFunction.BARTLETT: np.bartlett,
-    WindowFunction.BLACKMAN: np.blackman,
-}
-
-
-class SpectralTransform(OptionsEnum):
-    RAW_COMPLEX = "Complex FFT Output"
-    REAL = "Real Component of FFT"
-    IMAG = "Imaginary Component of FFT"
-    REL_POWER = "Relative Power"
-    REL_DB = "Log Power (Relative dB)"
-
-
-class SpectralOutput(OptionsEnum):
-    FULL = "Full Spectrum"
-    POSITIVE = "Positive Frequencies"
-    NEGATIVE = "Negative Frequencies"
-
-
-class SpectrumSettings(ez.Settings):
-    axis: Optional[str] = None
-    # n: Optional[int] = None # n parameter for fft
-    out_axis: Optional[str] = "freq"  # If none; don't change dim name
-    window: WindowFunction = WindowFunction.HAMMING
-    transform: SpectralTransform = SpectralTransform.REL_DB
-    output: SpectralOutput = SpectralOutput.POSITIVE
-
-
-class SpectrumState(ez.State):
-    cur_settings: SpectrumSettings
-
-
-class Spectrum(ez.Unit):
-    SETTINGS: SpectrumSettings
-    STATE: SpectrumState
-
-    INPUT_SETTINGS = ez.InputStream(SpectrumSettings)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
-
-    def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
-
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: SpectrumSettings):
-        self.STATE.cur_settings = msg
-
-    @ez.subscriber(INPUT_SIGNAL)
-    @ez.publisher(OUTPUT_SIGNAL)
-    async def on_data(self, message: AxisArray) -> AsyncGenerator:
-        axis_name = self.STATE.cur_settings.axis
-        if axis_name is None:
-            axis_name = message.dims[0]
-        axis_idx = message.get_axis_idx(axis_name)
-        axis = message.get_axis(axis_name)
-
-        spectrum = np.moveaxis(message.data, axis_idx, -1)
-
-        n_time = message.data.shape[axis_idx]
-        window = WINDOWS[self.STATE.cur_settings.window](n_time)
-
-        spectrum = np.fft.fft(spectrum * window) / n_time
-        spectrum = np.fft.fftshift(spectrum, axes=-1)
-        freqs = np.fft.fftshift(np.fft.fftfreq(n_time, d=axis.gain), axes=-1)
-
-        if self.STATE.cur_settings.transform != SpectralTransform.RAW_COMPLEX:
-            if self.STATE.cur_settings.transform == SpectralTransform.REAL:
-                spectrum = spectrum.real
-            elif self.STATE.cur_settings.transform == SpectralTransform.IMAG:
-                spectrum = spectrum.imag
-            else:
-                scale = np.sum(window**2.0) * axis.gain
-                spectrum = (2.0 * (np.abs(spectrum) ** 2.0)) / scale
-
-                if self.STATE.cur_settings.transform == SpectralTransform.REL_DB:
-                    spectrum = 10 * np.log10(spectrum)
-
-        axis_offset = freqs[0]
-        if self.STATE.cur_settings.output == SpectralOutput.POSITIVE:
-            axis_offset = freqs[n_time // 2]
-            spectrum = spectrum[..., n_time // 2 :]
-        elif self.STATE.cur_settings.output == SpectralOutput.NEGATIVE:
-            spectrum = spectrum[..., : n_time // 2]
-
-        spectrum = np.moveaxis(spectrum, axis_idx, -1)
-
-        out_axis = self.SETTINGS.out_axis
-        if out_axis is None:
-            out_axis = axis_name
-
-        freq_axis = AxisArray.Axis(
-            unit="Hz", gain=1.0 / (axis.gain * n_time), offset=axis_offset
-        )
-        new_axes = {**message.axes, **{out_axis: freq_axis}}
-
-        new_dims = [d for d in message.dims]
-        if self.SETTINGS.out_axis is not None:
-            new_dims[axis_idx] = self.SETTINGS.out_axis
-
-        out_msg = replace(message, data=spectrum, dims=new_dims, axes=new_axes)
-
-        yield self.OUTPUT_SIGNAL, out_msg
+from .spectrum import OptionsEnum as OptionsEnum
+from .spectrum import SpectralOutput as SpectralOutput
+from .spectrum import SpectralTransform as SpectralTransform
+from .spectrum import Spectrum as Spectrum
+from .spectrum import SpectrumSettings as SpectrumSettings
+from .spectrum import WindowFunction as WindowFunction

ezmsg/sigproc/spectrogram.py

@@ -0,0 +1,90 @@
+from typing import Generator
+
+import ezmsg.core as ez
+from ezmsg.baseproc import (
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+    CompositeProcessor,
+)
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.modify import modify_axis
+
+from .spectrum import (
+    SpectralOutput,
+    SpectralTransform,
+    SpectrumTransformer,
+    WindowFunction,
+)
+from .window import Anchor, WindowTransformer
+
+
+class SpectrogramSettings(ez.Settings):
+    """
+    Settings for :obj:`SpectrogramTransformer`.
+    """
+
+    window_dur: float | None = None
+    """window duration in seconds."""
+
+    window_shift: float | None = None
+    """"window step in seconds. If None, window_shift == window_dur"""
+
+    window_anchor: str | Anchor = Anchor.BEGINNING
+    """See :obj"`WindowTransformer`"""
+
+    window: WindowFunction = WindowFunction.HAMMING
+    """The :obj:`WindowFunction` to apply to the data slice prior to calculating the spectrum."""
+
+    transform: SpectralTransform = SpectralTransform.REL_DB
+    """The :obj:`SpectralTransform` to apply to the spectral magnitude."""
+
+    output: SpectralOutput = SpectralOutput.POSITIVE
+    """The :obj:`SpectralOutput` format."""
+
+
+class SpectrogramTransformer(CompositeProcessor[SpectrogramSettings, AxisArray, AxisArray]):
+    @staticmethod
+    def _initialize_processors(
+        settings: SpectrogramSettings,
+    ) -> dict[str, BaseStatefulProcessor | Generator[AxisArray, AxisArray, None]]:
+        return {
+            "windowing": WindowTransformer(
+                axis="time",
+                newaxis="win",
+                window_dur=settings.window_dur,
+                window_shift=settings.window_shift,
+                zero_pad_until="shift" if settings.window_shift is not None else "input",
+                anchor=settings.window_anchor,
+            ),
+            "spectrum": SpectrumTransformer(
+                axis="time",
+                window=settings.window,
+                transform=settings.transform,
+                output=settings.output,
+            ),
+            "modify_axis": modify_axis(name_map={"win": "time"}),
+        }
+
+
+class Spectrogram(BaseTransformerUnit[SpectrogramSettings, AxisArray, AxisArray, SpectrogramTransformer]):
+    SETTINGS = SpectrogramSettings
+
+
+def spectrogram(
+    window_dur: float | None = None,
+    window_shift: float | None = None,
+    window_anchor: str | Anchor = Anchor.BEGINNING,
+    window: WindowFunction = WindowFunction.HAMMING,
+    transform: SpectralTransform = SpectralTransform.REL_DB,
+    output: SpectralOutput = SpectralOutput.POSITIVE,
+) -> SpectrogramTransformer:
+    return SpectrogramTransformer(
+        SpectrogramSettings(
+            window_dur=window_dur,
+            window_shift=window_shift,
+            window_anchor=window_anchor,
+            window=window,
+            transform=transform,
+            output=output,
+        )
+    )