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

ezmsg/sigproc/butterworthfilter.py

@@ -1,18 +1,47 @@
+import typing
+
 import ezmsg.core as ez
 import scipy.signal
 import numpy as np
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
 
-from .filter import Filter, FilterState, FilterSettingsBase
-
-from typing import Optional, Tuple, Union
+from .filter import filtergen, Filter, FilterState, FilterSettingsBase
 
 
 class ButterworthFilterSettings(FilterSettingsBase):
+    """Settings for :obj:`ButterworthFilter`."""
+
     order: int = 0
-    cuton: Optional[float] = None  # Hz
-    cutoff: Optional[float] = None  # Hz
 
-    def filter_specs(self) -> Optional[Tuple[str, Union[float, Tuple[float, float]]]]:
+    cuton: typing.Optional[float] = None
+    """
+    Cuton frequency (Hz). If cutoff is not specified then this is the highpass corner, otherwise
+    if it is lower than cutoff then this is the beginning of the bandpass
+    or if it is greater than cuton then it is the end of the bandstop. 
+    """
+
+    cutoff: typing.Optional[float] = None
+    """
+    Cutoff frequency (Hz). If cuton is not specified then this is the lowpass corner, otherwise
+    if it is greater than cuton then this is the end of the bandpass,
+    or if it is less than cuton then it is the beginning of the bandstop. 
+    """
+
+    def filter_specs(
+        self,
+    ) -> typing.Optional[
+        typing.Tuple[str, typing.Union[float, typing.Tuple[float, float]]]
+    ]:
+        """
+        Determine the filter type given the corner frequencies.
+
+        Returns:
+            A tuple with the first element being a string indicating the filter type
+            (one of "lowpass", "highpass", "bandpass", "bandstop")
+            and the second element being the corner frequency or frequencies.
+
+        """
         if self.cuton is None and self.cutoff is None:
             return None
         elif self.cuton is None and self.cutoff is not None:
@@ -26,22 +55,84 @@ class ButterworthFilterSettings(FilterSettingsBase):
                 return "bandstop", (self.cutoff, self.cuton)
 
 
