ezmsg-sigproc 1.2.2__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,37 +1,65 @@
-import ezmsg.core as ez
-
-import scipy.signal
+import typing
 
+import ezmsg.core as ez
+from ezmsg.baseproc import BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 
+from .cheby import ChebyshevFilterSettings, ChebyshevFilterTransformer
 from .downsample import Downsample, DownsampleSettings
-from .filter import Filter, FilterCoefficients, FilterSettings
+from .filter import BACoeffs, SOSCoeffs
+
+
+class ChebyForDecimateTransformer(ChebyshevFilterTransformer[BACoeffs | SOSCoeffs]):
+    """
+    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 get_design_function(
+        self,
+    ) -> 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))
+            if ds_factor < 2:
+                return None
+            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):
-    SETTINGS: DownsampleSettings
+    """
+    A :obj:`Collection` chaining a :obj:`Filter` node configured as a lowpass Chebyshev filter
+    and a :obj:`Downsample` node.
+    """
+
+    SETTINGS = DownsampleSettings
 
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
-    FILTER = Filter()
+    FILTER = ChebyForDecimate()
     DOWNSAMPLE = Downsample()
 
     def configure(self) -> None:
+        cheby_settings = ChebyshevFilterSettings(
+            order=8,
+            ripple_tol=0.05,
+            Wn=0.4 * self.SETTINGS.target_rate,
+            btype="lowpass",
+            axis=self.SETTINGS.axis,
+            wn_hz=True,
+        )
+        self.FILTER.apply_settings(cheby_settings)
         self.DOWNSAMPLE.apply_settings(self.SETTINGS)
 
-        if self.SETTINGS.factor < 1:
-            raise ValueError("Decimation factor must be >= 1 (no decimation")
-        elif self.SETTINGS.factor == 1:
-            filt = FilterCoefficients()
-        else:
-            # See scipy.signal.decimate for IIR Filter Condition
-            b, a = scipy.signal.cheby1(8, 0.05, 0.8 / self.SETTINGS.factor)
-            system = scipy.signal.dlti(b, a)
-            filt = FilterCoefficients(b=system.num, a=system.den)  # type: ignore
-
-        self.FILTER.apply_settings(FilterSettings(filt=filt))
-
     def network(self) -> ez.NetworkDefinition:
         return (
             (self.INPUT_SIGNAL, self.FILTER.INPUT_SIGNAL),

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,63 +1,111 @@
-from dataclasses import replace
-
-from ezmsg.util.messages.axisarray import AxisArray
-
 import ezmsg.core as ez
 import numpy as np
-
-from typing import (
-    AsyncGenerator,
-    Optional,
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    replace,
+    slice_along_axis,
 )
 
 
 class DownsampleSettings(ez.Settings):
-    axis: Optional[str] = None
-    factor: int = 1
-
+    """
+    Settings for :obj:`Downsample` node.
+    """
 
-class DownsampleState(ez.State):
-    cur_settings: DownsampleSettings
-    s_idx: int = 0
-
-
-class Downsample(ez.Unit):
-    SETTINGS: DownsampleSettings
-    STATE: DownsampleState
+    axis: str = "time"
+    """The name of the axis along which to downsample."""
 
-    INPUT_SETTINGS = ez.InputStream(DownsampleSettings)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    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."""
 
-    def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
+    factor: int | None = None
+    """Explicitly specify downsample factor.  If specified, target_rate is ignored."""
 
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: DownsampleSettings) -> None:
-        self.STATE.cur_settings = msg
 
-    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
-    @ez.publisher(OUTPUT_SIGNAL)
-    async def on_signal(self, msg: AxisArray) -> AsyncGenerator:
-        if self.STATE.cur_settings.factor < 1:
-            raise ValueError("Downsample factor must be at least 1 (no downsampling)")
+@processor_state
+class DownsampleState:
+    q: int = 0
+    """The integer downsampling factor. It will be determined based on the target rate."""
 
-        axis_name = self.STATE.cur_settings.axis
-        if axis_name is None:
-            axis_name = msg.dims[0]
-        axis = msg.get_axis(axis_name)
-        axis_idx = msg.get_axis_idx(axis_name)
-
-        samples = np.arange(msg.data.shape[axis_idx]) + self.STATE.s_idx
-        samples = samples % self.STATE.cur_settings.factor
-        self.STATE.s_idx = samples[-1] + 1
+    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]):
+    """
+    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.
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        return hash((message.axes[self.settings.axis].gain, message.key))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis_info = message.get_axis(self.settings.axis)
+
+        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.
+            self._state.s_idx = samples[-1] + 1
 
         pub_samples = np.where(samples == 0)[0]
-        if len(pub_samples) != 0:
-            new_axes = {ax_name: msg.get_axis(ax_name) for ax_name in msg.dims}
-            new_offset = axis.offset + (axis.gain * pub_samples[0].item())
-            new_gain = axis.gain * self.STATE.cur_settings.factor
-            new_axes[axis_name] = replace(axis, gain=new_gain, offset=new_offset)
-            down_data = np.take(msg.data, pub_samples, axis_idx)
-            out_msg = replace(msg, data=down_data, dims=msg.dims, axes=new_axes)
-            yield self.OUTPUT_SIGNAL, out_msg
+        if len(pub_samples) > 0:
+            n_step = pub_samples[0].item()
+            data_slice = pub_samples
+        else:
+            n_step = 0
+            data_slice = slice(None, 0, None)
+        msg_out = replace(
+            message,
+            data=slice_along_axis(message.data, data_slice, axis=axis_idx),
+            axes={
+                **message.axes,
+                axis: replace(
+                    axis_info,
+                    gain=axis_info.gain * self._state.q,
+                    offset=axis_info.offset + axis_info.gain * n_step,
+                ),
+            },
+        )
+        return msg_out
+
+
+class Downsample(BaseTransformerUnit[DownsampleSettings, AxisArray, AxisArray, DownsampleTransformer]):
+    SETTINGS = DownsampleSettings
+
+
+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))