ezmsg-sigproc 2.4.1__py3-none-any.whl → 2.6.0__py3-none-any.whl

ezmsg/sigproc/butterworthfilter.py

@@ -5,11 +5,11 @@ import scipy.signal
 from scipy.signal import normalize
 
 from .filter import (
-    FilterBaseSettings,
     BACoeffs,
-    SOSCoeffs,
-    FilterByDesignTransformer,
     BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+    SOSCoeffs,
 )
 
 
@@ -27,14 +27,14 @@ class ButterworthFilterSettings(FilterBaseSettings):
     """
     Cuton frequency (Hz). If `cutoff` is not specified then this is the highpass corner. Otherwise,
     if this is lower than `cutoff` then this is the beginning of the bandpass
-    or if this is greater than `cutoff` then this is the end of the bandstop. 
+    or if this is greater than `cutoff` then this is the end of the bandstop.
     """
 
     cutoff: float | None = None
     """
     Cutoff frequency (Hz). If `cuton` is not specified then this is the lowpass corner. Otherwise,
     if this is greater than `cuton` then this is the end of the bandpass,
-    or if this is less than `cuton` then this is the beginning of the bandstop. 
+    or if this is less than `cuton` then this is the beginning of the bandstop.
     """
 
     wn_hz: bool = True
