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

ezmsg/sigproc/coordinatespaces.py

@@ -0,0 +1,159 @@
+"""
+Coordinate space transformations for streaming data.
+
+This module provides utilities and ezmsg nodes for transforming between
+Cartesian (x, y) and polar (r, theta) coordinate systems.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+from enum import Enum
+from typing import Tuple
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from array_api_compat import get_namespace, is_array_api_obj
+from ezmsg.baseproc import (
+    BaseTransformer,
+    BaseTransformerUnit,
+)
+from ezmsg.util.messages.axisarray import AxisArray, replace
+
+# -- Utility functions for coordinate transformations --
+
+
+def _get_namespace_or_numpy(*args: npt.ArrayLike):
+    """Get array namespace if any arg is an array, otherwise return numpy."""
+    for arg in args:
+        if is_array_api_obj(arg):
+            return get_namespace(arg)
+    return np
+
+
+def polar2z(r: npt.ArrayLike, theta: npt.ArrayLike) -> npt.ArrayLike:
+    """Convert polar coordinates to complex number representation."""
+    xp = _get_namespace_or_numpy(r, theta)
+    return r * xp.exp(1j * theta)
+
+
+def z2polar(z: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
+    """Convert complex number to polar coordinates (r, theta)."""
+    xp = _get_namespace_or_numpy(z)
+    return xp.abs(z), xp.atan2(xp.imag(z), xp.real(z))
+
+
+def cart2z(x: npt.ArrayLike, y: npt.ArrayLike) -> npt.ArrayLike:
+    """Convert Cartesian coordinates to complex number representation."""
+    return x + 1j * y
+
+
+def z2cart(z: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
+    """Convert complex number to Cartesian coordinates (x, y)."""
+    xp = _get_namespace_or_numpy(z)
+    return xp.real(z), xp.imag(z)
+
+
+def cart2pol(x: npt.ArrayLike, y: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
+    """Convert Cartesian coordinates (x, y) to polar coordinates (r, theta)."""
+    return z2polar(cart2z(x, y))
+
+
+def pol2cart(r: npt.ArrayLike, theta: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
+    """Convert polar coordinates (r, theta) to Cartesian coordinates (x, y)."""
+    return z2cart(polar2z(r, theta))
+
+
+# -- ezmsg transformer classes --
+
+
+class CoordinateMode(str, Enum):
+    """Transformation mode for coordinate conversion."""
+
+    CART2POL = "cart2pol"
+    """Convert Cartesian (x, y) to polar (r, theta)."""
+
+    POL2CART = "pol2cart"
+    """Convert polar (r, theta) to Cartesian (x, y)."""
+
+
+class CoordinateSpacesSettings(ez.Settings):
+    """
+    Settings for :obj:`CoordinateSpaces`.
+
+    See :obj:`coordinate_spaces` for argument details.
+    """
+
+    mode: CoordinateMode = CoordinateMode.CART2POL
+    """The transformation mode: 'cart2pol' or 'pol2cart'."""
+
+    axis: str | None = None
+    """
+    The name of the axis containing the coordinate components.
+    Defaults to the last axis. Must have exactly 2 elements (x,y or r,theta).
+    """
+
+
+class CoordinateSpacesTransformer(BaseTransformer[CoordinateSpacesSettings, AxisArray, AxisArray]):
+    """
+    Transform between Cartesian and polar coordinate systems.
+
+    The input must have exactly 2 elements along the specified axis:
+    - For cart2pol: expects (x, y), outputs (r, theta)
+    - For pol2cart: expects (r, theta), outputs (x, y)
+    """
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+
+        if message.data.shape[axis_idx] != 2:
+            raise ValueError(
+                f"Coordinate transformation requires exactly 2 elements along axis '{axis}', "
+                f"got {message.data.shape[axis_idx]}."
+            )
+
+        # Extract components along the specified axis
+        slices_a = [slice(None)] * message.data.ndim
+        slices_b = [slice(None)] * message.data.ndim
+        slices_a[axis_idx] = 0
+        slices_b[axis_idx] = 1
+
+        component_a = message.data[tuple(slices_a)]
+        component_b = message.data[tuple(slices_b)]
+
+        if self.settings.mode == CoordinateMode.CART2POL:
+            # Input: x, y -> Output: r, theta
+            out_a, out_b = cart2pol(component_a, component_b)
+        else:
+            # Input: r, theta -> Output: x, y
+            out_a, out_b = pol2cart(component_a, component_b)
+
+        # Stack results back along the same axis
+        result = xp.stack([out_a, out_b], axis=axis_idx)
+
+        # Update axis labels if present (use numpy for string labels)
+        axes = message.axes
+        if axis in axes and hasattr(axes[axis], "data"):
+            if self.settings.mode == CoordinateMode.CART2POL:
+                new_labels = np.array(["r", "theta"])
+            else:
+                new_labels = np.array(["x", "y"])
+            axes = {**axes, axis: replace(axes[axis], data=new_labels)}
+
+        return replace(message, data=result, axes=axes)
+
+
+class CoordinateSpaces(
+    BaseTransformerUnit[CoordinateSpacesSettings, AxisArray, AxisArray, CoordinateSpacesTransformer]
+):
+    """
+    Unit for transforming between Cartesian and polar coordinate systems.
+
+    See :obj:`CoordinateSpacesSettings` for configuration options.
+    """
+
+    SETTINGS = CoordinateSpacesSettings

ezmsg/sigproc/decimate.py

@@ -1,33 +1,39 @@
 import typing
 
 import ezmsg.core as ez
+from ezmsg.baseproc import BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 
-from .cheby import ChebyshevFilter, ChebyshevFilterSettings
+from .cheby import ChebyshevFilterSettings, ChebyshevFilterTransformer
 from .downsample import Downsample, DownsampleSettings
-from .filter import FilterCoefsMultiType
+from .filter import BACoeffs, SOSCoeffs
 
 
-class ChebyForDecimate(ChebyshevFilter):
+class ChebyForDecimateTransformer(ChebyshevFilterTransformer[BACoeffs | SOSCoeffs]):
     """
