ezmsg-sigproc 1.2.3__py3-none-any.whl → 1.3.2__py3-none-any.whl

ezmsg/sigproc/butterworthfilter.py

@@ -3,19 +3,45 @@ import typing
 import ezmsg.core as ez
 import scipy.signal
 import numpy as np
-
-from .filter import filtergen, Filter, FilterState, FilterSettingsBase
-
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.generator import consumer
 
+from .filter import filtergen, Filter, FilterState, FilterSettingsBase
+
 
 class ButterworthFilterSettings(FilterSettingsBase):
+    """Settings for :obj:`ButterworthFilter`."""
+
     order: int = 0
-    cuton: typing.Optional[float] = None  # Hz
-    cutoff: typing.Optional[float] = None  # Hz
 
-    def filter_specs(self) -> typing.Optional[typing.Tuple[str, typing.Union[float, typing.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:
@@ -37,28 +63,56 @@ def butter(
     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
-    axis_arr_in = AxisArray(np.array([]), dims=[""])
-    axis_arr_out = AxisArray(np.array([]), dims=[""])
+    msg_out = AxisArray(np.array([]), dims=[""])
 
+    # Check parameters
     btype, cutoffs = ButterworthFilterSettings(
         order=order, cuton=cuton, cutoff=cutoff
     ).filter_specs()
 
-    # We cannot calculate coefs yet because we do not know input sample rate
-    coefs = None
-    filter_gen = filtergen(axis, coefs, coef_type)  # Passthrough.
+    # 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:
-        axis_arr_in = yield axis_arr_out
-        if coefs is None and order > 0:
-            fs = 1 / axis_arr_in.axes[axis or axis_arr_in.dims[0]].gain
+        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=fs, output=coef_type
+                order,
+                Wn=cutoffs,
+                btype=btype,
+                fs=1 / msg_in.axes[axis].gain,
+                output=coef_type,
             )
             filter_gen = filtergen(axis, coefs, coef_type)
 
-        axis_arr_out = filter_gen.send(axis_arr_in)
+        msg_out = filter_gen.send(msg_in)
 
 
 class ButterworthFilterState(FilterState):
@@ -66,15 +120,17 @@ class ButterworthFilterState(FilterState):
 
 
 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) -> typing.Optional[typing.Tuple[np.ndarray, np.ndarray]]:
         specs = self.STATE.design.filter_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,89 +1,107 @@
 from dataclasses import replace
-import traceback
-from typing import AsyncGenerator, Optional, Generator
+import typing
 
 import numpy as np
-
-from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
 from ezmsg.util.generator import consumer
 import ezmsg.core as ez
 
+from .base import GenAxisArray
+
 
 @consumer
 def downsample(
-        axis: Optional[str] = None, factor: int = 1
-) -> Generator[AxisArray, AxisArray, None]:
-    axis_arr_in = AxisArray(np.array([]), dims=[""])
-    axis_arr_out = AxisArray(np.array([]), dims=[""])
+    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 = 0
+    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:
-        axis_arr_in = yield axis_arr_out
+        msg_in: AxisArray = yield msg_out
 
         if axis is None:
-            axis = axis_arr_in.dims[0]
-        axis_info = axis_arr_in.get_axis(axis)
-        axis_idx = axis_arr_in.get_axis_idx(axis)
-
-        samples = np.arange(axis_arr_in.data.shape[axis_idx]) + s_idx
-        samples = samples % factor
-        s_idx = samples[-1] + 1
+            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
 
         pub_samples = np.where(samples == 0)[0]
         if len(pub_samples) > 0:
-            new_axes = {ax_name: axis_arr_in.get_axis(ax_name) for ax_name in axis_arr_in.dims}
-            new_offset = axis_info.offset + (axis_info.gain * pub_samples[0].item())
-            new_gain = axis_info.gain * factor
-            new_axes[axis] = replace(axis_info, gain=new_gain, offset=new_offset)
-            down_data = np.take(axis_arr_in.data, pub_samples, axis=axis_idx)
-            axis_arr_out = replace(axis_arr_in, data=down_data, dims=axis_arr_in.dims, axes=new_axes)
+            n_step = pub_samples[0].item()
+            data_slice = pub_samples
         else:
-            axis_arr_out = None
+            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,
+                ),
+            },
+        )
 
 
 class DownsampleSettings(ez.Settings):
