ezmsg-sigproc 2.5.0__py3-none-any.whl → 2.7.0__py3-none-any.whl

ezmsg/sigproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.5.0'
-__version_tuple__ = version_tuple = (2, 5, 0)
+__version__ = version = '2.7.0'
+__version_tuple__ = version_tuple = (2, 7, 0)
 
 __commit_id__ = commit_id = None

ezmsg/sigproc/activation.py

@@ -1,10 +1,10 @@
-import scipy.special
 import ezmsg.core as ez
+import scipy.special
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
 from .spectral import OptionsEnum
-from .base import BaseTransformer, BaseTransformerUnit
 
 
 class ActivationFunction(OptionsEnum):
@@ -50,20 +50,14 @@ class ActivationTransformer(BaseTransformer[ActivationSettings, AxisArray, AxisA
             # str type handling
             function = self.settings.function.lower()
             if function not in ActivationFunction.options():
-                raise ValueError(
-                    f"Unrecognized activation function {function}. Must be one of {ACTIVATIONS.keys()}"
-                )
-            function = list(ACTIVATIONS.keys())[
-                ActivationFunction.options().index(function)
-            ]
+                raise ValueError(f"Unrecognized activation function {function}. Must be one of {ACTIVATIONS.keys()}")
+            function = list(ACTIVATIONS.keys())[ActivationFunction.options().index(function)]
             func = ACTIVATIONS[function]
 
         return replace(message, data=func(message.data))
 
 
-class Activation(
-    BaseTransformerUnit[ActivationSettings, AxisArray, AxisArray, ActivationTransformer]
-):
+class Activation(BaseTransformerUnit[ActivationSettings, AxisArray, AxisArray, ActivationTransformer]):
     SETTINGS = ActivationSettings
 
 

ezmsg/sigproc/adaptive_lattice_notch.py

@@ -1,12 +1,11 @@
+import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
 import scipy.signal
-import ezmsg.core as ez
+from ezmsg.baseproc import BaseStatefulTransformer, processor_state
 from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis
 from ezmsg.util.messages.util import replace
 
-from .base import processor_state, BaseStatefulTransformer
-
 
 class AdaptiveLatticeNotchFilterSettings(ez.Settings):
     """Settings for the Adaptive Lattice Notch Filter."""
@@ -76,9 +75,7 @@ class AdaptiveLatticeNotchFilterTransformer(
 
         fs = 1 / message.axes[self.settings.axis].gain
         init_f = (
-            self.settings.init_notch_freq
-            if self.settings.init_notch_freq is not None
-            else 0.07178314656435313 * fs
+            self.settings.init_notch_freq if self.settings.init_notch_freq is not None else 0.07178314656435313 * fs
         )
         init_omega = init_f * (2 * np.pi) / fs
         init_k1 = -np.cos(init_omega)
@@ -91,9 +88,7 @@ class AdaptiveLatticeNotchFilterTransformer(
         self._state.k1 = init_k1 + np.zeros(sample_shape, dtype=float)
         self._state.freq_template = CoordinateAxis(
             data=np.zeros((0,) + sample_shape, dtype=float),
-            dims=[self.settings.axis]
-            + message.dims[:ax_idx]
-            + message.dims[ax_idx + 1 :],
+            dims=[self.settings.axis] + message.dims[:ax_idx] + message.dims[ax_idx + 1 :],
             unit="Hz",
         )
 
@@ -147,9 +142,7 @@ class AdaptiveLatticeNotchFilterTransformer(
             for ix, k in enumerate(self._state.k1.flatten()):
                 # Filter to get s_n (notch filter state)
                 a_s = [1, k * gamma_plus_1, gamma]
-                s_n[:, ix], self._state.zi[:, ix] = scipy.signal.lfilter(
-                    [1], a_s, _x[:, ix], zi=self._state.zi[:, ix]
-                )
+                s_n[:, ix], self._state.zi[:, ix] = scipy.signal.lfilter([1], a_s, _x[:, ix], zi=self._state.zi[:, ix])
 
                 # Apply output filter to get y_out
                 b_y = [1, 2 * k, 1]
@@ -159,17 +152,11 @@ class AdaptiveLatticeNotchFilterTransformer(
             s_n_reshaped = s_n.reshape((s_n.shape[0],) + x_data.shape[1:])
             s_final = s_n_reshaped[-1]  # Current s_n
             s_final_1 = s_n_reshaped[-2]  # s_n_1
-            s_final_2 = (
-                s_n_reshaped[-3] if len(s_n_reshaped) > 2 else self._state.s_history[0]
-            )  # s_n_2
+            s_final_2 = s_n_reshaped[-3] if len(s_n_reshaped) > 2 else self._state.s_history[0]  # s_n_2
 
             # Update p and q using final values
-            self._state.p = eta * self._state.p + one_minus_eta * (
-                s_final_1 * (s_final + s_final_2)
-            )
-            self._state.q = eta * self._state.q + one_minus_eta * (
-                2 * (s_final_1 * s_final_1)
-            )
+            self._state.p = eta * self._state.p + one_minus_eta * (s_final_1 * (s_final + s_final_2))
+            self._state.q = eta * self._state.q + one_minus_eta * (2 * (s_final_1 * s_final_1))
 
             # Update reflection coefficient
             new_k1 = -self._state.p / (self._state.q + 1e-8)  # Avoid division by zero
@@ -199,17 +186,11 @@ class AdaptiveLatticeNotchFilterTransformer(
                 y_out[sample_ix] = s_n + 2 * self._state.k1 * s_n_1 + s_n_2
 
                 # Update filter parameters
-                self._state.p = eta * self._state.p + one_minus_eta * (
-                    s_n_1 * (s_n + s_n_2)
-                )
-                self._state.q = eta * self._state.q + one_minus_eta * (
-                    2 * (s_n_1 * s_n_1)
-                )
+                self._state.p = eta * self._state.p + one_minus_eta * (s_n_1 * (s_n + s_n_2))
+                self._state.q = eta * self._state.q + one_minus_eta * (2 * (s_n_1 * s_n_1))
 
                 # Update reflection coefficient
-                new_k1 = -self._state.p / (
-                    self._state.q + 1e-8
-                )  # Avoid division by zero
+                new_k1 = -self._state.p / (self._state.q + 1e-8)  # Avoid division by zero
                 new_k1 = np.clip(new_k1, -1, 1)  # Clip to prevent instability
                 self._state.k1 = mu * self._state.k1 + one_minus_mu * new_k1  # Smoothed
 

ezmsg/sigproc/affinetransform.py

@@ -1,18 +1,17 @@
 import os
 from pathlib import Path
 
+import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
-import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, AxisBase
-from ezmsg.util.messages.util import replace
-
-from .base import (
+from ezmsg.baseproc import (
     BaseStatefulTransformer,
-    BaseTransformerUnit,
     BaseTransformer,
+    BaseTransformerUnit,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import AxisArray, AxisBase
+from ezmsg.util.messages.util import replace
 
 
 class AffineTransformSettings(ez.Settings):
@@ -38,15 +37,12 @@ class AffineTransformState:
 
 
 class AffineTransformTransformer(
-    BaseStatefulTransformer[
-        AffineTransformSettings, AxisArray, AxisArray, AffineTransformState
-    ]
+    BaseStatefulTransformer[AffineTransformSettings, AxisArray, AxisArray, AffineTransformState]
 ):
     def __call__(self, message: AxisArray) -> AxisArray:
         # Override __call__ so we can shortcut if weights are None.
         if self.settings.weights is None or (
-            isinstance(self.settings.weights, str)
-            and self.settings.weights == "passthrough"
+            isinstance(self.settings.weights, str) and self.settings.weights == "passthrough"
         ):
             return message
         return super().__call__(message)
@@ -68,18 +64,12 @@ class AffineTransformTransformer(
         self._state.weights = weights
 
         axis = self.settings.axis or message.dims[-1]
-        if (
-            axis in message.axes
-            and hasattr(message.axes[axis], "data")
-            and weights.shape[0] != weights.shape[1]
-        ):
+        if axis in message.axes and hasattr(message.axes[axis], "data") and weights.shape[0] != weights.shape[1]:
             in_labels = message.axes[axis].data
             new_labels = []
             n_in, n_out = weights.shape
             if len(in_labels) != n_in:
-                ez.logger.warning(
-                    f"Received {len(in_labels)} for {n_in} inputs. Check upstream labels."
-                )
+                ez.logger.warning(f"Received {len(in_labels)} for {n_in} inputs. Check upstream labels.")
             else:
                 b_filled_outputs = np.any(weights, axis=0)
                 b_used_inputs = np.any(weights, axis=1)
@@ -97,9 +87,7 @@ class AffineTransformTransformer(
                 elif np.all(b_filled_outputs):
                     new_labels = np.array(in_labels)[b_used_inputs]
 
-            self._state.new_axis = replace(
-                message.axes[axis], data=np.array(new_labels)
-            )
+            self._state.new_axis = replace(message.axes[axis], data=np.array(new_labels))
 
     def _process(self, message: AxisArray) -> AxisArray:
         axis = self.settings.axis or message.dims[-1]
@@ -110,9 +98,7 @@ class AffineTransformTransformer(
             # 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
-            )
+            data = np.concatenate((data, np.ones(sample_shape).astype(data.dtype)), axis=axis_idx)
 
         if axis_idx in [-1, len(message.dims) - 1]:
             data = np.matmul(data, self._state.weights)
@@ -128,11 +114,7 @@ class AffineTransformTransformer(
         return replace(message, **replace_kwargs)
 
 
-class AffineTransform(
-    BaseTransformerUnit[
-        AffineTransformSettings, AxisArray, AxisArray, AffineTransformTransformer
-    ]
-):
+class AffineTransform(BaseTransformerUnit[AffineTransformSettings, AxisArray, AxisArray, AffineTransformTransformer]):
     SETTINGS = AffineTransformSettings
 
 
@@ -153,9 +135,7 @@ def affine_transform(
         :obj:`AffineTransformTransformer`.
     """
     return AffineTransformTransformer(
-        AffineTransformSettings(
-            weights=weights, axis=axis, right_multiply=right_multiply
-        )
+        AffineTransformSettings(weights=weights, axis=axis, right_multiply=right_multiply)
     )
 
 
@@ -178,9 +158,7 @@ class CommonRereferenceSettings(ez.Settings):
     """Set False to exclude each channel from participating in the calculation of its reference."""
 
 
-class CommonRereferenceTransformer(
-    BaseTransformer[CommonRereferenceSettings, AxisArray, AxisArray]
-):
+class CommonRereferenceTransformer(BaseTransformer[CommonRereferenceSettings, AxisArray, AxisArray]):
     def _process(self, message: AxisArray) -> AxisArray:
         if self.settings.mode == "passthrough":
             return message
@@ -188,9 +166,7 @@ class CommonRereferenceTransformer(
         axis = self.settings.axis or message.dims[-1]
         axis_idx = message.get_axis_idx(axis)
 
-        func = {"mean": np.mean, "median": np.median, "passthrough": zeros_for_noop}[
-            self.settings.mode
-        ]
+        func = {"mean": np.mean, "median": np.median, "passthrough": zeros_for_noop}[self.settings.mode]
 
         ref_data = func(message.data, axis=axis_idx, keepdims=True)
 
@@ -213,9 +189,7 @@ class CommonRereferenceTransformer(
 
 
 class CommonRereference(
-    BaseTransformerUnit[
-        CommonRereferenceSettings, AxisArray, AxisArray, CommonRereferenceTransformer
-    ]
+    BaseTransformerUnit[CommonRereferenceSettings, AxisArray, AxisArray, CommonRereferenceTransformer]
 ):
     SETTINGS = CommonRereferenceSettings
 

ezmsg/sigproc/aggregate.py

@@ -1,23 +1,23 @@
-from array_api_compat import get_namespace
 import typing
 
+import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
-import ezmsg.core as ez
+from array_api_compat import get_namespace
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 from ezmsg.util.messages.axisarray import (
     AxisArray,
-    slice_along_axis,
     AxisBase,
     replace,
+    slice_along_axis,
 )
 
 from .spectral import OptionsEnum
-from .base import (
-    BaseTransformer,
-    BaseStatefulTransformer,
-    BaseTransformerUnit,
-    processor_state,
-)
 
 
 class AggregationFunction(OptionsEnum):
@@ -89,9 +89,7 @@ class RangedAggregateState:
 
 
 class RangedAggregateTransformer(
-    BaseStatefulTransformer[
-        RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateState
-    ]
+    BaseStatefulTransformer[RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateState]
 ):
     def __call__(self, message: AxisArray) -> AxisArray:
         # Override for shortcut passthrough mode.
@@ -118,16 +116,12 @@ class RangedAggregateTransformer(
         if hasattr(target_axis, "data"):
             self._state.ax_vec = target_axis.data
         else:
-            self._state.ax_vec = target_axis.value(
-                np.arange(message.data.shape[ax_idx])
-            )
+            self._state.ax_vec = target_axis.value(np.arange(message.data.shape[ax_idx]))
 
         ax_dat = []
         slices = []
         for start, stop in self.settings.bands:
-            inds = np.where(
-                np.logical_and(self._state.ax_vec >= start, self._state.ax_vec <= stop)
-            )[0]
+            inds = np.where(np.logical_and(self._state.ax_vec >= start, self._state.ax_vec <= stop))[0]
             slices.append(np.s_[inds[0] : inds[-1] + 1])
             if hasattr(target_axis, "data"):
                 if self._state.ax_vec.dtype.type is np.str_:
@@ -164,8 +158,7 @@ class RangedAggregateTransformer(
             ]
         else:
             out_data = [
-                agg_func(slice_along_axis(message.data, sl, axis=ax_idx), axis=ax_idx)
-                for sl in self._state.slices
+                agg_func(slice_along_axis(message.data, sl, axis=ax_idx), axis=ax_idx) for sl in self._state.slices
             ]
 
         msg_out = replace(
@@ -187,11 +180,7 @@ class RangedAggregateTransformer(
         return msg_out
 
 
-class RangedAggregate(
-    BaseTransformerUnit[
-        RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateTransformer
-    ]
-):
+class RangedAggregate(BaseTransformerUnit[RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateTransformer]):
     SETTINGS = RangedAggregateSettings
 
 
@@ -212,9 +201,7 @@ def ranged_aggregate(
     Returns:
         :obj:`RangedAggregateTransformer`
     """
-    return RangedAggregateTransformer(
-        RangedAggregateSettings(axis=axis, bands=bands, operation=operation)
-    )
+    return RangedAggregateTransformer(RangedAggregateSettings(axis=axis, bands=bands, operation=operation))
 
 
 class AggregateSettings(ez.Settings):
@@ -242,9 +229,7 @@ class AggregateTransformer(BaseTransformer[AggregateSettings, AxisArray, AxisArr
         op = self.settings.operation
 
         if op == AggregationFunction.NONE:
-            raise ValueError(
-                "AggregationFunction.NONE is not supported for full-axis aggregation"
-            )
+            raise ValueError("AggregationFunction.NONE is not supported for full-axis aggregation")
 
         if op == AggregationFunction.TRAPEZOID:
             # Trapezoid integration requires x-coordinates
@@ -276,9 +261,7 @@ class AggregateTransformer(BaseTransformer[AggregateSettings, AxisArray, AxisArr
         )
 
 
-class AggregateUnit(
-    BaseTransformerUnit[AggregateSettings, AxisArray, AxisArray, AggregateTransformer]
-):
+class AggregateUnit(BaseTransformerUnit[AggregateSettings, AxisArray, AxisArray, AggregateTransformer]):
     """Unit that aggregates an entire axis using a specified operation."""
 
     SETTINGS = AggregateSettings

ezmsg/sigproc/bandpower.py

@@ -1,20 +1,20 @@
 from dataclasses import field
 
 import ezmsg.core as ez
+from ezmsg.baseproc import (
+    BaseProcessor,
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+    CompositeProcessor,
+)
 from ezmsg.util.messages.axisarray import AxisArray
 
-from .spectrogram import SpectrogramSettings, SpectrogramTransformer
 from .aggregate import (
     AggregationFunction,
-    RangedAggregateTransformer,
     RangedAggregateSettings,
+    RangedAggregateTransformer,
 )
-from .base import (
-    BaseProcessor,
-    CompositeProcessor,
-    BaseStatefulProcessor,
-    BaseTransformerUnit,
-)
+from .spectrogram import SpectrogramSettings, SpectrogramTransformer
 
 
 class BandPowerSettings(ez.Settings):
@@ -22,16 +22,12 @@ class BandPowerSettings(ez.Settings):
     Settings for ``BandPower``.
     """
 
-    spectrogram_settings: SpectrogramSettings = field(
-        default_factory=SpectrogramSettings
-    )
+    spectrogram_settings: SpectrogramSettings = field(default_factory=SpectrogramSettings)
     """
     Settings for spectrogram calculation.
     """
 
-    bands: list[tuple[float, float]] | None = field(
-        default_factory=lambda: [(17, 30), (70, 170)]
-    )
+    bands: list[tuple[float, float]] | None = field(default_factory=lambda: [(17, 30), (70, 170)])
     """
     (min, max) tuples of band limits in Hz.
     """
@@ -46,9 +42,7 @@ class BandPowerTransformer(CompositeProcessor[BandPowerSettings, AxisArray, Axis
         settings: BandPowerSettings,
     ) -> dict[str, BaseProcessor | BaseStatefulProcessor]:
         return {
-            "spectrogram": SpectrogramTransformer(
-                settings=settings.spectrogram_settings
-            ),
+            "spectrogram": SpectrogramTransformer(settings=settings.spectrogram_settings),
             "aggregate": RangedAggregateTransformer(
                 settings=RangedAggregateSettings(
                     axis="freq",
@@ -59,9 +53,7 @@ class BandPowerTransformer(CompositeProcessor[BandPowerSettings, AxisArray, Axis
         }
 
 
-class BandPower(
-    BaseTransformerUnit[BandPowerSettings, AxisArray, AxisArray, BandPowerTransformer]
-):
+class BandPower(BaseTransformerUnit[BandPowerSettings, AxisArray, AxisArray, BandPowerTransformer]):
     SETTINGS = BandPowerSettings