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

ezmsg/sigproc/scaler.py

@@ -0,0 +1,173 @@
+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
+from ezmsg.util.generator import consumer
+
+from .base import GenAxisArray
+
+
+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)
+
+
+@consumer
+def scaler(
+    time_constant: float = 1.0, axis: typing.Optional[str] = None
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Create a generator function that applies 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 `.send(axis_array)` and yields a
+        standardized, or "Z-scored" version of the input.
+    """
+    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)
+
+
+@consumer
+def scaler_np(
+    time_constant: float = 1.0, axis: typing.Optional[str] = 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.
+
+    Returns:
+        A primed generator object that expects `.send(axis_array)` and yields a
+        standardized, or "Z-scored" version of the input.
+    """
+    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
+
+    # Reset if input changes
+    check_input = {
+        "gain": None,  # Resets alpha
+        "shape": None,
+        "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"]
+        if b_reset:
+            check_input["shape"] = data.shape[1:]
+            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
+
+        result[np.isnan(result)] = 0.0
+        result = np.moveaxis(result, 0, axis_idx)
+        msg_out = replace(msg_in, 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: typing.Optional[str] = None
+
+
+class AdaptiveStandardScaler(GenAxisArray):
+    """Unit for :obj:`scaler_np`"""
+
+    SETTINGS = AdaptiveStandardScalerSettings
+
+    def construct_generator(self):
+        self.STATE.gen = scaler_np(
+            time_constant=self.SETTINGS.time_constant, axis=self.SETTINGS.axis
+        )

ezmsg/sigproc/signalinjector.py

@@ -0,0 +1,64 @@
+from dataclasses import replace
+import typing
+
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+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
+    amplitude: float = 1.0
+    mixing_seed: typing.Optional[int] = None
+
+
+class SignalInjectorState(ez.State):
+    cur_shape: typing.Optional[typing.Tuple[int, ...]] = None
+    cur_frequency: typing.Optional[float] = None
+    cur_amplitude: float
+    mixing: npt.NDArray
+
+
+class SignalInjector(ez.Unit):
+    SETTINGS = SignalInjectorSettings
+    STATE = SignalInjectorState
+
+    INPUT_FREQUENCY = ez.InputStream(typing.Optional[float])
+    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: typing.Optional[float]) -> None:
+        self.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)
+    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

ezmsg/sigproc/slicer.py

@@ -0,0 +1,133 @@
+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.generator import consumer
+
+from .base import GenAxisArray
+
+
+"""
+Slicer:Select a subset of data along a particular axis.
+"""
+
+
+def parse_slice(s: str) -> typing.Tuple[typing.Union[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
+
+    Args:
+        s: The string representation of the slice.
+
+    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:
+            return (int(parts[0]),)
+        return (slice(*(int(part.strip()) if part else None for part in parts)),)
+    suplist = [parse_slice(_) 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
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    # State variables
+    _slice: typing.Optional[typing.Union[slice, npt.NDArray]] = None
+    new_axis: typing.Optional[AxisArray.Axis] = 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)
+            if len(_slices) == 1:
+                _slice = _slices[0]
+                # Do we drop the sliced dimension?
+                b_change_dims = isinstance(_slice, int)
+            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
+
+            # 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
+            ):
+                new_labels = msg_in.axes[axis].labels[_slice]
+                new_axis = replace(msg_in.axes[axis], labels=new_labels)
+
+        replace_kwargs = {}
+        if b_change_dims:
+            # Dropping the target axis
+            replace_kwargs["dims"] = [
+                _ for dim_ix, _ in enumerate(msg_in.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()
+            }
+        msg_out = replace(
+            msg_in,
+            data=slice_along_axis(msg_in.data, _slice, axis_idx),
+            **replace_kwargs,
+        )
+
+
+class SlicerSettings(ez.Settings):
+    selection: str = ""
+    axis: typing.Optional[str] = None
+
+
+class Slicer(GenAxisArray):
+    SETTINGS = SlicerSettings
+
+    def construct_generator(self):
+        self.STATE.gen = slicer(
+            selection=self.SETTINGS.selection, axis=self.SETTINGS.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 WindowFunction as WindowFunction
+from .spectrum import SpectralTransform as SpectralTransform
+from .spectrum import SpectralOutput as SpectralOutput
+from .spectrum import SpectrumSettings as SpectrumSettings
+from .spectrum import Spectrum as Spectrum

ezmsg/sigproc/spectrogram.py

@@ -0,0 +1,86 @@
+import typing
+
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer, compose
+from ezmsg.util.messages.modify import modify_axis
+
+from .window import windowing
+from .spectrum import spectrum, WindowFunction, SpectralTransform, SpectralOutput
+from .base import GenAxisArray
+
+
+@consumer
+def spectrogram(
+    window_dur: typing.Optional[float] = None,
+    window_shift: typing.Optional[float] = None,
+    window: WindowFunction = WindowFunction.HANNING,
+    transform: SpectralTransform = SpectralTransform.REL_DB,
+    output: SpectralOutput = SpectralOutput.POSITIVE,
+) -> typing.Generator[typing.Optional[AxisArray], AxisArray, None]:
+    """
+    Calculate a spectrogram on streaming data.
+
+    Chains :obj:`ezmsg.sigproc.window.windowing` to apply a moving window on the data,
+    :obj:`ezmsg.sigproc.spectrum.spectrum` to calculate spectra for each window,
+    and finally :obj:`ezmsg.util.messages.modify.modify_axis` to convert the win axis back to time axis.
+
+    Args:
+        window_dur: See :obj:`ezmsg.sigproc.window.windowing`
+        window_shift: See :obj:`ezmsg.sigproc.window.windowing`
+        window: See :obj:`ezmsg.sigproc.spectrum.spectrum`
+        transform: See :obj:`ezmsg.sigproc.spectrum.spectrum`
+        output: See :obj:`ezmsg.sigproc.spectrum.spectrum`
+
+    Returns:
+        A primed generator object that expects `.send(axis_array)` of continuous data
+        and yields an AxisArray of time-frequency power values.
+    """
+
+    pipeline = compose(
+        windowing(
+            axis="time", newaxis="win", window_dur=window_dur, window_shift=window_shift
+        ),
+        spectrum(axis="time", window=window, transform=transform, output=output),
+        modify_axis(name_map={"win": "time"}),
+    )
+
+    # State variables
+    msg_out: typing.Optional[AxisArray] = None
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+        msg_out = pipeline(msg_in)
+
+
+class SpectrogramSettings(ez.Settings):
+    """
+    Settings for :obj:`Spectrogram`.
+    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 step in seconds. If None, window_shift == window_dur"""
+
+    # See SpectrumSettings for details of following settings:
+    window: WindowFunction = WindowFunction.HAMMING
+    transform: SpectralTransform = SpectralTransform.REL_DB
+    output: SpectralOutput = SpectralOutput.POSITIVE
+
+
+class Spectrogram(GenAxisArray):
+    """
+    Unit for :obj:`spectrogram`.
+    """
+
+    SETTINGS = SpectrogramSettings
+
+    def construct_generator(self):
+        self.STATE.gen = spectrogram(
+            window_dur=self.SETTINGS.window_dur,
+            window_shift=self.SETTINGS.window_shift,
+            window=self.SETTINGS.window,
+            transform=self.SETTINGS.transform,
+            output=self.SETTINGS.output,
+        )