-    A :obj:`ChebyshevFilter` node with a design filter method that additionally accepts a target sampling rate,
+    A :obj:`ChebyshevFilterTransformer` with a design filter method that additionally accepts a target sampling rate,
      and if the target rate cannot be achieved it returns None, else it returns the filter coefficients.
     """
 
-    def design_filter(
+    def get_design_function(
         self,
-    ) -> typing.Callable[[float], FilterCoefsMultiType | None]:
-        def cheby_opt_design_fun(fs: float) -> FilterCoefsMultiType | None:
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
+        def cheby_opt_design_fun(fs: float) -> BACoeffs | SOSCoeffs | None:
             if fs is None:
                 return None
-            ds_factor = int(fs / (2.5 * self.SETTINGS.Wn))
+            ds_factor = int(fs / (2.5 * self.settings.Wn))
             if ds_factor < 2:
                 return None
-            partial_fun = super(ChebyForDecimate, self).design_filter()
+            partial_fun = super(ChebyForDecimateTransformer, self).get_design_function()
             return partial_fun(fs)
+
         return cheby_opt_design_fun
 
 
+class ChebyForDecimate(BaseTransformerUnit[ChebyshevFilterSettings, AxisArray, AxisArray, ChebyForDecimateTransformer]):
+    SETTINGS = ChebyshevFilterSettings
+
+
 class Decimate(ez.Collection):
     """
     A :obj:`Collection` chaining a :obj:`Filter` node configured as a lowpass Chebyshev filter
@@ -43,7 +49,6 @@ class Decimate(ez.Collection):
     DOWNSAMPLE = Downsample()
 
     def configure(self) -> None:
