ezmsg-sigproc 2.9.0__tar.gz → 2.10.0__tar.gz

{ezmsg_sigproc-2.9.0 → ezmsg_sigproc-2.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 2.9.0
+Version: 2.10.0
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT

{ezmsg_sigproc-2.9.0 → ezmsg_sigproc-2.10.0}/docs/source/guides/tutorials/signalprocessing.rst

@@ -109,7 +109,7 @@ First, we need to install the `ezmsg-sigproc` package if we haven't already. Thi
 
 .. code-block:: bash
 
-    pip install "ezmsg[sigproc]"
+    pip install "ezmsg-sigproc"
 
 
 |ezmsg_logo_small| Creating the `Downsample` signal processor
@@ -237,7 +237,7 @@ The first two methods deal with the state of the processor (and are only require
 
 In order to implement these methods, we need to understand our preferred message format: `AxisArray`. This is a flexible and powerful class for handling multi-dimensional arrays with named axes, which is particularly useful for signal processing applications. I have already used `AxisArray` in our code as the input message and output message types.
 
-A detailed explanation of the `AxisArray` class is beyond the scope of this tutorial, but you can refer to the :doc:`AxisArray explainer <../explanations/axisarray>` as well as the :doc:`API reference <../reference/API/axisarray>` for more information.
+A detailed explanation of the `AxisArray` class is beyond the scope of this tutorial, but you can refer to the `AxisArray explainer <https://www.ezmsg.org/explanations/axisarray.html>`_ as well as the `API reference <https://www.ezmsg.org/ezmsg/reference/API/axisarray.html>`_ for more information.
 
 Brief Aside on AxisArray
 =================================
@@ -555,25 +555,6 @@ In a separate Python file in the same directory, you can test the `DownsampleTra
 
 Doing the above is very handy for unit testing your processor as well as for offline processing of data.
 
-.. note:: The `downsample` module in `ezmsg-sigproc` has a utility function for creating a `DownsampleTransformer` instance with the desired settings:
-
-    .. code-block:: python
-
-        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)
-            )
-
-    After importing this utility function, lines 8 and 9 in our code above could now read:
-
-    .. code-block:: python
-
-        downsampler = downsample(axis="time", target_rate=50)
-
 Of course, the real power of `ezmsg` comes from integrating your processor into an `ezmsg` Unit and using it in a processing pipeline. We will see how to do this next.
 
 
@@ -600,15 +581,15 @@ A lot of the behind-the-scenes work is done for you by the `BaseTransformerUnit`
         SETTINGS = DownsampleSettings
 
 
-Connecting it to other `Component`\ s and initialising the transformer are accomplished in the same way that we did in the :doc:`pipeline tutorial <pipeline>`.
+Connecting it to other `Component`\ s and initialising the transformer are accomplished in the same way that we did in the `Pipeline Tutorial <https://www.ezmsg.org/tutorials/pipeline.html>`_
 
 
 |ezmsg_logo_small| See Also
 ************************************
 
 - `Further examples <https://github.com/ezmsg-org/ezmsg/tree/master/examples>`_ can be found in the examples directory in `ezmsg`. These are examples of creating and using `ezmsg` Units and pipelines.
-- `ezmsg-sigproc` has a large number of already implemented signal processors. More information can be found at the :doc:`ezmsg-sigproc reference <../extensions/sigproc/content-sigproc>`.
-- `Downsample` class reference
+- `ezmsg-sigproc` has a large number of already implemented signal processors. More information can be found at the :doc:`ezmsg-sigproc reference <../sigproc/content-sigproc>`.
+- :doc:`Downsample class reference <../../api/generated/ezmsg.sigproc.downsample>`
 
 .. |ezmsg_logo_small| image:: ../_static/_images/ezmsg_logo.png
   :width: 40

{ezmsg_sigproc-2.9.0 → ezmsg_sigproc-2.10.0}/src/ezmsg/sigproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.9.0'
-__version_tuple__ = version_tuple = (2, 9, 0)
+__version__ = version = '2.10.0'
+__version_tuple__ = version_tuple = (2, 10, 0)
 
 __commit_id__ = commit_id = None

ezmsg_sigproc-2.10.0/src/ezmsg/sigproc/butterworthzerophase.py

@@ -0,0 +1,305 @@
+"""
+Streaming zero-phase Butterworth filter implemented as a two-stage composite processor.
+
+Stage 1: Forward causal Butterworth filter (from ezmsg.sigproc.butterworthfilter)
+Stage 2: Backward acausal filter with buffering (ButterworthBackwardFilterTransformer)
+
+The output is delayed by `pad_length` samples to ensure the backward pass has sufficient
+future context. The pad_length is computed analytically using scipy's heuristic.
+"""
+
+import functools
+import typing
+
+import numpy as np
+import scipy.signal
+from ezmsg.baseproc import BaseTransformerUnit
+from ezmsg.baseproc.composite import CompositeProcessor
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .butterworthfilter import (
+    ButterworthFilterSettings,
+    ButterworthFilterTransformer,
+    butter_design_fun,
+)
+from .filter import BACoeffs, FilterByDesignTransformer, SOSCoeffs
+from .util.axisarray_buffer import HybridAxisArrayBuffer
+
+
+class ButterworthZeroPhaseSettings(ButterworthFilterSettings):
+    """
+    Settings for :obj:`ButterworthZeroPhase`.
+
+    This implements a streaming zero-phase Butterworth filter using forward-backward
+    filtering. The output is delayed by `pad_length` samples to ensure the backward
+    pass has sufficient future context.
+
+    The pad_length is computed by finding where the filter's impulse response decays
+    to `settle_cutoff` fraction of its peak value. This accounts for the filter's
+    actual time constant rather than just its order.
+    """
+
+    # Inherits from ButterworthFilterSettings:
+    # axis, coef_type, order, cuton, cutoff, wn_hz
+
+    settle_cutoff: float = 0.01
+    """
+    Fraction of peak impulse response used to determine settling time.
+    The pad_length is set to the number of samples until the impulse response
+    decays to this fraction of its peak. Default is 0.01 (1% of peak).
+    """
+
+    max_pad_duration: float | None = None
+    """
+    Maximum pad duration in seconds. If set, the pad_length will be capped
+    at this value times the sampling rate. Use this to limit latency for
+    filters with very long impulse responses. Default is None (no limit).
+    """
+
+
+class ButterworthBackwardFilterTransformer(FilterByDesignTransformer[ButterworthFilterSettings, BACoeffs | SOSCoeffs]):
+    """
+    Backward (acausal) Butterworth filter with buffering.
+
+    This transformer buffers its input and applies the filter in reverse,
+    outputting only the "settled" portion where transients have decayed.
+    This introduces a lag of ``pad_length`` samples.
+
+    Intended to be used as stage 2 in a zero-phase filter pipeline, receiving
+    forward-filtered data from a ButterworthFilterTransformer.
+    """
+
+    # Instance attributes (initialized in _reset_state)
+    _buffer: HybridAxisArrayBuffer | None
+    _coefs_cache: BACoeffs | SOSCoeffs | None
+    _zi_tiled: np.ndarray | None
+    _pad_length: int
+
+    def get_design_function(
+        self,
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
+        return functools.partial(
+            butter_design_fun,
+            order=self.settings.order,
+            cuton=self.settings.cuton,
+            cutoff=self.settings.cutoff,
+            coef_type=self.settings.coef_type,
+            wn_hz=self.settings.wn_hz,
+        )
+
+    def _compute_pad_length(self, fs: float) -> int:
+        """
+        Compute pad length based on the filter's impulse response settling time.
+
+        The pad_length is determined by finding where the impulse response decays
+        to `settle_cutoff` fraction of its peak value. This is then optionally
+        capped by `max_pad_duration`.
+
+        Args:
+            fs: Sampling frequency in Hz.
+
+        Returns:
+            Number of samples for the pad length.
+        """
+        # Design the filter to compute impulse response
+        coefs = self.get_design_function()(fs)
+        if coefs is None:
+            # Filter design failed or is disabled
+            return 0
+
+        # Generate impulse response - use a generous length initially
+        # Start with scipy's heuristic as minimum, then extend if needed
+        if self.settings.coef_type == "ba":
+            min_length = 3 * (self.settings.order + 1)
+        else:
+            n_sections = (self.settings.order + 1) // 2
+            min_length = 3 * n_sections * 2
+
+        # Use 10x the minimum as initial impulse length, or at least 10000 samples
+        # (10000 samples allows for ~333ms at 30kHz, covering most practical cases)
+        impulse_length = max(min_length * 10, 10000)
+
+        # Cap impulse length computation if max_pad_duration is set
+        if self.settings.max_pad_duration is not None:
+            max_samples = int(self.settings.max_pad_duration * fs)
+            impulse_length = min(impulse_length, max_samples + 1)
+
+        impulse = np.zeros(impulse_length)
+        impulse[0] = 1.0
+
+        if self.settings.coef_type == "ba":
+            b, a = coefs
+            h = scipy.signal.lfilter(b, a, impulse)
+        else:
+            h = scipy.signal.sosfilt(coefs, impulse)
+
+        # Find where impulse response settles to settle_cutoff of peak
+        abs_h = np.abs(h)
+        peak = abs_h.max()
+        if peak == 0:
+            return min_length
+
+        threshold = self.settings.settle_cutoff * peak
+        above_threshold = np.where(abs_h > threshold)[0]
+
+        if len(above_threshold) == 0:
+            pad_length = min_length
+        else:
+            pad_length = above_threshold[-1] + 1
+
+        # Ensure at least the scipy heuristic minimum
+        pad_length = max(pad_length, min_length)
+
+        # Apply max_pad_duration cap if set
+        if self.settings.max_pad_duration is not None:
+            max_samples = int(self.settings.max_pad_duration * fs)
+            pad_length = min(pad_length, max_samples)
+
+        return pad_length
+
+    def _reset_state(self, message: AxisArray) -> None:
+        """Reset filter state when stream changes."""
+        self._coefs_cache = None
+        self._zi_tiled = None
+        self._buffer = None
+        # Compute pad_length based on the message's sampling rate
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        fs = 1 / message.axes[axis].gain
+        self._pad_length = self._compute_pad_length(fs)
+        self.state.needs_redesign = True
+
+    def _compute_zi_tiled(self, data: np.ndarray, ax_idx: int) -> None:
+        """Compute and cache the tiled zi for the given data shape.
+
+        Called once per stream (or after filter redesign). The result is
+        broadcast-ready for multiplication by the edge sample on each chunk.
+        """
+        if self.settings.coef_type == "ba":
+            b, a = self._coefs_cache
+            zi_base = scipy.signal.lfilter_zi(b, a)
+        else:  # sos
+            zi_base = scipy.signal.sosfilt_zi(self._coefs_cache)
+
+        n_tail = data.ndim - ax_idx - 1
+
+        if self.settings.coef_type == "ba":
+            zi_expand = (None,) * ax_idx + (slice(None),) + (None,) * n_tail
+            n_tile = data.shape[:ax_idx] + (1,) + data.shape[ax_idx + 1 :]
+        else:  # sos
+            zi_expand = (slice(None),) + (None,) * ax_idx + (slice(None),) + (None,) * n_tail
+            n_tile = (1,) + data.shape[:ax_idx] + (1,) + data.shape[ax_idx + 1 :]
+
+        self._zi_tiled = np.tile(zi_base[zi_expand], n_tile)
+
+    def _initialize_zi(self, data: np.ndarray, ax_idx: int) -> np.ndarray:
+        """Initialize filter state (zi) scaled by edge value."""
+        if self._zi_tiled is None:
+            self._compute_zi_tiled(data, ax_idx)
+        first_sample = np.take(data, [0], axis=ax_idx)
+        return self._zi_tiled * first_sample
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        ax_idx = message.get_axis_idx(axis)
+        fs = 1 / message.axes[axis].gain
+
+        # Check if we need to redesign filter
+        if self._coefs_cache is None or self.state.needs_redesign:
+            self._coefs_cache = self.get_design_function()(fs)
+            self._pad_length = self._compute_pad_length(fs)
+            self._zi_tiled = None  # Invalidate; recomputed on next use.
+            self.state.needs_redesign = False
+
+            # Initialize buffer with duration based on pad_length
+            # Add some margin to handle variable chunk sizes
+            buffer_duration = (self._pad_length + 1) / fs
+            self._buffer = HybridAxisArrayBuffer(duration=buffer_duration, axis=axis)
+
+        # Early exit if filter is effectively disabled
+        if self._coefs_cache is None or self.settings.order <= 0 or message.data.size <= 0:
+            return message
+
+        # Write new data to buffer
+        self._buffer.write(message)
+        n_available = self._buffer.available()
+        n_output = n_available - self._pad_length
+
+        # If we don't have enough data yet, return empty
+        if n_output <= 0:
+            new_shape = list(message.data.shape)
+            new_shape[ax_idx] = 0
+            empty_data = np.empty(new_shape, dtype=message.data.dtype)
+            return replace(message, data=empty_data)
+
+        # Peek all available data from buffer
+        # Note: HybridAxisArrayBuffer moves the target axis to position 0
+        buffered = self._buffer.peek(n_available)
+        combined = buffered.data
+        buffer_ax_idx = 0  # Buffer always puts time axis at position 0
+
+        # Backward filter on reversed data
+        combined_rev = np.flip(combined, axis=buffer_ax_idx)
+        backward_zi = self._initialize_zi(combined_rev, buffer_ax_idx)
+
+        if self.settings.coef_type == "ba":
+            b, a = self._coefs_cache
+            y_bwd_rev, _ = scipy.signal.lfilter(b, a, combined_rev, axis=buffer_ax_idx, zi=backward_zi)
+        else:  # sos
+            y_bwd_rev, _ = scipy.signal.sosfilt(self._coefs_cache, combined_rev, axis=buffer_ax_idx, zi=backward_zi)
+
+        # Reverse back to get output in correct time order
+        y_bwd = np.flip(y_bwd_rev, axis=buffer_ax_idx)
+
+        # Output the settled portion (first n_output samples)
+        y = y_bwd[:n_output]
+
+        # Advance buffer read head to discard output samples, keep pad_length
+        self._buffer.seek(n_output)
+
+        # Build output with adjusted time axis
+        # LinearAxis offset is already correct from the buffer
+        out_axis = buffered.axes[axis]
+
+        # Move axis back to original position if needed
+        if ax_idx != 0:
+            y = np.moveaxis(y, 0, ax_idx)
+
+        return replace(
+            message,
+            data=y,
+            axes={**message.axes, axis: out_axis},
+        )
+
+
+class ButterworthZeroPhaseTransformer(CompositeProcessor[ButterworthZeroPhaseSettings, AxisArray, AxisArray]):
+    """
+    Streaming zero-phase Butterworth filter as a composite of two stages.
+
+    Stage 1 (forward): Standard causal Butterworth filter with state
+    Stage 2 (backward): Acausal Butterworth filter with buffering
+
+    The output is delayed by ``pad_length`` samples.
+    """
+
+    @staticmethod
+    def _initialize_processors(
+        settings: ButterworthZeroPhaseSettings,
+    ) -> dict[str, typing.Any]:
+        # Both stages use the same filter design settings
+        return {
+            "forward": ButterworthFilterTransformer(settings),
+            "backward": ButterworthBackwardFilterTransformer(settings),
+        }
+
+    @classmethod
+    def get_message_type(cls, dir: str) -> type[AxisArray]:
+        if dir in ("in", "out"):
+            return AxisArray
+        raise ValueError(f"Invalid direction: {dir}. Must be 'in' or 'out'.")
+
+
+class ButterworthZeroPhase(
+    BaseTransformerUnit[ButterworthZeroPhaseSettings, AxisArray, AxisArray, ButterworthZeroPhaseTransformer]
+):
+    SETTINGS = ButterworthZeroPhaseSettings

{ezmsg_sigproc-2.9.0 → ezmsg_sigproc-2.10.0}/tests/integration/ezmsg/test_butterworthzerophase_system.py

@@ -2,12 +2,17 @@ import tempfile
 from pathlib import Path
 
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messagecodec import message_log
 from ezmsg.util.messagelogger import MessageLogger
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.terminate import TerminateOnTotal
 
-from ezmsg.sigproc.butterworthzerophase import ButterworthZeroPhase
+from ezmsg.sigproc.butterworthzerophase import (
+    ButterworthBackwardFilterTransformer,
+    ButterworthZeroPhase,
+    ButterworthZeroPhaseSettings,
+)
 from tests.helpers.synth import EEGSynth
 
 
@@ -16,6 +21,16 @@ def test_butterworth_zero_phase_system():
     n_time = 50
     n_total = 10
     n_channels = 96
+    order = 4
+    cuton = 30.0
+    cutoff = 45.0
+    coef_type = "sos"
+
+    # Compute expected pad_length for this filter configuration
+    settings = ButterworthZeroPhaseSettings(order=order, cuton=cuton, cutoff=cutoff, coef_type=coef_type)
+    backward = ButterworthBackwardFilterTransformer(settings)
+    pad_length = backward._compute_pad_length(fs)
+
     test_filename = Path(tempfile.gettempdir())
     test_filename = test_filename / Path("test_butterworth_zero_phase_system.txt")
     with open(test_filename, "w"):
@@ -25,10 +40,10 @@ def test_butterworth_zero_phase_system():
     comps = {
         "SRC": EEGSynth(n_time=n_time, fs=fs, n_ch=n_channels, alpha_freq=10.0),
         "BUTTER": ButterworthZeroPhase(
-            order=4,
-            cuton=30.0,
-            cutoff=45.0,
-            coef_type="sos",
+            order=order,
+            cuton=cuton,
+            cutoff=cutoff,
+            coef_type=coef_type,
         ),
         "LOG": MessageLogger(output=test_filename),
         "TERM": TerminateOnTotal(total=n_total),
@@ -45,7 +60,20 @@ def test_butterworth_zero_phase_system():
     messages = list(message_log(test_filename))
     assert len(messages) >= n_total
 
+    total_output_samples = 0
     for msg in messages:
         assert isinstance(msg, AxisArray)
+        # Non-time dimensions must always be preserved, even during warmup
         assert msg.data.shape[1] == n_channels
-        assert msg.data.shape[0] == n_time
+        # Time dimension may be 0 during warmup while buffering for backward pass
+        assert msg.data.shape[0] >= 0
+        # Data should be finite (for non-empty messages)
+        if msg.data.size > 0:
+            assert np.isfinite(msg.data).all()
+        total_output_samples += msg.data.shape[0]
+
+    # Total output should be input samples minus pad_length delay
+    total_input_samples = n_total * n_time
+    expected_output = total_input_samples - pad_length
+    # Allow some tolerance since we may have received slightly more than n_total messages
+    assert total_output_samples >= expected_output