+@consumer
+def butter(
+    axis: typing.Optional[str],
+    order: int = 0,
+    cuton: typing.Optional[float] = None,
+    cutoff: typing.Optional[float] = None,
+    coef_type: str = "ba",
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Apply Butterworth filter to streaming data. Uses :obj:`scipy.signal.butter` to design the filter.
+    See :obj:`ButterworthFilterSettings.filter_specs` for an explanation of specifying different
+    filter types (lowpass, highpass, bandpass, bandstop) from the parameters.
+
+    Args:
+        axis: The name of the axis to filter.
+        order: Filter order.
+        cuton: Corner frequency of the filter in Hz.
+        cutoff: Corner frequency of the filter in Hz.
+        coef_type: "ba" or "sos"
+
+    Returns:
+        A primed generator object which accepts .send(axis_array) and yields filtered axis array.
+
+    """
+    # IO
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    # Check parameters
+    btype, cutoffs = ButterworthFilterSettings(
+        order=order, cuton=cuton, cutoff=cutoff
+    ).filter_specs()
+
+    # State variables
+    # Initialize filtergen as passthrough until we can calculate coefs.
+    filter_gen = filtergen(axis, None, coef_type)
+
+    # Reset if these change.
+    check_input = {"gain": None}
+    # Key not checked because filter_gen will handle resetting if .key changes.
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+        axis = axis or msg_in.dims[0]
+
+        b_reset = msg_in.axes[axis].gain != check_input["gain"]
+        b_reset = b_reset and order > 0  # Not passthrough
+        if b_reset:
+            check_input["gain"] = msg_in.axes[axis].gain
+            coefs = scipy.signal.butter(
+                order,
+                Wn=cutoffs,
+                btype=btype,
+                fs=1 / msg_in.axes[axis].gain,
+                output=coef_type,
+            )
+            filter_gen = filtergen(axis, coefs, coef_type)
+
+        msg_out = filter_gen.send(msg_in)
+
+
 class ButterworthFilterState(FilterState):
     design: ButterworthFilterSettings
 
 
 class ButterworthFilter(Filter):
-    SETTINGS: ButterworthFilterSettings
-    STATE: ButterworthFilterState
+    """:obj:`Unit` for :obj:`butterworth`"""
+
+    SETTINGS = ButterworthFilterSettings
+    STATE = ButterworthFilterState
 
     INPUT_FILTER = ez.InputStream(ButterworthFilterSettings)
 
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         self.STATE.design = self.SETTINGS
         self.STATE.filt_designed = True
-        super().initialize()
+        await super().initialize()
 
-    def design_filter(self) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    def design_filter(self) -> typing.Optional[typing.Tuple[np.ndarray, np.ndarray]]:
         specs = self.STATE.design.filter_specs()
         if self.STATE.design.order > 0 and specs is not None:
             btype, cut = specs

ezmsg/sigproc/decimate.py

@@ -1,7 +1,5 @@
-import ezmsg.core as ez
-
 import scipy.signal
-
+import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 
 from .downsample import Downsample, DownsampleSettings
@@ -9,7 +7,12 @@ from .filter import Filter, FilterCoefficients, FilterSettings
 
 
 class Decimate(ez.Collection):
-    SETTINGS: DownsampleSettings
+    """
+    A :obj:`Collection` chaining a :obj:`Filter` node configured as a lowpass Chebyshev filter
+    and a :obj:`Downsample` node.
+    """
+
+    SETTINGS = DownsampleSettings
 
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)

ezmsg/sigproc/downsample.py

@@ -1,63 +1,107 @@
 from dataclasses import replace
+import typing
 
-from ezmsg.util.messages.axisarray import AxisArray
-
-import ezmsg.core as ez
 import numpy as np
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.generator import consumer
+import ezmsg.core as ez
 
-from typing import (
-    AsyncGenerator,
-    Optional,
-)
-
-
-class DownsampleSettings(ez.Settings):
-    axis: Optional[str] = None
-    factor: int = 1
-
-
-class DownsampleState(ez.State):
-    cur_settings: DownsampleSettings
-    s_idx: int = 0
-
+from .base import GenAxisArray
+
+
+@consumer
+def downsample(
+    axis: typing.Optional[str] = None, factor: int = 1
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Construct a generator that yields a downsampled version of the data .send() to it.
+    Downsampled data simply comprise every `factor`th sample.
+    This should only be used following appropriate lowpass filtering.
+    If your pipeline does not already have lowpass filtering then consider
+    using the :obj:`Decimate` collection instead.
+
+    Args:
+        axis: The name of the axis along which to downsample.
+        factor: Downsampling factor.
+
+    Returns:
+        A primed generator object ready to receive a `.send(axis_array)`
+        and yields the downsampled data.
+        Note that if a send chunk does not have sufficient samples to reach the
+        next downsample interval then `None` is yielded.
+
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    if factor < 1:
+        raise ValueError("Downsample factor must be at least 1 (no downsampling)")
+
+    # state variables
+    s_idx: int = 0  # Index of the next msg's first sample into the virtual rotating ds_factor counter.
+
+    check_input = {"gain": None, "key": None}
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+
+        if axis is None:
+            axis = msg_in.dims[0]
+        axis_info = msg_in.get_axis(axis)
+        axis_idx = msg_in.get_axis_idx(axis)
+
+        b_reset = (
+            msg_in.axes[axis].gain != check_input["gain"]
+            or msg_in.key != check_input["key"]
+        )
+        if b_reset:
+            check_input["gain"] = axis_info.gain
+            check_input["key"] = msg_in.key
+            # Reset state variables
+            s_idx = 0
+
+        n_samples = msg_in.data.shape[axis_idx]
+        samples = np.arange(s_idx, s_idx + n_samples) % factor
+        if n_samples > 0:
+            # Update state for next iteration.
+            s_idx = samples[-1] + 1
 
-class Downsample(ez.Unit):
-    SETTINGS: DownsampleSettings
-    STATE: DownsampleState
+        pub_samples = np.where(samples == 0)[0]
+        if len(pub_samples) > 0:
+            n_step = pub_samples[0].item()
+            data_slice = pub_samples
+        else:
+            n_step = 0
+            data_slice = slice(None, 0, None)
+        msg_out = replace(
+            msg_in,
+            data=slice_along_axis(msg_in.data, data_slice, axis=axis_idx),
+            axes={
+                **msg_in.axes,
+                axis: replace(
+                    axis_info,
+                    gain=axis_info.gain * factor,
+                    offset=axis_info.offset + axis_info.gain * n_step,
+                ),
+            },
+        )
 
-    INPUT_SETTINGS = ez.InputStream(DownsampleSettings)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
-    def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
+class DownsampleSettings(ez.Settings):
+    """
+    Settings for :obj:`Downsample` node.
+    See :obj:`downsample` documentation for a description of the parameters.
+    """
 
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: DownsampleSettings) -> None:
-        self.STATE.cur_settings = msg
+    axis: typing.Optional[str] = None
+    factor: int = 1
 
-    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
-    @ez.publisher(OUTPUT_SIGNAL)
-    async def on_signal(self, msg: AxisArray) -> AsyncGenerator:
-        if self.STATE.cur_settings.factor < 1:
-            raise ValueError("Downsample factor must be at least 1 (no downsampling)")
 
-        axis_name = self.STATE.cur_settings.axis
-        if axis_name is None:
-            axis_name = msg.dims[0]
-        axis = msg.get_axis(axis_name)
-        axis_idx = msg.get_axis_idx(axis_name)
+class Downsample(GenAxisArray):
+    """:obj:`Unit` for :obj:`bandpower`."""
 
-        samples = np.arange(msg.data.shape[axis_idx]) + self.STATE.s_idx
-        samples = samples % self.STATE.cur_settings.factor
-        self.STATE.s_idx = samples[-1] + 1
+    SETTINGS = DownsampleSettings
 
-        pub_samples = np.where(samples == 0)[0]
-        if len(pub_samples) != 0:
-            new_axes = {ax_name: msg.get_axis(ax_name) for ax_name in msg.dims}
-            new_offset = axis.offset + (axis.gain * pub_samples[0].item())
-            new_gain = axis.gain * self.STATE.cur_settings.factor
-            new_axes[axis_name] = replace(axis, gain=new_gain, offset=new_offset)
-            down_data = np.take(msg.data, pub_samples, axis_idx)
-            out_msg = replace(msg, data=down_data, dims=msg.dims, axes=new_axes)
-            yield self.OUTPUT_SIGNAL, out_msg
+    def construct_generator(self):
+        self.STATE.gen = downsample(
+            axis=self.SETTINGS.axis, factor=self.SETTINGS.factor
+        )

ezmsg/sigproc/ewmfilter.py

@@ -1,19 +1,20 @@
 import asyncio
 from dataclasses import replace
+import typing
 
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
-
 import numpy as np
 
 from .window import Window, WindowSettings
 
-from typing import AsyncGenerator, Optional
-
 
 class EWMSettings(ez.Settings):
-    axis: Optional[str] = None
-    zero_offset: bool = True  # If true, we assume zero DC offset
+    axis: typing.Optional[str] = None
+    """Name of the axis to accumulate."""
+
+    zero_offset: bool = True
+    """If true, we assume zero DC offset for input data."""
 
 
 class EWMState(ez.State):
@@ -28,14 +29,14 @@ class EWM(ez.Unit):
     References https://stackoverflow.com/a/42926270
     """
 
-    SETTINGS: EWMSettings
-    STATE: EWMState
+    SETTINGS = EWMSettings
+    STATE = EWMState
 
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     INPUT_BUFFER = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         self.STATE.signal_queue = asyncio.Queue()
         self.STATE.buffer_queue = asyncio.Queue()
 
@@ -48,7 +49,7 @@ class EWM(ez.Unit):
         self.STATE.buffer_queue.put_nowait(message)
 
     @ez.publisher(OUTPUT_SIGNAL)
-    async def sync_output(self) -> AsyncGenerator:
+    async def sync_output(self) -> typing.AsyncGenerator:
         while True:
             signal = await self.STATE.signal_queue.get()
             buffer = await self.STATE.buffer_queue.get()  # includes signal
@@ -73,9 +74,12 @@ class EWM(ez.Unit):
             buffer_data = buffer.data
             buffer_data = np.moveaxis(buffer_data, axis_idx, 0)
 
+            while scale_arr.ndim < buffer_data.ndim:
+                scale_arr = scale_arr[..., None]
+
             def ewma(data: np.ndarray) -> np.ndarray:
-                mult = scale_arr[:, np.newaxis] * data * pw0
-                out = scale_arr[::-1, np.newaxis] * mult.cumsum(axis=0)
+                mult = scale_arr * data * pw0
+                out = scale_arr[::-1] * mult.cumsum(axis=0)
 
                 if not self.SETTINGS.zero_offset:
                     out = (data[0, :, np.newaxis] * pows[1:]).T + out
@@ -93,13 +97,26 @@ class EWM(ez.Unit):
 
 
 class EWMFilterSettings(ez.Settings):
-    history_dur: float  # previous data to accumulate for standardization
-    axis: Optional[str] = None
-    zero_offset: bool = True  # If true, we assume zero DC offset for input data
+    history_dur: float
+    """Previous data to accumulate for standardization."""
+
+    axis: typing.Optional[str] = None
+    """Name of the axis to accumulate."""
+
+    zero_offset: bool = True
+    """If true, we assume zero DC offset for input data."""
 
 
 class EWMFilter(ez.Collection):
-    SETTINGS: EWMFilterSettings
+    """
+    A :obj:`Collection` that splits the input into a branch that
+    leads to :obj:`Window` which then feeds into :obj:`EWM` 's INPUT_BUFFER
+    and another branch that feeds directly into :obj:`EWM` 's INPUT_SIGNAL.
+
+    Consider :obj:`scaler` for a more efficient alternative.
+    """
+
+    SETTINGS = EWMFilterSettings
 
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
@@ -108,7 +125,12 @@ class EWMFilter(ez.Collection):
     EWM = EWM()
 
     def configure(self) -> None:
-        self.EWM.apply_settings(EWMSettings(axis=self.SETTINGS.axis, zero_offset=True))
+        self.EWM.apply_settings(
+            EWMSettings(
+                axis=self.SETTINGS.axis,
+                zero_offset=self.SETTINGS.zero_offset,
+            )
+        )
 
         self.WINDOW.apply_settings(
             WindowSettings(

ezmsg/sigproc/filter.py

@@ -1,13 +1,13 @@
+import asyncio
 from dataclasses import dataclass, replace, field
+import typing
 
 import ezmsg.core as ez
-import scipy.signal
-import numpy as np
-import asyncio
-
 from ezmsg.util.messages.axisarray import AxisArray
-
-from typing import AsyncGenerator, Optional, Tuple
+from ezmsg.util.generator import consumer
+import numpy as np
+import numpy.typing as npt
+import scipy.signal
 
 
 @dataclass
@@ -16,39 +16,124 @@ class FilterCoefficients:
     a: np.ndarray = field(default_factory=lambda: np.array([1.0, 0.0]))
 
 
+def _normalize_coefs(
+    coefs: typing.Union[
+        FilterCoefficients, typing.Tuple[npt.NDArray, npt.NDArray], npt.NDArray
+    ],
+) -> typing.Tuple[str, typing.Tuple[npt.NDArray, ...]]:
+    coef_type = "ba"
+    if coefs is not None:
+        # scipy.signal functions called with first arg `*coefs`.
+        # Make sure we have a tuple of coefficients.
+        if isinstance(coefs, npt.NDArray):
+            coef_type = "sos"
+            coefs = (coefs,)  # sos funcs just want a single ndarray.
+        elif isinstance(coefs, FilterCoefficients):
+            coefs = (FilterCoefficients.b, FilterCoefficients.a)
+    return coef_type, coefs
+
+
+@consumer
+def filtergen(
+    axis: str, coefs: typing.Optional[typing.Tuple[np.ndarray]], coef_type: str
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Construct a generic filter generator function.
+
+    Args:
+        axis: The name of the axis to operate on.
+        coefs: The pre-calculated filter coefficients.
+        coef_type: The type of filter coefficients. One of "ba" or "sos".
+
+    Returns:
+        A generator that expects .send(axis_array) and yields the filtered :obj:`AxisArray`.
+    """
+    # Massage inputs
+    if coefs is not None and not isinstance(coefs, tuple):
+        # scipy.signal functions called with first arg `*coefs`, but sos coefs are a single ndarray.
+        coefs = (coefs,)
+
+    # Init IO
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    filt_func = {"ba": scipy.signal.lfilter, "sos": scipy.signal.sosfilt}[coef_type]
+    zi_func = {"ba": scipy.signal.lfilter_zi, "sos": scipy.signal.sosfilt_zi}[coef_type]
+
+    # State variables
+    zi: typing.Optional[npt.NDArray] = None
+
+    # Reset if these change.
+    check_input = {"key": None, "shape": None}
+    # fs changing will be handled by caller that creates coefficients.
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+
+        if coefs is None:
+            # passthrough if we do not have a filter design.
+            msg_out = msg_in
+            continue
+
+        axis = msg_in.dims[0] if axis is None else axis
+        axis_idx = msg_in.get_axis_idx(axis)
+
+        # Re-calculate/reset zi if necessary
+        samp_shape = msg_in.data.shape[:axis_idx] + msg_in.data.shape[axis_idx + 1 :]
+        b_reset = samp_shape != check_input["shape"]
+        b_reset = b_reset or msg_in.key != check_input["key"]
+        if b_reset:
+            check_input["shape"] = samp_shape
+            check_input["key"] = msg_in.key
+
+            n_tail = msg_in.data.ndim - axis_idx - 1
+            zi = zi_func(*coefs)
+            zi_expand = (None,) * axis_idx + (slice(None),) + (None,) * n_tail
+            n_tile = (
+                msg_in.data.shape[:axis_idx] + (1,) + msg_in.data.shape[axis_idx + 1 :]
+            )
+            if coef_type == "sos":
+                # sos zi must keep its leading dimension (`order / 2` for low|high; `order` for bpass|bstop)
+                zi_expand = (slice(None),) + zi_expand
+                n_tile = (1,) + n_tile
+            zi = np.tile(zi[zi_expand], n_tile)
+
+        dat_out, zi = filt_func(*coefs, msg_in.data, axis=axis_idx, zi=zi)
+        msg_out = replace(msg_in, data=dat_out)
+
+
 class FilterSettingsBase(ez.Settings):
-    axis: Optional[str] = None
-    fs: Optional[float] = None
+    axis: typing.Optional[str] = None
+    fs: typing.Optional[float] = None
 
 
 class FilterSettings(FilterSettingsBase):
     # If you'd like to statically design a filter, define it in settings
-    filt: Optional[FilterCoefficients] = None
+    filt: typing.Optional[FilterCoefficients] = None
 
 
 class FilterState(ez.State):
-    axis: Optional[str] = None
-    zi: Optional[np.ndarray] = None
+    axis: typing.Optional[str] = None
+    zi: typing.Optional[np.ndarray] = None
     filt_designed: bool = False
-    filt: Optional[FilterCoefficients] = None
+    filt: typing.Optional[FilterCoefficients] = None
     filt_set: asyncio.Event = field(default_factory=asyncio.Event)
-    samp_shape: Optional[Tuple[int, ...]] = None
-    fs: Optional[float] = None  # Hz
+    samp_shape: typing.Optional[typing.Tuple[int, ...]] = None
+    fs: typing.Optional[float] = None  # Hz
 
 
 class Filter(ez.Unit):
-    SETTINGS: FilterSettingsBase
-    STATE: FilterState
+    SETTINGS = FilterSettingsBase
+    STATE = FilterState
 
     INPUT_FILTER = ez.InputStream(FilterCoefficients)
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
-    def design_filter(self) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    def design_filter(self) -> typing.Optional[typing.Tuple[np.ndarray, np.ndarray]]:
         raise NotImplementedError("Must implement 'design_filter' in Unit subclass!")
 
     # Set up filter with static initialization if specified
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         if self.SETTINGS.axis is not None:
             self.STATE.axis = self.SETTINGS.axis
 
@@ -84,7 +169,7 @@ class Filter(ez.Unit):
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
-    async def apply_filter(self, msg: AxisArray) -> AsyncGenerator:
+    async def apply_filter(self, msg: AxisArray) -> typing.AsyncGenerator:
         axis_name = msg.dims[0] if self.STATE.axis is None else self.STATE.axis
         axis_idx = msg.get_axis_idx(axis_name)
         axis = msg.get_axis(axis_name)
@@ -137,4 +222,7 @@ class Filter(ez.Unit):
         if one_dimensional:
             arr_out = np.squeeze(arr_out, axis=1)
 
-        yield self.OUTPUT_SIGNAL, replace(msg, data=arr_out),
+        yield (
+            self.OUTPUT_SIGNAL,
+            replace(msg, data=arr_out),
+        )