-
         cheby_settings = ChebyshevFilterSettings(
             order=8,
             ripple_tol=0.05,

ezmsg/sigproc/denormalize.py

@@ -0,0 +1,78 @@
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+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):
+    low_rate: float = 2.0
+    """Low end of probable rate after denormalization (Hz)."""
+
+    high_rate: float = 40.0
+    """High end of probable rate after denormalization (Hz)."""
+
+    distribution: str = "uniform"
+    """Distribution to sample rates from. Options are 'uniform', 'normal', or 'constant'."""
+
+
+@processor_state
+class DenormalizeState:
+    gains: npt.NDArray | None = None
+    offsets: npt.NDArray | None = None
+
+
+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
+    99.9% CIs between 0 and 2x the offset.
+
+    This is useful for simulating realistic firing rates from normalized data.
+    """
+
+    def _reset_state(self, message: AxisArray) -> None:
+        ax_ix = message.get_axis_idx("ch")
+        nch = message.data.shape[ax_ix]
+        arr_size = (nch, 1) if ax_ix == 0 else (1, nch)
+        if self.settings.distribution == "uniform":
+            self.state.offsets = np.random.uniform(2.0, 40.0, size=arr_size)
+        elif self.settings.distribution == "normal":
+            self.state.offsets = np.random.normal(
+                loc=(self.settings.low_rate + self.settings.high_rate) / 2.0,
+                scale=(self.settings.high_rate - self.settings.low_rate) / 6.0,
+                size=arr_size,
+            )
+            self.state.offsets = np.clip(
+                self.state.offsets,
+                a_min=self.settings.low_rate,
+                a_max=self.settings.high_rate,
+            )
+        elif self.settings.distribution == "constant":
+            self.state.offsets = np.full(
+                shape=arr_size,
+                fill_value=(self.settings.low_rate + self.settings.high_rate) / 2.0,
+            )
+        else:
+            raise ValueError(f"Invalid distribution: {self.settings.distribution}")
+        # Input has std == 1
+        # Desired output has range from 0 to 2*self.state.offsets within 99.9% confidence interval
+        # For a standard normal distribution, 99.9% of data is within +/- 3.29 std devs.
+        # So, gain = offset / 3.29 to scale the std dev appropriately.
+        self.state.gains = self.state.offsets / 3.29
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        denorm = message.data * self.state.gains + self.state.offsets
+        return replace(
+            message,
+            data=np.clip(denorm, a_min=0.0, a_max=None),
+        )
+
+
+class DenormalizeUnit(BaseTransformerUnit[DenormalizeSettings, AxisArray, AxisArray, DenormalizeTransformer]):
+    SETTINGS = DenormalizeSettings

ezmsg/sigproc/detrend.py

@@ -0,0 +1,28 @@
+import scipy.signal as sps
+from ezmsg.baseproc import BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray, replace
+
+from ezmsg.sigproc.ewma import EWMASettings, EWMATransformer
+
+
+class DetrendTransformer(EWMATransformer):
+    """
+    Detrend the data using an exponentially weighted moving average (EWMA)
+     estimate of the mean.
+    """
+
+    def _process(self, message):
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        means, self._state.zi = 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=message.data - means)
+
+
+class DetrendUnit(BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, DetrendTransformer]):
+    SETTINGS = EWMASettings

ezmsg/sigproc/diff.py

@@ -0,0 +1,82 @@
+"""
+Compute differences along an axis.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from array_api_compat import get_namespace
+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):
+    axis: str | None = None
+    scale_by_fs: bool = False
+
+
+@processor_state
+class DiffState:
+    last_dat: npt.NDArray | None = None
+    last_time: float | None = None
+
+
+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 :]
+        return hash((sample_shape, message.key))
+
+    def _reset_state(self, message) -> None:
+        ax_idx = message.get_axis_idx(self.settings.axis)
+        self.state.last_dat = slice_along_axis(message.data, slice(0, 1), axis=ax_idx)
+        if self.settings.scale_by_fs:
+            ax_info = message.get_axis(self.settings.axis)
+            if hasattr(ax_info, "data"):
+                if len(ax_info.data) > 1:
+                    self.state.last_time = 2 * ax_info.data[0] - ax_info.data[1]
+                else:
+                    self.state.last_time = ax_info.data[0] - 0.001
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        axis = self.settings.axis or message.dims[0]
+        ax_idx = message.get_axis_idx(axis)
+
+        diffs = xp.diff(
+            xp.concat((self.state.last_dat, message.data), axis=ax_idx),
+            axis=ax_idx,
+        )
+        # Prepare last_dat for next iteration
+        self.state.last_dat = slice_along_axis(message.data, slice(-1, None), axis=ax_idx)
+        # Scale by fs if requested. This converts 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"):
+                # ax_info.data is typically numpy for metadata, so use np.diff here
+                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)
+                diffs /= xp.asarray(dt[exp_sl])
+                self.state.last_time = ax_info.data[-1]  # For next iteration
+            else:
+                diffs /= ax_info.gain
+
+        return replace(message, data=diffs)
+
+
+class DiffUnit(BaseTransformerUnit[DiffSettings, AxisArray, AxisArray, DiffTransformer]):
+    SETTINGS = DiffSettings
+
+
+def diff(axis: str = "time", scale_by_fs: bool = False) -> DiffTransformer:
+    return DiffTransformer(DiffSettings(axis=axis, scale_by_fs=scale_by_fs))