-    axis: Optional[str] = None
-    factor: int = 1
-
+    """
+    Settings for :obj:`Downsample` node.
+    See :obj:`downsample` documentation for a description of the parameters.
+    """
 
-class DownsampleState(ez.State):
-    cur_settings: DownsampleSettings
-    gen: Generator
+    axis: typing.Optional[str] = None
+    factor: int = 1
 
 
-class Downsample(ez.Unit):
-    SETTINGS: DownsampleSettings
-    STATE: DownsampleState
+class Downsample(GenAxisArray):
+    """:obj:`Unit` for :obj:`bandpower`."""
 
-    INPUT_SETTINGS = ez.InputStream(DownsampleSettings)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    SETTINGS = DownsampleSettings
 
     def construct_generator(self):
-        self.STATE.gen = downsample(axis=self.STATE.cur_settings.axis, factor=self.STATE.cur_settings.factor)
-
-    def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
-        self.construct_generator()
-
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: DownsampleSettings) -> None:
-        self.STATE.cur_settings = msg
-        self.construct_generator()
-
-    @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)")
-
-        try:
-            out_msg = self.STATE.gen.send(msg)
-            if out_msg is not None:
-                yield self.OUTPUT_SIGNAL, out_msg
-        except (StopIteration, GeneratorExit):
-            ez.logger.debug(f"Downsample closed in {self.address}")
-        except Exception:
-            ez.logger.info(traceback.format_exc())
+        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
@@ -96,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)
@@ -113,7 +127,7 @@ class EWMFilter(ez.Collection):
     def configure(self) -> None:
         self.EWM.apply_settings(
             EWMSettings(
-                axis=self.SETTINGS.axis, 
+                axis=self.SETTINGS.axis,
                 zero_offset=self.SETTINGS.zero_offset,
             )
         )

ezmsg/sigproc/filter.py

@@ -1,25 +1,26 @@
 import asyncio
-import typing
-
 from dataclasses import dataclass, replace, field
+import typing
 
 import ezmsg.core as ez
-import scipy.signal
-
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
 import numpy as np
 import numpy.typing as npt
+import scipy.signal
 
-from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.generator import consumer
 
 @dataclass
 class FilterCoefficients:
     b: np.ndarray = field(default_factory=lambda: np.array([1.0, 0.0]))
     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,...]]:
+    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`.
@@ -31,57 +32,73 @@ def _normalize_coefs(
             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
-    axis_arr_in = AxisArray(np.array([]), dims=[""])
-    axis_arr_out = AxisArray(np.array([]), dims=[""])
+    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
-    axis_idx = None
-    zi = None
-    expected_shape = None
+    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:
-        axis_arr_in = yield axis_arr_out
+        msg_in: AxisArray = yield msg_out
 
         if coefs is None:
             # passthrough if we do not have a filter design.
-            axis_arr_out = axis_arr_in
+            msg_out = msg_in
             continue
 
-        if axis_idx is None:
-            axis_name = axis_arr_in.dims[0] if axis is None else axis
-            axis_idx = axis_arr_in.get_axis_idx(axis_name)
-
-        dat_in = axis_arr_in.data
+        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 = dat_in.shape[:axis_idx] + dat_in.shape[axis_idx + 1 :]
-        if zi is None or samp_shape != expected_shape:
-            expected_shape = samp_shape
-            n_tail = dat_in.ndim - axis_idx - 1
+        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 = dat_in.shape[:axis_idx] + (1,) + dat_in.shape[axis_idx + 1 :]
+            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, dat_in, axis=axis_idx, zi=zi)
-        axis_arr_out = replace(axis_arr_in, data=dat_out)
+        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):
@@ -105,8 +122,8 @@ class FilterState(ez.State):
 
 
 class Filter(ez.Unit):
-    SETTINGS: FilterSettingsBase
-    STATE: FilterState
+    SETTINGS = FilterSettingsBase
+    STATE = FilterState
 
     INPUT_FILTER = ez.InputStream(FilterCoefficients)
     INPUT_SIGNAL = ez.InputStream(AxisArray)
@@ -116,7 +133,7 @@ class Filter(ez.Unit):
         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
 
@@ -205,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),
+        )