ezmsg-sigproc 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

ezmsg/sigproc/filter.py

@@ -1,14 +1,16 @@
-import asyncio
 from dataclasses import dataclass, field
 import typing
 
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 import numpy as np
 import numpy.typing as npt
 import scipy.signal
 
+from ezmsg.sigproc.base import GenAxisArray
+
 
 @dataclass
 class FilterCoefficients:
@@ -17,10 +19,8 @@ class FilterCoefficients:
 
 
 def _normalize_coefs(
-    coefs: typing.Union[
-        FilterCoefficients, typing.Tuple[npt.NDArray, npt.NDArray], npt.NDArray
-    ],
-) -> typing.Tuple[str, typing.Tuple[npt.NDArray, ...]]:
+    coefs: FilterCoefficients | tuple[npt.NDArray, npt.NDArray] | npt.NDArray,
+) -> tuple[str, tuple[npt.NDArray, ...]]:
     coef_type = "ba"
     if coefs is not None:
         # scipy.signal functions called with first arg `*coefs`.
@@ -35,7 +35,7 @@ def _normalize_coefs(
 
 @consumer
 def filtergen(
-    axis: str, coefs: typing.Optional[typing.Tuple[np.ndarray]], coef_type: str
+    axis: str, coefs: npt.NDArray | tuple[npt.NDArray] | None, coef_type: str
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Filter data using the provided coefficients.
@@ -61,7 +61,7 @@ def filtergen(
     zi_func = {"ba": scipy.signal.lfilter_zi, "sos": scipy.signal.sosfilt_zi}[coef_type]
 
     # State variables
-    zi: typing.Optional[npt.NDArray] = None
+    zi: npt.NDArray | None = None
 
     # Reset if these change.
     check_input = {"key": None, "shape": None}
@@ -105,128 +105,95 @@ def filtergen(
         msg_out = replace(msg_in, data=dat_out)
 
 
-class FilterSettingsBase(ez.Settings):
-    axis: typing.Optional[str] = None
-    fs: typing.Optional[float] = None
+# Type aliases
+BACoeffs = tuple[npt.NDArray, npt.NDArray]
+SOSCoeffs = npt.NDArray
+FilterCoefsMultiType = BACoeffs | SOSCoeffs
 
 
-class FilterSettings(FilterSettingsBase):
-    # If you'd like to statically design a filter, define it in settings
-    filt: typing.Optional[FilterCoefficients] = None
+@consumer
+def filter_gen_by_design(
+    axis: str,
+    coef_type: str,
+    design_fun: typing.Callable[[float], FilterCoefsMultiType | None],
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Filter data using a filter whose coefficients are calculated using the provided design function.
 
+    Args:
+        axis: The name of the axis to filter.
+            Note: The axis must be represented in the message .axes and be of type AxisArray.LinearAxis.
+        coef_type: "ba" or "sos"
+        design_fun: A callable that takes "fs" as its only argument and returns a tuple of filter coefficients.
+          If the design_fun returns None then the filter will act as a passthrough.
+          Hint: To make a design function that only requires fs, use functools.partial to set other parameters.
+          See butterworthfilter for an example.
 
-class FilterState(ez.State):
-    axis: typing.Optional[str] = None
-    zi: typing.Optional[np.ndarray] = None
-    filt_designed: bool = False
-    filt: typing.Optional[FilterCoefficients] = None
-    filt_set: asyncio.Event = field(default_factory=asyncio.Event)
-    samp_shape: typing.Optional[typing.Tuple[int, ...]] = None
-    fs: typing.Optional[float] = None  # Hz
+    Returns:
 
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
 
-class Filter(ez.Unit):
-    SETTINGS = FilterSettingsBase
-    STATE = FilterState
+    # State variables
+    # Initialize filtergen as passthrough until we receive a message that allows us to design the filter.
+    filter_gen = filtergen(axis, None, coef_type)
 
-    INPUT_FILTER = ez.InputStream(FilterCoefficients)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    # Reset if these change.
+    check_input = {"gain": None}
+    # No need to check parameters that don't affect the design; filter_gen should check most of its parameters.
 
-    def design_filter(self) -> typing.Optional[typing.Tuple[np.ndarray, np.ndarray]]:
-        raise NotImplementedError("Must implement 'design_filter' in Unit subclass!")
+    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"]
+        if b_reset:
+            check_input["gain"] = msg_in.axes[axis].gain
+            coefs = design_fun(1 / msg_in.axes[axis].gain)
+            filter_gen = filtergen(axis, coefs, coef_type)
 
-    # Set up filter with static initialization if specified
-    async def initialize(self) -> None:
-        if self.SETTINGS.axis is not None:
-            self.STATE.axis = self.SETTINGS.axis
+        msg_out = filter_gen.send(msg_in)
 
-        if isinstance(self.SETTINGS, FilterSettings):
-            if self.SETTINGS.filt is not None:
-                self.STATE.filt = self.SETTINGS.filt
-                self.STATE.filt_set.set()
-        else:
-            self.STATE.filt_set.clear()
 
-        if self.SETTINGS.fs is not None:
-            try:
-                self.update_filter()
-            except NotImplementedError:
-                ez.logger.debug("Using filter coefficients.")
+class FilterBaseSettings(ez.Settings):
+    axis: str | None = None
+    coef_type: str = "ba"
 
-    @ez.subscriber(INPUT_FILTER)
-    async def redesign(self, message: FilterCoefficients):
-        self.STATE.filt = message
-
-    def update_filter(self):
-        try:
-            coefs = self.design_filter()
-            self.STATE.filt = (
-                FilterCoefficients() if coefs is None else FilterCoefficients(*coefs)
-            )
-            self.STATE.filt_set.set()
-            self.STATE.filt_designed = True
-        except NotImplementedError as e:
-            raise e
-        except Exception as e:
-            ez.logger.warning(f"Error when designing filter: {e}")
-
-    @ez.subscriber(INPUT_SIGNAL)
-    @ez.publisher(OUTPUT_SIGNAL)
-    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)
-        fs = 1.0 / axis.gain
-
-        if self.STATE.fs != fs and self.STATE.filt_designed is True:
-            self.STATE.fs = fs
-            self.update_filter()
-
-        # Ensure filter is defined
-        # TODO: Maybe have me be a passthrough filter until coefficients are received
-        if self.STATE.filt is None:
-            self.STATE.filt_set.clear()
-            ez.logger.info("Awaiting filter coefficients...")
-            await self.STATE.filt_set.wait()
-            ez.logger.info("Filter coefficients received.")
-
-        assert self.STATE.filt is not None
-
-        arr_in = msg.data
-
-        # If the array is one dimensional, add a temporary second dimension so that the math works out
-        one_dimensional = False
-        if arr_in.ndim == 1:
-            arr_in = np.expand_dims(arr_in, axis=1)
-            one_dimensional = True
-
-        # We will perform filter with time dimension as last axis
-        arr_in = np.moveaxis(arr_in, axis_idx, -1)
-        samp_shape = arr_in[..., 0].shape
 
-        # Re-calculate/reset zi if necessary
-        if self.STATE.zi is None or samp_shape != self.STATE.samp_shape:
-            zi: np.ndarray = scipy.signal.lfilter_zi(
-                self.STATE.filt.b, self.STATE.filt.a
-            )
-            self.STATE.samp_shape = samp_shape
-            self.STATE.zi = np.array([zi] * np.prod(self.STATE.samp_shape))
-            self.STATE.zi = self.STATE.zi.reshape(
-                tuple(list(self.STATE.samp_shape) + [zi.shape[0]])
-            )
+class FilterBase(GenAxisArray):
+    SETTINGS = FilterBaseSettings
+
+    # Backwards-compatible with `Filter` unit
+    INPUT_FILTER = ez.InputStream(FilterCoefsMultiType)
+
+    def design_filter(
+        self,
+    ) -> typing.Callable[[float], FilterCoefsMultiType | None]:
+        raise NotImplementedError("Must implement 'design_filter' in Unit subclass!")
 
-        arr_out, self.STATE.zi = scipy.signal.lfilter(
-            self.STATE.filt.b, self.STATE.filt.a, arr_in, zi=self.STATE.zi
+    def construct_generator(self):
+        design_fun = self.design_filter()
+        self.STATE.gen = filter_gen_by_design(
+            self.SETTINGS.axis, self.SETTINGS.coef_type, design_fun
         )
 
-        arr_out = np.moveaxis(arr_out, -1, axis_idx)
+    @ez.subscriber(INPUT_FILTER)
+    async def redesign(self, message: FilterBaseSettings) -> None:
+        self.apply_settings(message)
+        self.construct_generator()
 
-        # Remove temporary first dimension if necessary
-        if one_dimensional:
-            arr_out = np.squeeze(arr_out, axis=1)
 
-        yield (
-            self.OUTPUT_SIGNAL,
-            replace(msg, data=arr_out),
-        )
+class FilterSettings(FilterBaseSettings):
+    # If you'd like to statically design a filter, define it in settings
+    coefs: FilterCoefficients | None = None
+    # Note: coef_type = "ba" is assumed for this class.
+
+
+class Filter(FilterBase):
+    SETTINGS = FilterSettings
+
+    INPUT_FILTER = ez.InputStream(FilterCoefficients)
+
+    def design_filter(self) -> typing.Callable[[float], BACoeffs | None]:
+        if self.SETTINGS.coefs is None:
+            return lambda fs: None
+        return lambda fs: (self.SETTINGS.coefs.b, self.SETTINGS.coefs.a)

ezmsg/sigproc/filterbank.py

@@ -8,7 +8,8 @@ import scipy.fft as sp_fft
 from scipy.special import lambertw
 import numpy.typing as npt
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
 from .base import GenAxisArray
@@ -35,7 +36,7 @@ class MinPhaseMode(OptionsEnum):
 
 @consumer
 def filterbank(
-    kernels: typing.Union[list[npt.NDArray], tuple[npt.NDArray, ...]],
+    kernels: list[npt.NDArray] | tuple[npt.NDArray, ...],
     mode: FilterbankMode = FilterbankMode.CONV,
     min_phase: MinPhaseMode = MinPhaseMode.NONE,
     axis: str = "time",
@@ -62,10 +63,10 @@ def filterbank(
      with the data payload containing the absolute value of the input :obj:`AxisArray` data.
 
     """
-    msg_out: typing.Optional[AxisArray] = None
+    msg_out: AxisArray | None = None
 
     # State variables
-    template: typing.Optional[AxisArray] = None
+    template: AxisArray | None = None
 
     # Reset if these change
     check_input = {
@@ -257,7 +258,7 @@ def filterbank(
 
 
 class FilterbankSettings(ez.Settings):
-    kernels: typing.Union[list[npt.NDArray], tuple[npt.NDArray, ...]]
+    kernels: list[npt.NDArray] | tuple[npt.NDArray, ...]
     mode: FilterbankMode = FilterbankMode.CONV
     min_phase: MinPhaseMode = MinPhaseMode.NONE
     axis: str = "time"

ezmsg/sigproc/math/abs.py

@@ -3,7 +3,8 @@ import typing
 import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.generator import consumer
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from ..base import GenAxisArray
 

ezmsg/sigproc/math/clip.py

@@ -3,7 +3,8 @@ import typing
 import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.generator import consumer
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from ..base import GenAxisArray
 

ezmsg/sigproc/math/difference.py

@@ -3,7 +3,8 @@ import typing
 import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.generator import consumer
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from ..base import GenAxisArray
 

ezmsg/sigproc/math/invert.py

@@ -3,7 +3,8 @@ import typing
 import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.generator import consumer
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from ..base import GenAxisArray
 

ezmsg/sigproc/math/log.py

@@ -3,7 +3,8 @@ import typing
 import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.generator import consumer
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from ..base import GenAxisArray
 

ezmsg/sigproc/math/scale.py

@@ -3,7 +3,8 @@ import typing
 import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.generator import consumer
-from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 from ..base import GenAxisArray
 

ezmsg/sigproc/messages.py

@@ -1,6 +1,5 @@
 import warnings
 import time
-import typing
 
 import numpy.typing as npt
 from ezmsg.util.messages.axisarray import AxisArray
@@ -20,7 +19,7 @@ def TSMessage(
     data: npt.NDArray,
     fs: float = 1.0,
     time_dim: int = 0,
-    timestamp: typing.Optional[float] = None,
+    timestamp: float | None = None,
 ) -> AxisArray:
     dims = [f"dim_{i}" for i in range(data.ndim)]
     dims[time_dim] = "time"

ezmsg/sigproc/sampler.py

@@ -10,8 +10,8 @@ import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import (
     AxisArray,
     slice_along_axis,
-    replace,
 )
+from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
 
@@ -20,7 +20,7 @@ class SampleTriggerMessage:
     timestamp: float = field(default_factory=time.time)
     """Time of the trigger, in seconds. The Clock depends on the input but defaults to time.time"""
 
-    period: typing.Optional[typing.Tuple[float, float]] = None
+    period: tuple[float, float] | None = None
     """The period around the timestamp, in seconds"""
 
     value: typing.Any = None
@@ -39,13 +39,11 @@ class SampleMessage:
 @consumer
 def sampler(
     buffer_dur: float,
-    axis: typing.Optional[str] = None,
-    period: typing.Optional[typing.Tuple[float, float]] = None,
+    axis: str | None = None,
+    period: tuple[float, float] | None = None,
     value: typing.Any = None,
     estimate_alignment: bool = True,
-) -> typing.Generator[
-    typing.List[SampleMessage], typing.Union[AxisArray, SampleTriggerMessage], None
-]:
+) -> typing.Generator[list[SampleMessage], AxisArray | SampleTriggerMessage, None]:
     """
     Sample data into a buffer, accept triggers, and return slices of sampled
     data around the trigger time.
@@ -74,7 +72,7 @@ def sampler(
 
     # State variables (most shared between trigger- and data-processing.
     triggers: deque[SampleTriggerMessage] = deque()
-    buffer: typing.Optional[npt.NDArray] = None
+    buffer: npt.NDArray | None = None
     n_samples: int = 0
     offset: float = 0.0
 
@@ -230,8 +228,8 @@ class SamplerSettings(ez.Settings):
     """
 
     buffer_dur: float
-    axis: typing.Optional[str] = None
-    period: typing.Optional[typing.Tuple[float, float]] = None
+    axis: str | None = None
+    period: tuple[float, float] | None = None
     """Optional default period if unspecified in SampleTriggerMessage"""
 
     value: typing.Any = None
@@ -248,9 +246,7 @@ class SamplerSettings(ez.Settings):
 
 class SamplerState(ez.State):
     cur_settings: SamplerSettings
-    gen: typing.Generator[
-        typing.Union[AxisArray, SampleTriggerMessage], typing.List[SampleMessage], None
-    ]
+    gen: typing.Generator[AxisArray | SampleTriggerMessage, list[SampleMessage], None]
 
 
 class Sampler(ez.Unit):
@@ -295,7 +291,7 @@ class Sampler(ez.Unit):
 
 
 class TriggerGeneratorSettings(ez.Settings):
-    period: typing.Tuple[float, float]
+    period: tuple[float, float]
     """The period around the trigger event."""
 
     prewait: float = 0.5