ezmsg/sigproc/downsample.py

@@ -1,82 +1,81 @@
-import typing
-
+import ezmsg.core as ez
 import numpy as np
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 from ezmsg.util.messages.axisarray import (
     AxisArray,
-    slice_along_axis,
     replace,
+    slice_along_axis,
 )
-from ezmsg.util.generator import consumer
-import ezmsg.core as ez
 
-from .base import GenAxisArray
 
+class DownsampleSettings(ez.Settings):
+    """
+    Settings for :obj:`Downsample` node.
+    """
 
-@consumer
-def downsample(
-    axis: str | None = None, target_rate: float | None = None
-) -> typing.Generator[AxisArray, AxisArray, None]:
+    axis: str = "time"
+    """The name of the axis along which to downsample."""
+
+    target_rate: float | None = None
+    """Desired rate after downsampling. The actual rate will be the nearest integer factor of the
+            input rate that is the same or higher than the target rate."""
+
+    factor: int | None = None
+    """Explicitly specify downsample factor.  If specified, target_rate is ignored."""
+
+
+@processor_state
+class DownsampleState:
+    q: int = 0
+    """The integer downsampling factor. It will be determined based on the target rate."""
+
+    s_idx: int = 0
+    """Index of the next msg's first sample into the virtual rotating ds_factor counter."""
+
+
+class DownsampleTransformer(BaseStatefulTransformer[DownsampleSettings, AxisArray, AxisArray, DownsampleState]):
     """
-    Construct a generator that yields a downsampled version of the data .send() to it.
     Downsampled data simply comprise every `factor`th sample.
     This should only be used following appropriate lowpass filtering.
     If your pipeline does not already have lowpass filtering then consider
     using the :obj:`Decimate` collection instead.
-
-    Args:
-        axis: The name of the axis along which to downsample.
-            Note: The axis must exist in the message .axes and be of type AxisArray.LinearAxis.
-        target_rate: Desired rate after downsampling. The actual rate will be the nearest integer factor of the
-            input rate that is the same or higher than the target rate.
-
-    Returns:
-        A primed generator object ready to receive an :obj:`AxisArray` via `.send(axis_array)`
-        and yields an :obj:`AxisArray` with its data downsampled.
-        Note that if a send chunk does not have sufficient samples to reach the
-        next downsample interval then an :obj:`AxisArray` with size-zero data is yielded.
-
     """