@@ -96,9 +96,7 @@ def butter_design_fun(
     """
     coefs = None
     if order > 0:
-        btype, cutoffs = ButterworthFilterSettings(
-            order=order, cuton=cuton, cutoff=cutoff
-        ).filter_specs()
+        btype, cutoffs = ButterworthFilterSettings(order=order, cuton=cuton, cutoff=cutoff).filter_specs()
         coefs = scipy.signal.butter(
             order,
             Wn=cutoffs,
@@ -111,9 +109,7 @@ def butter_design_fun(
     return coefs
 
 
-class ButterworthFilterTransformer(
-    FilterByDesignTransformer[ButterworthFilterSettings, BACoeffs | SOSCoeffs]
-):
+class ButterworthFilterTransformer(FilterByDesignTransformer[ButterworthFilterSettings, BACoeffs | SOSCoeffs]):
     def get_design_function(
         self,
     ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
@@ -127,11 +123,7 @@ class ButterworthFilterTransformer(
         )
 
 
-class ButterworthFilter(
-    BaseFilterByDesignTransformerUnit[
-        ButterworthFilterSettings, ButterworthFilterTransformer
-    ]
-):
+class ButterworthFilter(BaseFilterByDesignTransformerUnit[ButterworthFilterSettings, ButterworthFilterTransformer]):
     SETTINGS = ButterworthFilterSettings
 
 

ezmsg/sigproc/butterworthzerophase.py

@@ -0,0 +1,123 @@
+import functools
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import scipy.signal
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from ezmsg.sigproc.base import SettingsType
+from ezmsg.sigproc.butterworthfilter import ButterworthFilterSettings, butter_design_fun
+from ezmsg.sigproc.filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+    FilterByDesignTransformer,
+    SOSCoeffs,
+)
+
+
+class ButterworthZeroPhaseSettings(ButterworthFilterSettings):
+    """Settings for :obj:`ButterworthZeroPhase`."""
+
+    # axis, coef_type, order, cuton, cutoff, wn_hz are inherited from ButterworthFilterSettings
+    padtype: str | None = None
+    """
+    Padding type to use in `scipy.signal.filtfilt`.
+    Must be one of {'odd', 'even', 'constant', None}.
+    Default is None for no padding.
+    """
+
+    padlen: int | None = 0
+    """
+    Length of the padding to use in `scipy.signal.filtfilt`.
+    If None, SciPy's default padding is used.
+    """
+
+
+class ButterworthZeroPhaseTransformer(FilterByDesignTransformer[ButterworthZeroPhaseSettings, BACoeffs | SOSCoeffs]):
+    """Zero-phase (filtfilt) Butterworth using your design function."""
+
+    def get_design_function(
+        self,
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
+        return functools.partial(
+            butter_design_fun,
+            order=self.settings.order,
+            cuton=self.settings.cuton,
+            cutoff=self.settings.cutoff,
+            coef_type=self.settings.coef_type,
+            wn_hz=self.settings.wn_hz,
+        )
+
+    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
+        self._coefs_cache = None
+        self._fs_cache = None
+        self.state.needs_redesign = True
+
+    def _reset_state(self, message: AxisArray) -> None:
+        self._coefs_cache = None
+        self._fs_cache = None
+        self.state.needs_redesign = True
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        ax_idx = message.get_axis_idx(axis)
+        fs = 1 / message.axes[axis].gain
+
+        if (
+            self._coefs_cache is None
+            or self.state.needs_redesign
+            or (self._fs_cache is None or not np.isclose(self._fs_cache, fs))
+        ):
+            self._coefs_cache = self.get_design_function()(fs)
+            self._fs_cache = fs
+            self.state.needs_redesign = False
+
+        if self._coefs_cache is None or self.settings.order <= 0 or message.data.size <= 0:
+            return message
+
+        x = message.data
+        if self.settings.coef_type == "sos":
+            y = scipy.signal.sosfiltfilt(
+                self._coefs_cache,
+                x,
+                axis=ax_idx,
+                padtype=self.settings.padtype,
+                padlen=self.settings.padlen,
+            )
+        elif self.settings.coef_type == "ba":
+            b, a = self._coefs_cache
+            y = scipy.signal.filtfilt(
+                b,
+                a,
+                x,
+                axis=ax_idx,
+                padtype=self.settings.padtype,
+                padlen=self.settings.padlen,
+            )
+        else:
+            ez.logger.error("coef_type must be 'sos' or 'ba'.")
+            raise ValueError("coef_type must be 'sos' or 'ba'.")
+
+        return replace(message, data=y)
+
+
+class ButterworthZeroPhase(
+    BaseFilterByDesignTransformerUnit[ButterworthZeroPhaseSettings, ButterworthZeroPhaseTransformer]
+):
+    SETTINGS = ButterworthZeroPhaseSettings

ezmsg/sigproc/cheby.py

@@ -5,11 +5,11 @@ import scipy.signal
 from scipy.signal import normalize
 
 from .filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
     FilterBaseSettings,
     FilterByDesignTransformer,
-    BACoeffs,
     SOSCoeffs,
-    BaseFilterByDesignTransformerUnit,
 )
 
 
@@ -104,9 +104,7 @@ def cheby_design_fun(
     return coefs
 
 
-class ChebyshevFilterTransformer(
-    FilterByDesignTransformer[ChebyshevFilterSettings, BACoeffs | SOSCoeffs]
-):
+class ChebyshevFilterTransformer(FilterByDesignTransformer[ChebyshevFilterSettings, BACoeffs | SOSCoeffs]):
     def get_design_function(
         self,
     ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
@@ -123,9 +121,5 @@ class ChebyshevFilterTransformer(
         )
 
 
-class ChebyshevFilter(
-    BaseFilterByDesignTransformerUnit[
-        ChebyshevFilterSettings, ChebyshevFilterTransformer
-    ]
-):
+class ChebyshevFilter(BaseFilterByDesignTransformerUnit[ChebyshevFilterSettings, ChebyshevFilterTransformer]):
     SETTINGS = ChebyshevFilterSettings

ezmsg/sigproc/combfilter.py

@@ -1,15 +1,16 @@
 import functools
 import typing
+
 import numpy as np
 import scipy.signal
 from scipy.signal import normalize
 
 from .filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
     FilterBaseSettings,
     FilterByDesignTransformer,
-    BACoeffs,
     SOSCoeffs,
-    BaseFilterByDesignTransformerUnit,
 )
 
 
@@ -103,9 +104,7 @@ def comb_design_fun(
     return combined_sos
 
 
-class CombFilterTransformer(
-    FilterByDesignTransformer[CombFilterSettings, BACoeffs | SOSCoeffs]
-):
+class CombFilterTransformer(FilterByDesignTransformer[CombFilterSettings, BACoeffs | SOSCoeffs]):
     def get_design_function(
         self,
     ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
@@ -120,9 +119,7 @@ class CombFilterTransformer(
         )
 
 
-class CombFilterUnit(
-    BaseFilterByDesignTransformerUnit[CombFilterSettings, CombFilterTransformer]
-):
+class CombFilterUnit(BaseFilterByDesignTransformerUnit[CombFilterSettings, CombFilterTransformer]):
     SETTINGS = CombFilterSettings
 
 

ezmsg/sigproc/decimate.py

@@ -4,7 +4,7 @@ import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 
 from .base import BaseTransformerUnit
-from .cheby import ChebyshevFilterTransformer, ChebyshevFilterSettings
+from .cheby import ChebyshevFilterSettings, ChebyshevFilterTransformer
 from .downsample import Downsample, DownsampleSettings
 from .filter import BACoeffs, SOSCoeffs
 
@@ -30,11 +30,7 @@ class ChebyForDecimateTransformer(ChebyshevFilterTransformer[BACoeffs | SOSCoeff
         return cheby_opt_design_fun
 
 
-class ChebyForDecimate(
-    BaseTransformerUnit[
-        ChebyshevFilterSettings, AxisArray, AxisArray, ChebyForDecimateTransformer
-    ]
-):
+class ChebyForDecimate(BaseTransformerUnit[ChebyshevFilterSettings, AxisArray, AxisArray, ChebyForDecimateTransformer]):
     SETTINGS = ChebyshevFilterSettings
 
 

ezmsg/sigproc/denormalize.py

@@ -1,13 +1,14 @@
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
 from ezmsg.sigproc.base import (
-    BaseTransformerUnit,
     BaseStatefulTransformer,
+    BaseTransformerUnit,
     processor_state,
 )
-from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.messages.util import replace
 
 
 class DenormalizeSettings(ez.Settings):
@@ -27,9 +28,7 @@ class DenormalizeState:
     offsets: npt.NDArray | None = None
 
 
-class DenormalizeTransformer(
-    BaseStatefulTransformer[DenormalizeSettings, AxisArray, AxisArray, DenormalizeState]
-):
+class DenormalizeTransformer(BaseStatefulTransformer[DenormalizeSettings, AxisArray, AxisArray, DenormalizeState]):
     """
     Scales data from a normalized distribution (mean=0, std=1) to a denormalized
     distribution using random per-channel offsets and gains designed to keep the
@@ -76,9 +75,5 @@ class DenormalizeTransformer(
         )
 
 
-class DenormalizeUnit(
-    BaseTransformerUnit[
-        DenormalizeSettings, AxisArray, AxisArray, DenormalizeTransformer
-    ]
-):
+class DenormalizeUnit(BaseTransformerUnit[DenormalizeSettings, AxisArray, AxisArray, DenormalizeTransformer]):
     SETTINGS = DenormalizeSettings

ezmsg/sigproc/detrend.py

@@ -1,7 +1,8 @@
 import scipy.signal as sps
 from ezmsg.util.messages.axisarray import AxisArray, replace
-from ezmsg.sigproc.ewma import EWMATransformer, EWMASettings
+
 from ezmsg.sigproc.base import BaseTransformerUnit
+from ezmsg.sigproc.ewma import EWMASettings, EWMATransformer
 
 
 class DetrendTransformer(EWMATransformer):
@@ -23,7 +24,5 @@ class DetrendTransformer(EWMATransformer):
         return replace(message, data=message.data - means)
 
 
-class DetrendUnit(
-    BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, DetrendTransformer]
-):
+class DetrendUnit(BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, DetrendTransformer]):
     SETTINGS = EWMASettings

