ezmsg-sigproc 2.7.0__tar.gz → 2.8.0__tar.gz

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 2.7.0
+Version: 2.8.0
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/docs/source/guides/explanations/sigproc.rst

@@ -329,7 +329,7 @@ Often, all that is required is the following (e.g., for a custom transformer):
 
   import ezmsg.core as ez
   from ezmsg.util.messages.axisarray import AxisArray
-  from ezmsg.sigproc.base import BaseTransformer, BaseTransformerUnit
+  from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 
 
   class CustomTransformerSettings(ez.Settings):

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/docs/source/guides/tutorials/signalprocessing.rst

@@ -132,7 +132,7 @@ Add the following import statements to the top of the `downsample.py` file:
     )
     import ezmsg.core as ez
 
-    from ezmsg.sigproc.base import (
+    from ezmsg.baseproc import (
         BaseStatefulTransformer,
         BaseTransformerUnit,
         processor_state,

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/__version__.py

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

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/affinetransform.py

@@ -1,3 +1,13 @@
+"""Affine transformations via matrix multiplication: y = Ax or y = Ax + B.
+
+For full matrix transformations where channels are mixed (off-diagonal weights),
+use :obj:`AffineTransformTransformer` or the `AffineTransform` unit.
+
+For simple per-channel scaling and offset (diagonal weights only), use
+:obj:`LinearTransformTransformer` from :mod:`ezmsg.sigproc.linear` instead,
+which is more efficient as it avoids matrix multiplication.
+"""
+
 import os
 from pathlib import Path
 
@@ -17,7 +27,6 @@ from ezmsg.util.messages.util import replace
 class AffineTransformSettings(ez.Settings):
     """
     Settings for :obj:`AffineTransform`.
-    See :obj:`affine_transform` for argument details.
     """
 
     weights: np.ndarray | str | Path
@@ -39,6 +48,19 @@ class AffineTransformState:
 class AffineTransformTransformer(
     BaseStatefulTransformer[AffineTransformSettings, AxisArray, AxisArray, AffineTransformState]
 ):
+    """Apply affine transformation via matrix multiplication: y = Ax or y = Ax + B.
+
+    Use this transformer when you need full matrix transformations that mix
+    channels (off-diagonal weights), such as spatial filters or projections.
+
+    For simple per-channel scaling and offset where each output channel depends
+    only on its corresponding input channel (diagonal weight matrix), use
+    :obj:`LinearTransformTransformer` instead, which is more efficient.
+
+    The weights matrix can include an offset row (stacked as [A|B]) where the
+    input is automatically augmented with a column of ones to compute y = Ax + B.
+    """
+
     def __call__(self, message: AxisArray) -> AxisArray:
         # Override __call__ so we can shortcut if weights are None.
         if self.settings.weights is None or (

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/butterworthzerophase.py

@@ -4,10 +4,10 @@ import typing
 import ezmsg.core as ez
 import numpy as np
 import scipy.signal
+from ezmsg.baseproc import SettingsType
 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,

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/denormalize.py

@@ -1,14 +1,13 @@
 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 (
+from ezmsg.baseproc import (
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 
 class DenormalizeSettings(ez.Settings):

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/detrend.py

@@ -1,7 +1,7 @@
 import scipy.signal as sps
+from ezmsg.baseproc import BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray, replace
 
-from ezmsg.sigproc.base import BaseTransformerUnit
 from ezmsg.sigproc.ewma import EWMASettings, EWMATransformer
 
 

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/diff.py

@@ -1,14 +1,13 @@
 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 (
+from ezmsg.baseproc import (
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.util import replace
 
 
 class DiffSettings(ez.Settings):

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/ewma.py

@@ -139,8 +139,15 @@ class EWMA_Deprecated:
 
 class EWMASettings(ez.Settings):
     time_constant: float = 1.0
+    """The amount of time for the smoothed response of a unit step function to reach 1 - 1/e approx-eq 63.2%."""
+
     axis: str | None = None
 
+    accumulate: bool = True
+    """If True, update the EWMA state with each sample. If False, only apply
+    the current EWMA estimate without updating state (useful for inference
+    periods where you don't want to adapt statistics)."""
+
 
 @processor_state
 class EWMAState:
@@ -166,15 +173,45 @@ class EWMATransformer(BaseStatefulTransformer[EWMASettings, AxisArray, AxisArray
             return message
         axis = self.settings.axis or message.dims[0]
         axis_idx = message.get_axis_idx(axis)
-        expected, self._state.zi = sps.lfilter(
-            [self._state.alpha],
-            [1.0, self._state.alpha - 1.0],
-            message.data,
-            axis=axis_idx,
-            zi=self._state.zi,
-        )
+        if self.settings.accumulate:
+            # Normal behavior: update state with new samples
+            expected, self._state.zi = sps.lfilter(
+                [self._state.alpha],
+                [1.0, self._state.alpha - 1.0],
+                message.data,
+                axis=axis_idx,
+                zi=self._state.zi,
+            )
+        else:
+            # Process-only: compute output without updating state
+            expected, _ = sps.lfilter(
+                [self._state.alpha],
+                [1.0, self._state.alpha - 1.0],
+                message.data,
+                axis=axis_idx,
+                zi=self._state.zi,
+            )
         return replace(message, data=expected)
 
 
 class EWMAUnit(BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, EWMATransformer]):
     SETTINGS = EWMASettings
+
+    @ez.subscriber(BaseTransformerUnit.INPUT_SETTINGS)
+    async def on_settings(self, msg: EWMASettings) -> None:
+        """
+        Handle settings updates with smart reset behavior.
+
+        Only resets state if `axis` changes (structural change).
+        Changes to `time_constant` or `accumulate` are applied without
+        resetting accumulated state.
+        """
+        old_axis = self.SETTINGS.axis
+        self.apply_settings(msg)
+
+        if msg.axis != old_axis:
+            # Axis changed - need full reset
+            self.create_processor()
+        else:
+            # Only accumulate or time_constant changed - keep state
+            self.processor.settings = msg

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/extract_axis.py

@@ -1,9 +1,8 @@
 import ezmsg.core as ez
 import numpy as np
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray, replace
 
-from ezmsg.sigproc.base import BaseTransformer, BaseTransformerUnit
-
 
 class ExtractAxisSettings(ez.Settings):
     axis: str = "freq"

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/filter.py

@@ -6,10 +6,7 @@ import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
 import scipy.signal
-from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.messages.util import replace
-
-from ezmsg.sigproc.base import (
+from ezmsg.baseproc import (
     BaseConsumerUnit,
     BaseStatefulTransformer,
     BaseTransformerUnit,
@@ -17,6 +14,8 @@ from ezmsg.sigproc.base import (
     TransformerType,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 
 
 @dataclass

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/fir_hilbert.py

@@ -4,10 +4,10 @@ import typing
 import ezmsg.core as ez
 import numpy as np
 import scipy.signal as sps
+from ezmsg.baseproc import BaseStatefulTransformer, processor_state
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ezmsg.sigproc.base import BaseStatefulTransformer, processor_state
 from ezmsg.sigproc.filter import (
     BACoeffs,
     BaseFilterByDesignTransformerUnit,

ezmsg_sigproc-2.8.0/src/ezmsg/sigproc/linear.py

@@ -0,0 +1,118 @@
+"""Apply a linear transformation: output = scale * input + offset.
+
+Supports per-element scale and offset along a specified axis.
+Uses Array API for compatibility with numpy, cupy, torch, etc.
+
+For full matrix transformations, use :obj:`AffineTransformTransformer` instead.
+"""
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from array_api_compat import get_namespace
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .base import BaseStatefulTransformer, BaseTransformerUnit, processor_state
+
+
+class LinearTransformSettings(ez.Settings):
+    scale: float | list[float] | npt.ArrayLike = 1.0
+    """Scale factor(s). Can be a scalar (applied to all elements) or an array
+    matching the size of the specified axis for per-element scaling."""
+
+    offset: float | list[float] | npt.ArrayLike = 0.0
+    """Offset value(s). Can be a scalar (applied to all elements) or an array
+    matching the size of the specified axis for per-element offset."""
+
+    axis: str | None = None
+    """Axis along which to apply per-element scale/offset. If None, scalar
+    scale/offset are broadcast to all elements."""
+
+
+@processor_state
+class LinearTransformState:
+    scale: npt.NDArray = None
+    """Prepared scale array for broadcasting."""
+
+    offset: npt.NDArray = None
+    """Prepared offset array for broadcasting."""
+
+
+class LinearTransformTransformer(
+    BaseStatefulTransformer[LinearTransformSettings, AxisArray, AxisArray, LinearTransformState]
+):
+    """Apply linear transformation: output = scale * input + offset.
+
+    This transformer is optimized for element-wise linear operations with
+    optional per-channel (or per-axis) coefficients. For full matrix
+    transformations, use :obj:`AffineTransformTransformer` instead.
+
+    Examples:
+        # Uniform scaling and offset
+        >>> transformer = LinearTransformTransformer(LinearTransformSettings(scale=2.0, offset=1.0))
+
+        # Per-channel scaling (e.g., for 3-channel data along "ch" axis)
+        >>> transformer = LinearTransformTransformer(LinearTransformSettings(
+        ...     scale=[0.5, 1.0, 2.0],
+        ...     offset=[0.0, 0.1, 0.2],
+        ...     axis="ch"
+        ... ))
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        """Hash based on shape and axis to detect when broadcast shapes need recalculation."""
+        axis = self.settings.axis
+        if axis is not None:
+            axis_idx = message.get_axis_idx(axis)
+            return hash((message.data.ndim, axis_idx, message.data.shape[axis_idx]))
+        return hash(message.data.ndim)
+
+    def _reset_state(self, message: AxisArray) -> None:
+        """Prepare scale/offset arrays with proper broadcast shapes."""
+        xp = get_namespace(message.data)
+        ndim = message.data.ndim
+
+        scale = self.settings.scale
+        offset = self.settings.offset
+
+        # Convert settings to arrays
+        if isinstance(scale, (list, np.ndarray)):
+            scale = xp.asarray(scale, dtype=xp.float64)
+        else:
+            # Scalar: create a 0-d array
+            scale = xp.asarray(float(scale), dtype=xp.float64)
+
+        if isinstance(offset, (list, np.ndarray)):
+            offset = xp.asarray(offset, dtype=xp.float64)
+        else:
+            # Scalar: create a 0-d array
+            offset = xp.asarray(float(offset), dtype=xp.float64)
+
+        # If axis is specified and we have 1-d arrays, reshape for proper broadcasting
+        if self.settings.axis is not None and ndim > 0:
+            axis_idx = message.get_axis_idx(self.settings.axis)
+
+            if scale.ndim == 1:
+                # Create shape for broadcasting: all 1s except at axis_idx
+                broadcast_shape = [1] * ndim
+                broadcast_shape[axis_idx] = scale.shape[0]
+                scale = xp.reshape(scale, broadcast_shape)
+
+            if offset.ndim == 1:
+                broadcast_shape = [1] * ndim
+                broadcast_shape[axis_idx] = offset.shape[0]
+                offset = xp.reshape(offset, broadcast_shape)
+
+        self._state.scale = scale
+        self._state.offset = offset
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        result = message.data * self._state.scale + self._state.offset
+        return replace(message, data=result)
+
+
+class LinearTransform(BaseTransformerUnit[LinearTransformSettings, AxisArray, AxisArray, LinearTransformTransformer]):
+    """Unit wrapper for LinearTransformTransformer."""
+
+    SETTINGS = LinearTransformSettings

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/rollingscaler.py

@@ -3,14 +3,14 @@ from collections import deque
 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 (
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
 from ezmsg.sigproc.sampler import SampleMessage
 
 

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/src/ezmsg/sigproc/scaler.py

@@ -1,5 +1,6 @@
 import typing
 
+import ezmsg.core as ez
 import numpy as np
 from ezmsg.baseproc import (
     BaseStatefulTransformer,
@@ -82,11 +83,38 @@ class AdaptiveStandardScalerTransformer(
     ]
 ):
     def _reset_state(self, message: AxisArray) -> None:
-        self._state.samps_ewma = EWMATransformer(time_constant=self.settings.time_constant, axis=self.settings.axis)
-        self._state.vars_sq_ewma = EWMATransformer(time_constant=self.settings.time_constant, axis=self.settings.axis)
+        self._state.samps_ewma = EWMATransformer(
+            time_constant=self.settings.time_constant,
+            axis=self.settings.axis,
+            accumulate=self.settings.accumulate,
+        )
+        self._state.vars_sq_ewma = EWMATransformer(
+            time_constant=self.settings.time_constant,
+            axis=self.settings.axis,
+            accumulate=self.settings.accumulate,
+        )
+
+    @property
+    def accumulate(self) -> bool:
+        """Whether to accumulate statistics from incoming samples."""
+        return self.settings.accumulate
+
+    @accumulate.setter
+    def accumulate(self, value: bool) -> None:
+        """
+        Set the accumulate mode and propagate to child EWMA transformers.
+
+        Args:
+            value: If True, update statistics with each sample.
+                   If False, only apply current statistics without updating.
+        """
+        if self._state.samps_ewma is not None:
+            self._state.samps_ewma.settings = replace(self._state.samps_ewma.settings, accumulate=value)
+        if self._state.vars_sq_ewma is not None:
+            self._state.vars_sq_ewma.settings = replace(self._state.vars_sq_ewma.settings, accumulate=value)
 
     def _process(self, message: AxisArray) -> AxisArray:
-        # Update step
+        # Update step (respects accumulate setting via child EWMAs)
         mean_message = self._state.samps_ewma(message)
         var_sq_message = self._state.vars_sq_ewma(replace(message, data=message.data**2))
 
@@ -108,6 +136,27 @@ class AdaptiveStandardScaler(
 ):
     SETTINGS = AdaptiveStandardScalerSettings
 
+    @ez.subscriber(BaseTransformerUnit.INPUT_SETTINGS)
+    async def on_settings(self, msg: AdaptiveStandardScalerSettings) -> None:
+        """
+        Handle settings updates with smart reset behavior.
+
+        Only resets state if `axis` changes (structural change).
+        Changes to `time_constant` or `accumulate` are applied without
+        resetting accumulated statistics.
+        """
+        old_axis = self.SETTINGS.axis
+        self.apply_settings(msg)
+
+        if msg.axis != old_axis:
+            # Axis changed - need full reset
+            self.create_processor()
+        else:
+            # Update accumulate on processor (propagates to child EWMAs)
+            self.processor.accumulate = msg.accumulate
+            # Also update own settings reference
+            self.processor.settings = msg
+
 
 # Backwards compatibility...
 def scaler_np(time_constant: float = 1.0, axis: str | None = None) -> AdaptiveStandardScalerTransformer:

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/tests/unit/test_activation.py

@@ -4,7 +4,12 @@ import scipy.special
 from ezmsg.util.messages.axisarray import AxisArray
 from frozendict import frozendict
 
-from ezmsg.sigproc.activation import ACTIVATIONS, ActivationFunction, activation
+from ezmsg.sigproc.activation import (
+    ACTIVATIONS,
+    ActivationFunction,
+    ActivationSettings,
+    ActivationTransformer,
+)
 
 
 @pytest.mark.parametrize("function", [_ for _ in ActivationFunction] + ActivationFunction.options())
@@ -25,8 +30,8 @@ def test_activation(function: str):
             )
             yield msg
 
-    proc = activation(function=function)
-    out_msgs = [proc.send(_) for _ in msg_generator()]
+    proc = ActivationTransformer(ActivationSettings(function=function))
+    out_msgs = [proc(_) for _ in msg_generator()]
     out_dat = AxisArray.concatenate(*out_msgs, dim="time").data
 
     if function in ACTIVATIONS:

{ezmsg_sigproc-2.7.0 → ezmsg_sigproc-2.8.0}/tests/unit/test_affine_transform.py

@@ -4,11 +4,16 @@ from pathlib import Path
 import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray
 
-from ezmsg.sigproc.affinetransform import affine_transform, common_rereference
+from ezmsg.sigproc.affinetransform import (
+    AffineTransformSettings,
+    AffineTransformTransformer,
+    CommonRereferenceSettings,
+    CommonRereferenceTransformer,
+)
 from tests.helpers.util import assert_messages_equal
 
 
-def test_affine_generator():
+def test_affine_transform():
     n_times = 13
     n_chans = 64
     in_dat = np.arange(n_times * n_chans).reshape(n_times, n_chans)
@@ -20,16 +25,16 @@ def test_affine_generator():
 
     backup = [copy.deepcopy(msg_in)]
 
-    gen = affine_transform(weights=np.eye(n_chans), axis="ch")
-    msg_out = gen.send(msg_in)
+    xformer = AffineTransformTransformer(AffineTransformSettings(weights=np.eye(n_chans), axis="ch"))
+    msg_out = xformer(msg_in)
     assert msg_out.data.shape == in_dat.shape
     assert np.allclose(msg_out.data, in_dat)
     assert not np.may_share_memory(msg_out.data, in_dat)
 
     assert_messages_equal([msg_in], backup)
 
-    # Send again just to make sure the generator doesn't crash
-    _ = gen.send(msg_in)
+    # Call again just to make sure the transformer doesn't crash
+    _ = xformer(msg_in)
 
     # Test with weights from a CSV file.
     csv_path = Path(__file__).parents[1] / "resources" / "xform.csv"
@@ -37,27 +42,29 @@ def test_affine_generator():
     expected_out = in_dat @ weights.T
     # Same result: expected_out = np.vstack([(step[None, :] * weights).sum(axis=1) for step in in_dat])
 
-    gen = affine_transform(weights=csv_path, axis="ch", right_multiply=False)
-    msg_out = gen.send(msg_in)
+    xformer = AffineTransformTransformer(AffineTransformSettings(weights=csv_path, axis="ch", right_multiply=False))
+    msg_out = xformer(msg_in)
     assert np.allclose(msg_out.data, expected_out)
     assert len(msg_out.axes["ch"].data) == weights.shape[0]
     assert (msg_out.axes["ch"].data[:-1] == msg_in.axes["ch"].data).all()
 
     # Try again as str, not Path
-    gen = affine_transform(weights=str(csv_path), axis="ch", right_multiply=False)
-    msg_out = gen.send(msg_in)
+    xformer = AffineTransformTransformer(
+        AffineTransformSettings(weights=str(csv_path), axis="ch", right_multiply=False)
+    )
+    msg_out = xformer(msg_in)
     assert np.allclose(msg_out.data, expected_out)
     assert len(msg_out.axes["ch"].data) == weights.shape[0]
 
     # Try again as direct ndarray
-    gen = affine_transform(weights=weights, axis="ch", right_multiply=False)
-    msg_out = gen.send(msg_in)
+    xformer = AffineTransformTransformer(AffineTransformSettings(weights=weights, axis="ch", right_multiply=False))
+    msg_out = xformer(msg_in)
     assert np.allclose(msg_out.data, expected_out)
     assert len(msg_out.axes["ch"].data) == weights.shape[0]
 
     # One more time, but we pre-transpose the weights and do not override right_multiply
-    gen = affine_transform(weights=weights.T, axis="ch", right_multiply=True)
-    msg_out = gen.send(msg_in)
+    xformer = AffineTransformTransformer(AffineTransformSettings(weights=weights.T, axis="ch", right_multiply=True))
+    msg_out = xformer(msg_in)
     assert np.allclose(msg_out.data, expected_out)
     assert len(msg_out.axes["ch"].data) == weights.shape[0]
 
@@ -70,9 +77,9 @@ def test_affine_passthrough():
 
     backup = [copy.deepcopy(msg_in)]
 
-    gen = affine_transform(weights="passthrough", axis="does not matter")
-    msg_out = gen.send(msg_in)
-    # We wouldn't want out_data is in_dat ezmsg pipeline but it's fine for the generator
+    xformer = AffineTransformTransformer(AffineTransformSettings(weights="passthrough", axis="does not matter"))
+    msg_out = xformer(msg_in)
+    # We wouldn't want out_data is in_dat ezmsg pipeline but it's fine for the transformer
     assert msg_out.data is in_dat
     assert_messages_equal([msg_out], backup)
 
@@ -85,8 +92,8 @@ def test_common_rereference():
 
     backup = [copy.deepcopy(msg_in)]
 
-    gen = common_rereference(mode="mean", axis="ch", include_current=True)
-    msg_out = gen.send(msg_in)
+    xformer = CommonRereferenceTransformer(CommonRereferenceSettings(mode="mean", axis="ch", include_current=True))
+    msg_out = xformer(msg_in)
     assert np.array_equal(
         msg_out.data,
         msg_in.data - np.mean(msg_in.data, axis=1, keepdims=True),
@@ -103,17 +110,17 @@ def test_common_rereference():
         expected_out.append(msg_in.data[..., ch_ix] - np.mean(msg_in.data[..., idx], axis=1))
     expected_out = np.stack(expected_out).T
 
-    gen = common_rereference(mode="mean", axis="ch", include_current=False)
-    msg_out = gen.send(msg_in)  # 41 us
+    xformer = CommonRereferenceTransformer(CommonRereferenceSettings(mode="mean", axis="ch", include_current=False))
+    msg_out = xformer(msg_in)  # 41 us
     assert np.allclose(msg_out.data, expected_out)
 
-    # Instead of CAR, we could use affine_transform with weights that reproduce CAR.
+    # Instead of CAR, we could use AffineTransformTransformer with weights that reproduce CAR.
     # However, this method is 30x slower than above. (Actual difference varies depending on data shape).
     if False:
         weights = -np.ones((n_chans, n_chans)) / (n_chans - 1)
         np.fill_diagonal(weights, 1)
-        gen = affine_transform(weights=weights, axis="ch")
-        msg_out = gen.send(msg_in)
+        xformer = AffineTransformTransformer(AffineTransformSettings(weights=weights, axis="ch"))
+        msg_out = xformer(msg_in)
         assert np.allclose(msg_out.data, expected_out)
 
 
@@ -123,7 +130,7 @@ def test_car_passthrough():
     in_dat = np.arange(n_times * n_chans).reshape(n_times, n_chans)
     msg_in = AxisArray(in_dat, dims=["time", "ch"])
 
-    gen = common_rereference(mode="passthrough")
-    msg_out = gen.send(msg_in)
+    xformer = CommonRereferenceTransformer(CommonRereferenceSettings(mode="passthrough"))
+    msg_out = xformer(msg_in)
     assert np.array_equal(msg_out.data, in_dat)
     assert np.may_share_memory(msg_out.data, in_dat)