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

ezmsg/sigproc/__init__.py

@@ -1 +1,4 @@
-from .__version__ import __version__
+import importlib.metadata
+
+
+__version__ = importlib.metadata.version("ezmsg-sigproc")

ezmsg/sigproc/affinetransform.py

@@ -0,0 +1,124 @@
+from dataclasses import replace
+import os
+from pathlib import Path
+from typing import Generator, Optional, Union
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer, GenAxisArray
+
+
+@consumer
+def affine_transform(
+    weights: Union[np.ndarray, str, Path],
+    axis: Optional[str] = None,
+    right_multiply: bool = True,
+) -> Generator[AxisArray, AxisArray, None]:
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = AxisArray(np.array([]), dims=[""])
+
+    if isinstance(weights, str):
+        weights = Path(os.path.abspath(os.path.expanduser(weights)))
+    if isinstance(weights, Path):
+        weights = np.loadtxt(weights, delimiter=",")
+    if not right_multiply:
+        weights = weights.T
+    weights = np.ascontiguousarray(weights)
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+
+        if axis is None:
+            axis = axis_arr_in.dims[-1]
+            axis_idx = -1
+        else:
+            axis_idx = axis_arr_in.get_axis_idx(axis)
+
+        data = axis_arr_in.data
+
+        if data.shape[axis_idx] == (weights.shape[0] - 1):
+            # The weights are stacked A|B where A is the transform and B is a single row
+            #  in the equation y = Ax + B. This supports NeuroKey's weights matrices.
+            sample_shape = data.shape[:axis_idx] + (1,) + data.shape[axis_idx+1:]
+            data = np.concatenate((data, np.ones(sample_shape).astype(data.dtype)), axis=axis_idx)
+
+        if axis_idx in [-1, len(axis_arr_in.dims) - 1]:
+            data = np.matmul(data, weights)
+        else:
+            data = np.moveaxis(data, axis_idx, -1)
+            data = np.matmul(data, weights)
+            data = np.moveaxis(data, -1, axis_idx)
+        axis_arr_out = replace(axis_arr_in, data=data)
+
+
+class AffineTransformSettings(ez.Settings):
+    weights: Union[np.ndarray, str, Path]
+    axis: Optional[str] = None
+    right_multiply: bool = True
+
+
+class AffineTransform(GenAxisArray):
+    SETTINGS: AffineTransformSettings
+
+    def construct_generator(self):
+        self.STATE.gen = affine_transform(
+            weights=self.SETTINGS.weights,
+            axis=self.SETTINGS.axis,
+            right_multiply=self.SETTINGS.right_multiply,
+        )
+
+
+@consumer
+def common_rereference(
+    mode: str = "mean", axis: Optional[str] = None, include_current: bool = True
+) -> Generator[AxisArray, AxisArray, None]:
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = AxisArray(np.array([]), dims=[""])
+
+    func = {"mean": np.mean, "median": np.median}[mode]
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+
+        if axis is None:
+            axis = axis_arr_in.dims[-1]
+            axis_idx = -1
+        else:
+            axis_idx = axis_arr_in.get_axis_idx(axis)
+
+        ref_data = func(axis_arr_in.data, axis=axis_idx, keepdims=True)
+
+        if not include_current:
+            # Typical `CAR = x[0]/N + x[1]/N + ... x[i-1]/N + x[i]/N + x[i+1]/N + ... + x[N-1]/N`
+            # and is the same for all i, so it is calculated only once in `ref_data`.
+            # However, if we had excluded the current channel,
+            # then we would have omitted the contribution of the current channel:
+            # `CAR[i] = x[0]/(N-1) + x[1]/(N-1) + ... x[i-1]/(N-1) + x[i+1]/(N-1) + ... + x[N-1]/(N-1)`
+            # The majority of the calculation is the same as when the current channel is included;
+            # we need only rescale CAR so the divisor is `N-1` instead of `N`, then subtract the contribution
+            # from the current channel (i.e., `x[i] / (N-1)`)
+            #  i.e., `CAR[i] = (N / (N-1)) * common_CAR - x[i]/(N-1)`
+            # We can use broadcasting subtraction instead of looping over channels.
+            N = axis_arr_in.data.shape[axis_idx]
+            ref_data = (N / (N - 1)) * ref_data - axis_arr_in.data / (N - 1)
+            # Side note: I profiled using affine_transform and it's about 30x slower than this implementation.
+
+        axis_arr_out = replace(axis_arr_in, data=axis_arr_in.data - ref_data)
+
+
+class CommonRereferenceSettings(ez.Settings):
+    mode: str = "mean"
+    axis: Optional[str] = None
+    include_current: bool = True
+
+
+class CommonRereference(GenAxisArray):
+    SETTINGS: CommonRereferenceSettings
+
+    def construct_generator(self):
+        self.STATE.gen = common_rereference(
+            mode=self.SETTINGS.mode,
+            axis=self.SETTINGS.axis,
+            include_current=self.SETTINGS.include_current,
+        )

