ezmsg-sigproc 2.4.1__py3-none-any.whl → 2.6.0__py3-none-any.whl

ezmsg/sigproc/rollingscaler.py

@@ -0,0 +1,232 @@
+from collections import deque
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from ezmsg.sigproc.base import (
+    BaseAdaptiveTransformer,
+    BaseAdaptiveTransformerUnit,
+    processor_state,
+)
+from ezmsg.sigproc.sampler import SampleMessage
+
+
+class RollingScalerSettings(ez.Settings):
+    axis: str = "time"
+    """
+    Axis along which samples are arranged.
+    """
+
+    k_samples: int | None = 20
+    """
+    Rolling window size in number of samples.
+    """
+
+    window_size: float | None = None
+    """
+    Rolling window size in seconds.
+    If set, overrides `k_samples`.
+    `update_with_signal` likely should be True if using this option.
+    """
+
+    update_with_signal: bool = False
+    """
+    If True, update rolling statistics using the incoming process stream.
+    """
+
+    min_samples: int = 1
+    """
+    Minimum number of samples required to compute statistics.
+    Used when `window_size` is not set.
+    """
+
+    min_seconds: float = 1.0
+    """
+    Minimum duration in seconds required to compute statistics.
+    Used when `window_size` is set.
+    """
+
+    artifact_z_thresh: float | None = None
+    """
+    Threshold for z-score based artifact detection.
+    If set, samples with any channel exceeding this z-score will be excluded
+    from updating the rolling statistics.
+    """
+
+    clip: float | None = 10.0
+    """
+    If set, clip the output values to the range [-clip, clip].
+    """
+
+
+@processor_state
+class RollingScalerState:
+    mean: npt.NDArray | None = None
+    N: int = 0
+    M2: npt.NDArray | None = None
+    samples: deque | None = None
+    k_samples: int | None = None
+    min_samples: int | None = None
+
+
+class RollingScalerProcessor(BaseAdaptiveTransformer[RollingScalerSettings, AxisArray, AxisArray, RollingScalerState]):
+    """
+    Processor for rolling z-score normalization of input `AxisArray` messages.
+
+    The processor maintains rolling statistics (mean and variance) over the last `k_samples`
+    samples received via the `partial_fit()` method. When processing an `AxisArray` message,
+    it normalizes the data using the current rolling statistics.
+
+    The input `AxisArray` messages are expected to have shape `(time, ch)`, where `ch` is the
+    channel axis. The processor computes the z-score for each channel independently.
+
+    Note: You should consider instead using the AdaptiveStandardScalerTransformer which
+    is computationally more efficient and uses less memory. This RollingScalerProcessor
+    is primarily provided to reproduce processing in the literature.
+
+    Settings:
+    ---------
+    k_samples: int
+        Number of previous samples to use for rolling statistics.
+
+    Example:
+    -----------------------------
+    ```python
+    processor = RollingScalerProcessor(
+        settings=RollingScalerSettings(
+            k_samples=20  # Number of previous samples to use for rolling statistics
+        )
+    )
+    ```
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        gain = message.axes[axis].gain if hasattr(message.axes[axis], "gain") else 1
+        axis_idx = message.get_axis_idx(axis)
+        samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        return hash((message.key, samp_shape, gain))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        ch = message.data.shape[-1]
+        self._state.mean = np.zeros(ch)
+        self._state.N = 0
+        self._state.M2 = np.zeros(ch)
+        self._state.k_samples = (
+            int(np.ceil(self.settings.window_size / message.axes[self.settings.axis].gain))
+            if self.settings.window_size is not None
+            else self.settings.k_samples
+        )
+        if self._state.k_samples is not None and self._state.k_samples < 1:
+            ez.logger.warning("window_size smaller than sample gain; setting k_samples to 1.")
+            self._state.k_samples = 1
+        elif self._state.k_samples is None:
+            ez.logger.warning("k_samples is None; z-score accumulation will be unbounded.")
+        self._state.samples = deque(maxlen=self._state.k_samples)
+        self._state.min_samples = (
+            int(np.ceil(self.settings.min_seconds / message.axes[self.settings.axis].gain))
+            if self.settings.window_size is not None
+            else self.settings.min_samples
+        )
+        if self._state.k_samples is not None and self._state.min_samples > self._state.k_samples:
+            ez.logger.warning("min_samples is greater than k_samples; adjusting min_samples to k_samples.")
+            self._state.min_samples = self._state.k_samples
+
+    def _add_batch_stats(self, x: npt.NDArray) -> None:
+        x = np.asarray(x, dtype=np.float64)
+        n_b = x.shape[0]
+        mean_b = np.mean(x, axis=0)
+        M2_b = np.sum((x - mean_b) ** 2, axis=0)
+
+        if self._state.k_samples is not None and len(self._state.samples) == self._state.k_samples:
+            n_old, mean_old, M2_old = self._state.samples.popleft()
+            N_T = self._state.N
+            N_new = N_T - n_old
+
+            if N_new <= 0:
+                self._state.N = 0
+                self._state.mean = np.zeros_like(self._state.mean)
+                self._state.M2 = np.zeros_like(self._state.M2)
+            else:
+                delta = mean_old - self._state.mean
+                self._state.N = N_new
+                self._state.mean = (N_T * self._state.mean - n_old * mean_old) / N_new
+                self._state.M2 = self._state.M2 - M2_old - (delta * delta) * (N_T * n_old / N_new)
+
+        N_A = self._state.N
+        N = N_A + n_b
+        delta = mean_b - self._state.mean
+        self._state.mean = self._state.mean + delta * (n_b / N)
+        self._state.M2 = self._state.M2 + M2_b + (delta * delta) * (N_A * n_b / N)
+        self._state.N = N
+
+        self._state.samples.append((n_b, mean_b, M2_b))
+
+    def partial_fit(self, message: SampleMessage) -> None:
+        x = message.sample.data
+        self._add_batch_stats(x)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        if self._state.N == 0 or self._state.N < self._state.min_samples:
+            if self.settings.update_with_signal:
+                x = message.data
+                if self.settings.artifact_z_thresh is not None and self._state.N > 0:
+                    varis = self._state.M2 / self._state.N
+                    std = np.maximum(np.sqrt(varis), 1e-8)
+                    z = np.abs((x - self._state.mean) / std)
+                    mask = np.any(z > self.settings.artifact_z_thresh, axis=1)
+                    x = x[~mask]
+                if x.size > 0:
+                    self._add_batch_stats(x)
+            return message
+
+        varis = self._state.M2 / self._state.N
+        std = np.maximum(np.sqrt(varis), 1e-8)
+        with np.errstate(divide="ignore", invalid="ignore"):
+            result = (message.data - self._state.mean) / std
+        result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
+        if self.settings.clip is not None:
+            result = np.clip(result, -self.settings.clip, self.settings.clip)
+
+        if self.settings.update_with_signal:
+            x = message.data
+            if self.settings.artifact_z_thresh is not None:
+                z_scores = np.abs((x - self._state.mean) / std)
+                mask = np.any(z_scores > self.settings.artifact_z_thresh, axis=1)
+                x = x[~mask]
+            if x.size > 0:
+                self._add_batch_stats(x)
+
+        return replace(message, data=result)
+
+
+class RollingScalerUnit(
+    BaseAdaptiveTransformerUnit[
+        RollingScalerSettings,
+        AxisArray,
+        AxisArray,
+        RollingScalerProcessor,
+    ]
+):
+    """
+    Unit wrapper for :obj:`RollingScalerProcessor`.
+
+    This unit performs rolling z-score normalization on incoming `AxisArray` messages. The unit maintains rolling
+    statistics (mean and variance) over the last `k_samples` samples received. When processing an `AxisArray` message,
+    it normalizes the data using the current rolling statistics.
+
+    Example:
+    -----------------------------
+    ```python
+    unit = RollingScalerUnit(
+        settings=RollingScalerSettings(
+            k_samples=20  # Number of previous samples to use for rolling statistics
+        )
+    )
+    ```
+    """
+
+    SETTINGS = RollingScalerSettings

ezmsg/sigproc/sampler.py

@@ -1,28 +1,28 @@
 import asyncio
-from collections import deque
 import copy
 import traceback
 import typing
+from collections import deque
 
-import numpy as np
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messages.axisarray import (
     AxisArray,
 )
 from ezmsg.util.messages.util import replace
 
-from .util.profile import profile_subpub
-from .util.axisarray_buffer import HybridAxisArrayBuffer
-from .util.buffer import UpdateStrategy
-from .util.message import SampleMessage, SampleTriggerMessage
 from .base import (
-    BaseStatefulTransformer,
     BaseConsumerUnit,
-    BaseTransformerUnit,
-    BaseStatefulProducer,
     BaseProducerUnit,
+    BaseStatefulProducer,
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
     processor_state,
 )
+from .util.axisarray_buffer import HybridAxisArrayBuffer
+from .util.buffer import UpdateStrategy
+from .util.message import SampleMessage, SampleTriggerMessage
+from .util.profile import profile_subpub
 
 
 class SamplerSettings(ez.Settings):
@@ -74,12 +74,8 @@ class SamplerState:
     triggers: deque[SampleTriggerMessage] | None = None
 
 
-class SamplerTransformer(
-    BaseStatefulTransformer[SamplerSettings, AxisArray, AxisArray, SamplerState]
-):
-    def __call__(
-        self, message: AxisArray | SampleTriggerMessage
-    ) -> list[SampleMessage]:
+class SamplerTransformer(BaseStatefulTransformer[SamplerSettings, AxisArray, AxisArray, SamplerState]):
+    def __call__(self, message: AxisArray | SampleTriggerMessage) -> list[SampleMessage]:
         # TODO: Currently we have a single entry point that accepts both
         #  data and trigger messages and we choose a code path based on
         #  the message type. However, in the future we will likely replace
@@ -99,9 +95,7 @@ class SamplerTransformer(
         # Compute hash based on message properties that require state reset
         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 :]
-        )
+        sample_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
         return hash((sample_shape, message.key))
 
     def _reset_state(self, message: AxisArray) -> None:
@@ -193,20 +187,14 @@ class SamplerTransformer(
         trigger_ts: float = message.timestamp
         if not self.settings.estimate_alignment:
             # Override the trigger timestamp with the next sample's likely timestamp.
-            trigger_ts = (
-                self._state.buffer.axis_final_value + self._state.buffer.axis_gain
-            )
+            trigger_ts = self._state.buffer.axis_final_value + self._state.buffer.axis_gain
 
-        new_trig_msg = replace(
-            message, timestamp=trigger_ts, period=_period, value=_value
-        )
+        new_trig_msg = replace(message, timestamp=trigger_ts, period=_period, value=_value)
         self._state.triggers.append(new_trig_msg)
         return []
 
 
-class Sampler(
-    BaseTransformerUnit[SamplerSettings, AxisArray, AxisArray, SamplerTransformer]
-):
+class Sampler(BaseTransformerUnit[SamplerSettings, AxisArray, AxisArray, SamplerTransformer]):
     SETTINGS = SamplerSettings
 
     INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage)
@@ -269,19 +257,13 @@ class TriggerGeneratorState:
     output: int = 0
 
 
-class TriggerProducer(
-    BaseStatefulProducer[
-        TriggerGeneratorSettings, SampleTriggerMessage, TriggerGeneratorState
-    ]
-):
+class TriggerProducer(BaseStatefulProducer[TriggerGeneratorSettings, SampleTriggerMessage, TriggerGeneratorState]):
     def _reset_state(self) -> None:
         self._state.output = 0
 
     async def _produce(self) -> SampleTriggerMessage:
         await asyncio.sleep(self.settings.publish_period)
-        out_msg = SampleTriggerMessage(
-            period=self.settings.period, value=self._state.output
-        )
+        out_msg = SampleTriggerMessage(period=self.settings.period, value=self._state.output)
         self._state.output += 1
         return out_msg
 

ezmsg/sigproc/scaler.py

@@ -1,27 +1,25 @@
 import typing
 
 import numpy as np
+from ezmsg.util.generator import consumer
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
-from ezmsg.util.generator import consumer
 
 from .base import (
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
 )
-from .ewma import EWMATransformer, EWMASettings, _alpha_from_tau
 
 # 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 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]:
+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.
@@ -85,19 +83,13 @@ class AdaptiveStandardScalerTransformer(
     ]
 ):
     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
-        )
+        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
         mean_message = self._state.samps_ewma(message)
-        var_sq_message = self._state.vars_sq_ewma(
-            replace(message, data=message.data**2)
-        )
+        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
@@ -119,9 +111,7 @@ class AdaptiveStandardScaler(
 
 
 # Backwards compatibility...
-def scaler_np(
-    time_constant: float = 1.0, axis: str | None = None
-) -> AdaptiveStandardScalerTransformer:
+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,8 +1,8 @@
 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 ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from .base import (
     BaseAsyncTransformer,
@@ -27,15 +27,11 @@ class SignalInjectorState:
 
 
 class SignalInjectorTransformer(
-    BaseAsyncTransformer[
-        SignalInjectorSettings, AxisArray, AxisArray, SignalInjectorState
-    ]
+    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 :]
-        )
+        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:
@@ -44,9 +40,7 @@ class SignalInjectorTransformer(
         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 :]
-        )
+        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
@@ -63,11 +57,7 @@ class SignalInjectorTransformer(
         return out_msg
 
 
-class SignalInjector(
-    BaseTransformerUnit[
-        SignalInjectorSettings, AxisArray, AxisArray, SignalInjectorTransformer
-    ]
-):
+class SignalInjector(BaseTransformerUnit[SignalInjectorSettings, AxisArray, AxisArray, SignalInjectorTransformer]):
     SETTINGS = SignalInjectorSettings
     INPUT_FREQUENCY = ez.InputStream(float | None)
     INPUT_AMPLITUDE = ez.InputStream(float)

ezmsg/sigproc/slicer.py

@@ -1,11 +1,11 @@
+import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
-import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import (
     AxisArray,
-    slice_along_axis,
     AxisBase,
     replace,
+    slice_along_axis,
 )
 
 from .base import (
@@ -49,11 +49,7 @@ def parse_slice(
     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
-            ):
+            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)),)
@@ -76,9 +72,7 @@ class SlicerState:
     b_change_dims: bool = False
 
 
-class SlicerTransformer(
-    BaseStatefulTransformer[SlicerSettings, AxisArray, AxisArray, SlicerState]
-):
+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)
@@ -101,11 +95,7 @@ class SlicerTransformer(
             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
-        ):
+        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]
@@ -119,17 +109,10 @@ class SlicerTransformer(
 
         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
-            }
+            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()
-            }
+            replace_kwargs["axes"] = {k: (v if k != axis else self._state.new_axis) for k, v in message.axes.items()}
 
         return replace(
             message,
@@ -138,9 +121,7 @@ class SlicerTransformer(
         )
 
 
-class Slicer(
-    BaseTransformerUnit[SlicerSettings, AxisArray, AxisArray, SlicerTransformer]
-):
+class Slicer(BaseTransformerUnit[SlicerSettings, AxisArray, AxisArray, SlicerTransformer]):
     SETTINGS = SlicerSettings
 
 

ezmsg/sigproc/spectral.py

@@ -1,6 +1,6 @@
 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 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

@@ -1,20 +1,21 @@
 from typing import Generator
+
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.modify import modify_axis
 
-from .window import Anchor, WindowTransformer
-from .spectrum import (
-    WindowFunction,
-    SpectralTransform,
-    SpectralOutput,
-    SpectrumTransformer,
-)
 from .base import (
-    CompositeProcessor,
     BaseStatefulProcessor,
     BaseTransformerUnit,
+    CompositeProcessor,
+)
+from .spectrum import (
+    SpectralOutput,
+    SpectralTransform,
+    SpectrumTransformer,
+    WindowFunction,
 )
+from .window import Anchor, WindowTransformer
 
 
 class SpectrogramSettings(ez.Settings):
@@ -41,9 +42,7 @@ class SpectrogramSettings(ez.Settings):
     """The :obj:`SpectralOutput` format."""
 
 
-class SpectrogramTransformer(
-    CompositeProcessor[SpectrogramSettings, AxisArray, AxisArray]
-):
+class SpectrogramTransformer(CompositeProcessor[SpectrogramSettings, AxisArray, AxisArray]):
     @staticmethod
     def _initialize_processors(
         settings: SpectrogramSettings,
@@ -54,9 +53,7 @@ class SpectrogramTransformer(
                 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",
+                zero_pad_until="shift" if settings.window_shift is not None else "input",
                 anchor=settings.window_anchor,
             ),
             "spectrum": SpectrumTransformer(
@@ -69,11 +66,7 @@ class SpectrogramTransformer(
         }
 
 
-class Spectrogram(
-    BaseTransformerUnit[
-        SpectrogramSettings, AxisArray, AxisArray, SpectrogramTransformer
-    ]
-):
+class Spectrogram(BaseTransformerUnit[SpectrogramSettings, AxisArray, AxisArray, SpectrogramTransformer]):
     SETTINGS = SpectrogramSettings