-    msg_out = AxisArray(np.array([]), dims=[""])
-
-    # state variables
-    factor: int = 0  # The integer downsampling factor. It will be determined based on the target rate.
-    s_idx: int = 0  # Index of the next msg's first sample into the virtual rotating ds_factor counter.
-
-    check_input = {"gain": None, "key": None}
 
-    while True:
-        msg_in: AxisArray = yield msg_out
+    def _hash_message(self, message: AxisArray) -> int:
+        return hash((message.axes[self.settings.axis].gain, message.key))
 
-        if axis is None:
-            axis = msg_in.dims[0]
-        axis_info = msg_in.get_axis(axis)
-        axis_idx = msg_in.get_axis_idx(axis)
+    def _reset_state(self, message: AxisArray) -> None:
+        axis_info = message.get_axis(self.settings.axis)
 
-        b_reset = (
-            msg_in.axes[axis].gain != check_input["gain"]
-            or msg_in.key != check_input["key"]
-        )
-        if b_reset:
-            check_input["gain"] = axis_info.gain
-            check_input["key"] = msg_in.key
-            # Reset state variables
-            s_idx = 0
-            if target_rate is None:
-                factor = 1
-            else:
-                factor = int(1 / (axis_info.gain * target_rate))
-            if factor < 1:
-                ez.logger.warning(
-                    f"Target rate {target_rate} cannot be achieved with input rate of {1/axis_info.gain}."
-                    "Setting factor to 1."
-                )
-                factor = 1
-
-        n_samples = msg_in.data.shape[axis_idx]
-        samples = np.arange(s_idx, s_idx + n_samples) % factor
+        if self.settings.factor is not None:
+            q = self.settings.factor
+        elif self.settings.target_rate is None:
+            q = 1
+        else:
+            q = int(1 / (axis_info.gain * self.settings.target_rate))
+        if q < 1:
+            ez.logger.warning(
+                f"Target rate {self.settings.target_rate} cannot be achieved with input rate of {1 / axis_info.gain}."
+                "Setting factor to 1."
+            )
+            q = 1
+        self._state.q = q
+        self._state.s_idx = 0
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis
+        axis_info = message.get_axis(axis)
+        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
         if n_samples > 0:
             # Update state for next iteration.
-            s_idx = samples[-1] + 1
+            self._state.s_idx = samples[-1] + 1
 
         pub_samples = np.where(samples == 0)[0]
         if len(pub_samples) > 0:
@@ -86,35 +85,27 @@ def downsample(
             n_step = 0
             data_slice = slice(None, 0, None)
         msg_out = replace(
-            msg_in,
-            data=slice_along_axis(msg_in.data, data_slice, axis=axis_idx),
+            message,
+            data=slice_along_axis(message.data, data_slice, axis=axis_idx),
             axes={
-                **msg_in.axes,
+                **message.axes,
                 axis: replace(
                     axis_info,
-                    gain=axis_info.gain * factor,
+                    gain=axis_info.gain * self._state.q,
                     offset=axis_info.offset + axis_info.gain * n_step,
                 ),
             },
         )
+        return msg_out
 
 
-class DownsampleSettings(ez.Settings):
-    """
-    Settings for :obj:`Downsample` node.
-    See :obj:`downsample` documentation for a description of the parameters.
-    """
-
-    axis: str | None = None
-    target_rate: float | None = None
-
-
-class Downsample(GenAxisArray):
-    """:obj:`Unit` for :obj:`bandpower`."""
-
+class Downsample(BaseTransformerUnit[DownsampleSettings, AxisArray, AxisArray, DownsampleTransformer]):
     SETTINGS = DownsampleSettings
 
-    def construct_generator(self):
-        self.STATE.gen = downsample(
-            axis=self.SETTINGS.axis, target_rate=self.SETTINGS.target_rate
-        )
+
+def downsample(
+    axis: str = "time",
+    target_rate: float | None = None,
+    factor: int | None = None,
+) -> DownsampleTransformer:
+    return DownsampleTransformer(DownsampleSettings(axis=axis, target_rate=target_rate, factor=factor))