ezmsg/sigproc/aggregate.py

@@ -0,0 +1,103 @@
+from dataclasses import replace
+import typing
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.generator import consumer, GenAxisArray
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.sigproc.spectral import OptionsEnum
+
+
+class AggregationFunction(OptionsEnum):
+    NONE = "None (all)"
+    MAX = "max"
+    MIN = "min"
+    MEAN = "mean"
+    MEDIAN = "median"
+    STD = "std"
+    NANMAX = "nanmax"
+    NANMIN = "nanmin"
+    NANMEAN = "nanmean"
+    NANMEDIAN = "nanmedian"
+    NANSTD = "nanstd"
+
+
+AGGREGATORS = {
+    AggregationFunction.NONE: np.all,
+    AggregationFunction.MAX: np.max,
+    AggregationFunction.MIN: np.min,
+    AggregationFunction.MEAN: np.mean,
+    AggregationFunction.MEDIAN: np.median,
+    AggregationFunction.STD: np.std,
+    AggregationFunction.NANMAX: np.nanmax,
+    AggregationFunction.NANMIN: np.nanmin,
+    AggregationFunction.NANMEAN: np.nanmean,
+    AggregationFunction.NANMEDIAN: np.nanmedian,
+    AggregationFunction.NANSTD: np.nanstd
+}
+
+
+@consumer
+def ranged_aggregate(
+    axis: typing.Optional[str] = None,
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = None,
+    operation: AggregationFunction = AggregationFunction.MEAN
+):
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = AxisArray(np.array([]), dims=[""])
+
+    target_axis: typing.Optional[AxisArray.Axis] = None
+    out_axis = AxisArray.Axis()
+    slices: typing.Optional[typing.List[typing.Tuple[typing.Any, ...]]] = None
+    axis_name = ""
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+        if bands is None:
+            axis_arr_out = axis_arr_in
+        else:
+            if slices is None or target_axis != axis_arr_in.get_axis(axis_name):
+                # Calculate the slices. If we are operating on time axis then
+                axis_name = axis or axis_arr_in.dims[0]
+                ax_idx = axis_arr_in.get_axis_idx(axis_name)
+                target_axis = axis_arr_in.axes[axis_name]
+
+                ax_vec = target_axis.offset + np.arange(axis_arr_in.data.shape[ax_idx]) * target_axis.gain
+                slices = []
+                mids = []
+                for (start, stop) in bands:
+                    inds = np.where(np.logical_and(ax_vec >= start, ax_vec <= stop))[0]
+                    mids.append(np.mean(inds) * target_axis.gain + target_axis.offset)
+                    slices.append(np.s_[inds[0]:inds[-1] + 1])
+                out_axis = AxisArray.Axis(
+                    unit=target_axis.unit, offset=mids[0], gain=(mids[1] - mids[0]) if len(mids) > 1 else 1.0
+                )
+
+            agg_func = AGGREGATORS[operation]
+            out_data = [
+                agg_func(slice_along_axis(axis_arr_in.data, sl, axis=ax_idx), axis=ax_idx)
+                for sl in slices
+            ]
+            new_axes = {**axis_arr_in.axes, axis_name: out_axis}
+            axis_arr_out = replace(
+                axis_arr_in,
+                data=np.stack(out_data, axis=ax_idx),
+                axes=new_axes
+            )
+
+
+class RangedAggregateSettings(ez.Settings):
+    axis: typing.Optional[str] = None
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = None
+    operation: AggregationFunction = AggregationFunction.MEAN
+
+
+class RangedAggregate(GenAxisArray):
+    SETTINGS: RangedAggregateSettings
+
+    def construct_generator(self):
+        self.STATE.gen = ranged_aggregate(
+            axis=self.SETTINGS.axis,
+            bands=self.SETTINGS.bands,
+            operation=self.SETTINGS.operation
+        )

ezmsg/sigproc/bandpower.py