ezmsg/sigproc/diff.py

@@ -1,13 +1,14 @@
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.util import replace
+
 from ezmsg.sigproc.base import (
+    BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
-    BaseStatefulTransformer,
 )
-from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
-from ezmsg.util.messages.util import replace
 
 
 class DiffSettings(ez.Settings):
@@ -21,9 +22,7 @@ class DiffState:
     last_time: float | None = None
 
 
-class DiffTransformer(
-    BaseStatefulTransformer[DiffSettings, AxisArray, AxisArray, DiffState]
-):
+class DiffTransformer(BaseStatefulTransformer[DiffSettings, AxisArray, AxisArray, DiffState]):
     def _hash_message(self, message: AxisArray) -> int:
         ax_idx = message.get_axis_idx(self.settings.axis)
         sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
@@ -49,20 +48,14 @@ class DiffTransformer(
             axis=ax_idx,
         )
         # Prepare last_dat for next iteration
-        self.state.last_dat = slice_along_axis(
-            message.data, slice(-1, None), axis=ax_idx
-        )
+        self.state.last_dat = slice_along_axis(message.data, slice(-1, None), axis=ax_idx)
         # Scale by fs if requested. This convers the diff to a derivative. e.g., diff of position becomes velocity.
         if self.settings.scale_by_fs:
             ax_info = message.get_axis(axis)
             if hasattr(ax_info, "data"):
                 dt = np.diff(np.concatenate(([self.state.last_time], ax_info.data)))
                 # Expand dt dims to match diffs
