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

ezmsg/sigproc/filter.py

@@ -1,13 +1,21 @@
-from dataclasses import dataclass, replace, field
+import typing
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
 
 import ezmsg.core as ez
-import scipy.signal
 import numpy as np
-import asyncio
-
+import numpy.typing as npt
+import scipy.signal
+from ezmsg.baseproc import (
+    BaseConsumerUnit,
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    SettingsType,
+    TransformerType,
+    processor_state,
+)
 from ezmsg.util.messages.axisarray import AxisArray
-
-from typing import AsyncGenerator, Optional, Tuple
+from ezmsg.util.messages.util import replace
 
 
 @dataclass
@@ -16,125 +24,282 @@ class FilterCoefficients:
     a: np.ndarray = field(default_factory=lambda: np.array([1.0, 0.0]))
 
 
-class FilterSettingsBase(ez.Settings):
-    axis: Optional[str] = None
-    fs: Optional[float] = None
+# Type aliases
+BACoeffs = tuple[npt.NDArray, npt.NDArray]
+SOSCoeffs = npt.NDArray
+FilterCoefsType = typing.TypeVar("FilterCoefsType", BACoeffs, SOSCoeffs)
+
+
+def _normalize_coefs(
+    coefs: FilterCoefficients | tuple[npt.NDArray, npt.NDArray] | npt.NDArray | None,
+) -> tuple[str, tuple[npt.NDArray, ...] | None]:
+    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, np.ndarray):
+            coef_type = "sos"
+            coefs = (coefs,)  # sos funcs just want a single ndarray.
+        elif isinstance(coefs, FilterCoefficients):
+            coefs = (coefs.b, coefs.a)
+        elif not isinstance(coefs, tuple):
+            coefs = (coefs,)
+    return coef_type, coefs
+
+
+class FilterBaseSettings(ez.Settings):
+    axis: str | None = None
+    """The name of the axis to operate on."""
+
+    coef_type: str = "ba"
+    """The type of filter coefficients. One of "ba" or "sos"."""
+
+
+class FilterSettings(FilterBaseSettings):
+    coefs: FilterCoefficients | None = None
+    """The pre-calculated filter coefficients."""
+
+    # Note: coef_type = "ba" is assumed for this class.
+
+
+@processor_state
+class FilterState:
+    zi: npt.NDArray | None = None
+
+
+class FilterTransformer(BaseStatefulTransformer[FilterSettings, AxisArray, AxisArray, FilterState]):
+    """
+    Filter data using the provided coefficients.
+    """
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        if self.settings.coefs is None:
+            return message
+        if self._state.zi is None:
+            self._reset_state(message)
+            self._hash = self._hash_message(message)
+        return super().__call__(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        axis_idx = message.get_axis_idx(axis)
+        samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        return hash((message.key, samp_shape))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        axis_idx = message.get_axis_idx(axis)
+        n_tail = message.data.ndim - axis_idx - 1
+        _, coefs = _normalize_coefs(self.settings.coefs)
+
+        if self.settings.coef_type == "ba":
+            b, a = coefs
+            if len(a) == 1 or np.allclose(a[1:], 0):
+                # For FIR filters, use lfiltic with zero initial conditions
+                zi = scipy.signal.lfiltic(b, a, [])
+            else:
+                # For IIR filters...
+                zi = scipy.signal.lfilter_zi(b, a)
+        else:
+            # For second-order sections (SOS) filters, use sosfilt_zi
+            zi = scipy.signal.sosfilt_zi(*coefs)
+
+        zi_expand = (None,) * axis_idx + (slice(None),) + (None,) * n_tail
+        n_tile = message.data.shape[:axis_idx] + (1,) + message.data.shape[axis_idx + 1 :]
+
+        if self.settings.coef_type == "sos":
+            zi_expand = (slice(None),) + zi_expand
+            n_tile = (1,) + n_tile
+
+        self.state.zi = np.tile(zi[zi_expand], n_tile)
+
+    def update_coefficients(
+        self,
+        coefs: FilterCoefficients | tuple[npt.NDArray, npt.NDArray] | npt.NDArray,
+        coef_type: str | None = None,
+    ) -> None:
+        """
+        Update filter coefficients.
+
+        If the new coefficients have the same length as the current ones, only the coefficients are updated.
+        If the lengths differ, the filter state is also reset to handle the new filter order.
+
+        Args:
+            coefs: New filter coefficients
+        """
+        old_coefs = self.settings.coefs
+
+        # Update settings with new coefficients
+        self.settings = replace(self.settings, coefs=coefs)
+        if coef_type is not None:
+            self.settings = replace(self.settings, coef_type=coef_type)
+
+        # Check if we need to reset the state
+        if self.state.zi is not None:
+            reset_needed = False
+
+            if self.settings.coef_type == "ba":
+                if isinstance(old_coefs, FilterCoefficients) and isinstance(coefs, FilterCoefficients):
+                    if len(old_coefs.b) != len(coefs.b) or len(old_coefs.a) != len(coefs.a):
+                        reset_needed = True
+                elif isinstance(old_coefs, tuple) and isinstance(coefs, tuple):
+                    if len(old_coefs[0]) != len(coefs[0]) or len(old_coefs[1]) != len(coefs[1]):
+                        reset_needed = True
+                else:
+                    reset_needed = True
+            elif self.settings.coef_type == "sos":
+                if isinstance(old_coefs, np.ndarray) and isinstance(coefs, np.ndarray):
+                    if old_coefs.shape != coefs.shape:
+                        reset_needed = True
+                else:
+                    reset_needed = True
+
+            if reset_needed:
+                self.state.zi = None  # This will trigger _reset_state on the next call
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        if message.data.size > 0:
+            axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+            axis_idx = message.get_axis_idx(axis)
+            _, coefs = _normalize_coefs(self.settings.coefs)
+            filt_func = {"ba": scipy.signal.lfilter, "sos": scipy.signal.sosfilt}[self.settings.coef_type]
+            dat_out, self.state.zi = filt_func(*coefs, message.data, axis=axis_idx, zi=self.state.zi)
+        else:
+            dat_out = message.data
+
+        return replace(message, data=dat_out)
 
 
-class FilterSettings(FilterSettingsBase):
-    # If you'd like to statically design a filter, define it in settings
-    filt: Optional[FilterCoefficients] = None
+class Filter(BaseTransformerUnit[FilterSettings, AxisArray, AxisArray, FilterTransformer]):
+    SETTINGS = FilterSettings
 
 
-class FilterState(ez.State):
-    axis: Optional[str] = None
-    zi: Optional[np.ndarray] = None
-    filt_designed: bool = False
-    filt: Optional[FilterCoefficients] = None
-    filt_set: asyncio.Event = field(default_factory=asyncio.Event)
-    samp_shape: Optional[Tuple[int, ...]] = None
-    fs: Optional[float] = None  # Hz
+def filtergen(axis: str, coefs: npt.NDArray | tuple[npt.NDArray] | None, coef_type: str) -> FilterTransformer:
+    """
+    Filter data using the provided coefficients.
 
+    Returns:
+        :obj:`FilterTransformer`.
+    """
+    return FilterTransformer(FilterSettings(axis=axis, coefs=coefs, coef_type=coef_type))
 
-class Filter(ez.Unit):
-    SETTINGS: FilterSettingsBase
-    STATE: FilterState
 
-    INPUT_FILTER = ez.InputStream(FilterCoefficients)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+@processor_state
+class FilterByDesignState:
+    filter: FilterTransformer | None = None
+    needs_redesign: bool = False
 
-    def design_filter(self) -> Optional[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:
-        if self.SETTINGS.axis is not None:
-            self.STATE.axis = self.SETTINGS.axis
+class FilterByDesignTransformer(
+    BaseStatefulTransformer[SettingsType, AxisArray, AxisArray, FilterByDesignState],
+    ABC,
+    typing.Generic[SettingsType, FilterCoefsType],
+):
+    """Abstract base class for filter design transformers."""
 
-        if isinstance(self.SETTINGS, FilterSettings):
-            if self.SETTINGS.filt is not None:
-                self.STATE.filt = self.SETTINGS.filt
-                self.STATE.filt_set.set()
+    @classmethod
+    def get_message_type(cls, dir: str) -> type[AxisArray]:
+        if dir in ("in", "out"):
+            return AxisArray
+        else:
+            raise ValueError(f"Invalid direction: {dir}. Must be 'in' or 'out'.")
+
+    @abstractmethod
+    def get_design_function(self) -> typing.Callable[[float], FilterCoefsType | None]:
+        """Return a function that takes sampling frequency and returns filter coefficients."""
+        ...
+
+    def update_settings(self, new_settings: typing.Optional[SettingsType] = None, **kwargs) -> None:
+        """
+        Update settings and mark that filter coefficients need to be recalculated.
+
+        Args:
+            new_settings: Complete new settings object to replace current settings
+            **kwargs: Individual settings to update
+        """
+        # Update settings
+        if new_settings is not None:
+            self.settings = new_settings
+        else:
+            self.settings = replace(self.settings, **kwargs)
+
+        # Set flag to trigger recalculation on next message
+        if self.state.filter is not None:
+            self.state.needs_redesign = True
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        # Offer a shortcut when there is no design function or order is 0.
+        if hasattr(self.settings, "order") and not self.settings.order:
+            return message
+        design_fun = self.get_design_function()
+        if design_fun is None:
+            return message
+
+        # Check if filter exists but needs redesign due to settings change
+        if self.state.filter is not None and self.state.needs_redesign:
+            axis = self.state.filter.settings.axis
+            fs = 1 / message.axes[axis].gain
+            coefs = design_fun(fs)
+
+            # Convert BA to SOS if requested
+            if coefs is not None and self.settings.coef_type == "sos":
+                if isinstance(coefs, tuple) and len(coefs) == 2:
+                    # It's BA format, convert to SOS
+                    b, a = coefs
+                    coefs = scipy.signal.tf2sos(b, a)
+
+            self.state.filter.update_coefficients(coefs, coef_type=self.settings.coef_type)
+            self.state.needs_redesign = False
+
+        return super().__call__(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        gain = message.axes[axis].gain if hasattr(message.axes[axis], "gain") else 1
+        axis_idx = message.get_axis_idx(axis)
+        samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        return hash((message.key, samp_shape, gain))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        design_fun = self.get_design_function()
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        fs = 1 / message.axes[axis].gain
+        coefs = design_fun(fs)
+
+        # Convert BA to SOS if requested
+        if coefs is not None and self.settings.coef_type == "sos":
+            if isinstance(coefs, tuple) and len(coefs) == 2:
+                # It's BA format, convert to SOS
+                b, a = coefs
+                coefs = scipy.signal.tf2sos(b, a)
+
+        new_settings = FilterSettings(axis=axis, coef_type=self.settings.coef_type, coefs=coefs)
+        self.state.filter = FilterTransformer(settings=new_settings)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        return self.state.filter(message)
+
+
+class BaseFilterByDesignTransformerUnit(
+    BaseTransformerUnit[SettingsType, AxisArray, AxisArray, FilterByDesignTransformer],
+    typing.Generic[SettingsType, TransformerType],
+):
+    @ez.subscriber(BaseConsumerUnit.INPUT_SETTINGS)
+    async def on_settings(self, msg: SettingsType) -> None:
+        """
+        Receive a settings message, override self.SETTINGS, and re-create the processor.
+        Child classes that wish to have fine-grained control over whether the
+        core processor resets on settings changes should override this method.
+
+        Args:
+            msg: a settings message.
+        """
+        self.apply_settings(msg)
+
+        # Check if processor exists yet
+        if hasattr(self, "processor") and self.processor is not None:
+            # Update the existing processor with new settings
+            self.processor.update_settings(self.SETTINGS)
         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.")
-
-    @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) -> 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]])
-            )
-
-        arr_out, self.STATE.zi = scipy.signal.lfilter(
-            self.STATE.filt.b, self.STATE.filt.a, arr_in, zi=self.STATE.zi
-        )
-
-        arr_out = np.moveaxis(arr_out, -1, axis_idx)
-
-        # 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),
+            # Processor doesn't exist yet, create a new one
+            self.create_processor()