@@ -0,0 +1,53 @@
+from dataclasses import field
+import typing
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer, compose, GenAxisArray
+
+from .spectrogram import spectrogram, SpectrogramSettings
+from .aggregate import ranged_aggregate, AggregationFunction
+
+
+@consumer
+def bandpower(
+    spectrogram_settings: SpectrogramSettings,
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = [(17, 30), (70, 170)]
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = AxisArray(np.array([]), dims=[""])
+
+    f_spec = spectrogram(
+        window_dur=spectrogram_settings.window_dur,
+        window_shift=spectrogram_settings.window_shift,
+        window=spectrogram_settings.window,
+        transform=spectrogram_settings.transform,
+        output=spectrogram_settings.output
+    )
+    f_agg = ranged_aggregate(
+        axis="freq",
+        bands=bands,
+        operation=AggregationFunction.MEAN
+    )
+    pipeline = compose(f_spec, f_agg)
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+        axis_arr_out = pipeline(axis_arr_in)
+
+
+class BandPowerSettings(ez.Settings):
+    spectrogram_settings: SpectrogramSettings = field(default_factory=SpectrogramSettings)
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = (
+        field(default_factory=lambda: [(17, 30), (70, 170)]))
+
+
+class BandPower(GenAxisArray):
+    SETTINGS: BandPowerSettings
+
+    def construct_generator(self):
+        self.STATE.gen = bandpower(
+            spectrogram_settings=self.SETTINGS.spectrogram_settings,
+            bands=self.SETTINGS.bands
+        )

ezmsg/sigproc/butterworthfilter.py

@@ -1,18 +1,21 @@
+import typing
+
 import ezmsg.core as ez
 import scipy.signal
 import numpy as np
 
-from .filter import Filter, FilterState, FilterSettingsBase
+from .filter import filtergen, Filter, FilterState, FilterSettingsBase
 
-from typing import Optional, Tuple, Union
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
 
 
 class ButterworthFilterSettings(FilterSettingsBase):
     order: int = 0
-    cuton: Optional[float] = None  # Hz
-    cutoff: Optional[float] = None  # Hz
+    cuton: typing.Optional[float] = None  # Hz
+    cutoff: typing.Optional[float] = None  # Hz
 
-    def filter_specs(self) -> Optional[Tuple[str, Union[float, Tuple[float, float]]]]:
+    def filter_specs(self) -> typing.Optional[typing.Tuple[str, typing.Union[float, typing.Tuple[float, float]]]]:
         if self.cuton is None and self.cutoff is None:
             return None
         elif self.cuton is None and self.cutoff is not None:
@@ -26,6 +29,38 @@ 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]:
+    # IO
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = AxisArray(np.array([]), dims=[""])
+
+    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.
+
+    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
+            coefs = scipy.signal.butter(
+                order, Wn=cutoffs, btype=btype, fs=fs, output=coef_type
+            )
+            filter_gen = filtergen(axis, coefs, coef_type)
+
+        axis_arr_out = filter_gen.send(axis_arr_in)
+
+
 class ButterworthFilterState(FilterState):
     design: ButterworthFilterSettings
 
@@ -41,7 +76,7 @@ class ButterworthFilter(Filter):
         self.STATE.filt_designed = True
         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
@@ -57,6 +92,9 @@ class ButterworthFilter(Filter):
 
     @ez.subscriber(INPUT_FILTER)
     async def redesign(self, message: ButterworthFilterSettings) -> None:
+        if type(message) is not ButterworthFilterSettings:
+            return
+
         if self.STATE.design.order != message.order:
             self.STATE.zi = None
         self.STATE.design = message

ezmsg/sigproc/downsample.py

@@ -1,14 +1,46 @@
 from dataclasses import replace
+import traceback
+from typing import AsyncGenerator, Optional, Generator
 
-from ezmsg.util.messages.axisarray import AxisArray
+import numpy as np
 
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
 import ezmsg.core as ez
-import numpy as np
 
-from typing import (
-    AsyncGenerator,
-    Optional,
-)
+
+@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=[""])
+
+    # state variables
+    s_idx = 0
+
+    while True:
+        axis_arr_in = yield axis_arr_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
+
+        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)
+        else:
+            axis_arr_out = None
 
 
 class DownsampleSettings(ez.Settings):
@@ -18,7 +50,7 @@ class DownsampleSettings(ez.Settings):
 
 class DownsampleState(ez.State):
     cur_settings: DownsampleSettings
-    s_idx: int = 0
+    gen: Generator
 
 
 class Downsample(ez.Unit):
@@ -29,12 +61,17 @@ class Downsample(ez.Unit):
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
+    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)
@@ -42,22 +79,11 @@ class Downsample(ez.Unit):
         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)
-
-        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
-
-        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
+        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())

ezmsg/sigproc/ewmfilter.py

@@ -73,9 +73,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
@@ -108,7 +111,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,39 +1,107 @@
+import asyncio
+import typing
+
 from dataclasses import dataclass, replace, field
 
 import ezmsg.core as ez
 import scipy.signal
+
 import numpy as np
-import asyncio
+import numpy.typing as npt
 
 from ezmsg.util.messages.axisarray import AxisArray
-
-from typing import AsyncGenerator, Optional, Tuple
-
+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,...]]:
+    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]:
+    # 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=[""])
+
+    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
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+
+        if coefs is None:
+            # passthrough if we do not have a filter design.
+            axis_arr_out = axis_arr_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
+
+        # 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
+            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 :]
+            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)
+
 
 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):
@@ -44,7 +112,7 @@ class Filter(ez.Unit):
     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
@@ -84,7 +152,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)