-                exp_sl = (
-                    (None,) * ax_idx
-                    + (Ellipsis,)
-                    + (None,) * (message.data.ndim - ax_idx - 1)
-                )
+                exp_sl = (None,) * ax_idx + (Ellipsis,) + (None,) * (message.data.ndim - ax_idx - 1)
                 diffs /= dt[exp_sl]
                 self.state.last_time = ax_info.data[-1]  # For next iteration
             else:
@@ -71,9 +64,7 @@ class DiffTransformer(
         return replace(message, data=diffs)
 
 
-class DiffUnit(
-    BaseTransformerUnit[DiffSettings, AxisArray, AxisArray, DiffTransformer]
-):
+class DiffUnit(BaseTransformerUnit[DiffSettings, AxisArray, AxisArray, DiffTransformer]):
     SETTINGS = DiffSettings
 
 

ezmsg/sigproc/downsample.py

@@ -1,10 +1,10 @@
+import ezmsg.core as ez
 import numpy as np
 from ezmsg.util.messages.axisarray import (
     AxisArray,
-    slice_along_axis,
     replace,
+    slice_along_axis,
 )
-import ezmsg.core as ez
 
 from .base import (
     BaseStatefulTransformer,
@@ -38,9 +38,7 @@ class DownsampleState:
     """Index of the next msg's first sample into the virtual rotating ds_factor counter."""
 
 
-class DownsampleTransformer(
-    BaseStatefulTransformer[DownsampleSettings, AxisArray, AxisArray, DownsampleState]
-):
+class DownsampleTransformer(BaseStatefulTransformer[DownsampleSettings, AxisArray, AxisArray, DownsampleState]):
     """
     Downsampled data simply comprise every `factor`th sample.
     This should only be used following appropriate lowpass filtering.
@@ -75,9 +73,7 @@ class DownsampleTransformer(
         axis_idx = message.get_axis_idx(axis)
 
         n_samples = message.data.shape[axis_idx]
-        samples = (
-            np.arange(self.state.s_idx, self.state.s_idx + n_samples) % self._state.q
-        )
+        samples = np.arange(self.state.s_idx, self.state.s_idx + n_samples) % self._state.q
         if n_samples > 0:
             # Update state for next iteration.
             self._state.s_idx = samples[-1] + 1
@@ -104,9 +100,7 @@ class DownsampleTransformer(
         return msg_out
 
 
-class Downsample(
-    BaseTransformerUnit[DownsampleSettings, AxisArray, AxisArray, DownsampleTransformer]
-):
+class Downsample(BaseTransformerUnit[DownsampleSettings, AxisArray, AxisArray, DownsampleTransformer]):
     SETTINGS = DownsampleSettings
 
 
@@ -115,6 +109,4 @@ def downsample(
     target_rate: float | None = None,
     factor: int | None = None,
 ) -> DownsampleTransformer:
-    return DownsampleTransformer(
-        DownsampleSettings(axis=axis, target_rate=target_rate, factor=factor)
-    )
+    return DownsampleTransformer(DownsampleSettings(axis=axis, target_rate=target_rate, factor=factor))

ezmsg/sigproc/ewma.py

@@ -1,14 +1,14 @@
-from dataclasses import field
 import functools
+from dataclasses import field
 
+import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
 import scipy.signal as sps
-import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
 from ezmsg.util.messages.util import replace
 
-from .base import BaseStatefulTransformer, processor_state, BaseTransformerUnit
+from .base import BaseStatefulTransformer, BaseTransformerUnit, processor_state
 
 
 def _tau_from_alpha(alpha: float, dt: float) -> float:
@@ -29,9 +29,7 @@ def _alpha_from_tau(tau: float, dt: float) -> float:
     return 1 - np.exp(-dt / tau)
 
 
-def ewma_step(
-    sample: npt.NDArray, zi: npt.NDArray, alpha: float, beta: float | None = None
-):
+def ewma_step(sample: npt.NDArray, zi: npt.NDArray, alpha: float, beta: float | None = None):
     """
     Do an exponentially weighted moving average step.
 
@@ -97,9 +95,7 @@ class EWMA_Deprecated:
         if self.prev is None:
             self.prev = arr[:1]
 
-        out += self.prev * np.expand_dims(
-            self.weights[1 : n + 1], list(range(1, arr.ndim))
-        )
+        out += self.prev * np.expand_dims(self.weights[1 : n + 1], list(range(1, arr.ndim)))
 
         self.prev = out[-1:]
 
@@ -128,9 +124,7 @@ class EWMA_Deprecated:
         if self.prev is None:
             self.prev = arr[:1]
 
-        result += self.prev * np.expand_dims(
-            self.weights[1 : n + 1], list(range(1, arr.ndim))
-        )
+        result += self.prev * np.expand_dims(self.weights[1 : n + 1], list(range(1, arr.ndim)))
 
         # Store the result back into prev
         self.prev = result[-1]
@@ -155,25 +149,17 @@ class EWMAState:
     zi: npt.NDArray | None = None
 
 
-class EWMATransformer(
-    BaseStatefulTransformer[EWMASettings, AxisArray, AxisArray, EWMAState]
-):
+class EWMATransformer(BaseStatefulTransformer[EWMASettings, AxisArray, AxisArray, EWMAState]):
     def _hash_message(self, message: AxisArray) -> int:
         axis = self.settings.axis or message.dims[0]
         axis_idx = message.get_axis_idx(axis)
-        sample_shape = (
-            message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
-        )
+        sample_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
         return hash((sample_shape, message.axes[axis].gain, message.key))
 
     def _reset_state(self, message: AxisArray) -> None:
         axis = self.settings.axis or message.dims[0]
-        self._state.alpha = _alpha_from_tau(
-            self.settings.time_constant, message.axes[axis].gain
-        )
-        sub_dat = slice_along_axis(
-            message.data, slice(None, 1, None), axis=message.get_axis_idx(axis)
-        )
+        self._state.alpha = _alpha_from_tau(self.settings.time_constant, message.axes[axis].gain)
+        sub_dat = slice_along_axis(message.data, slice(None, 1, None), axis=message.get_axis_idx(axis))
         self._state.zi = (1 - self._state.alpha) * sub_dat
 
     def _process(self, message: AxisArray) -> AxisArray:
@@ -191,7 +177,5 @@ class EWMATransformer(
         return replace(message, data=expected)
 
 
-class EWMAUnit(
-    BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, EWMATransformer]
-):
+class EWMAUnit(BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, EWMATransformer]):
     SETTINGS = EWMASettings

ezmsg/sigproc/ewmfilter.py

@@ -2,9 +2,9 @@ import asyncio
 import typing
 
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
-import numpy as np
 
 from .window import Window, WindowSettings
 

ezmsg/sigproc/extract_axis.py

@@ -1,6 +1,7 @@
-import numpy as np
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray, replace
+
 from ezmsg.sigproc.base import BaseTransformer, BaseTransformerUnit
 
 
@@ -35,7 +36,5 @@ class ExtractAxisData(BaseTransformer[ExtractAxisSettings, AxisArray, AxisArray]
             )
 
 
-class ExtractAxisDataUnit(
-    BaseTransformerUnit[ExtractAxisSettings, AxisArray, AxisArray, ExtractAxisData]
-):
+class ExtractAxisDataUnit(BaseTransformerUnit[ExtractAxisSettings, AxisArray, AxisArray, ExtractAxisData]):
     SETTINGS = ExtractAxisSettings