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

ezmsg/sigproc/spectrum.py

@@ -0,0 +1,277 @@
+import enum
+import typing
+from functools import partial
+
+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,
+    replace,
+    slice_along_axis,
+)
+
+
+class OptionsEnum(enum.Enum):
+    @classmethod
+    def options(cls):
+        return list(map(lambda c: c.value, cls))
+
+
+class WindowFunction(OptionsEnum):
+    """Windowing function prior to calculating spectrum."""
+
+    NONE = "None (Rectangular)"
+    """None."""
+
+    HAMMING = "Hamming"
+    """:obj:`numpy.hamming`"""
+
+    HANNING = "Hanning"
+    """:obj:`numpy.hanning`"""
+
+    BARTLETT = "Bartlett"
+    """:obj:`numpy.bartlett`"""
+
+    BLACKMAN = "Blackman"
+    """:obj:`numpy.blackman`"""
+
+
+WINDOWS = {
+    WindowFunction.NONE: np.ones,
+    WindowFunction.HAMMING: np.hamming,
+    WindowFunction.HANNING: np.hanning,
+    WindowFunction.BARTLETT: np.bartlett,
+    WindowFunction.BLACKMAN: np.blackman,
+}
+
+
+class SpectralTransform(OptionsEnum):
+    """Additional transformation functions to apply to the spectral result."""
+
+    RAW_COMPLEX = "Complex FFT Output"
+    REAL = "Real Component of FFT"
+    IMAG = "Imaginary Component of FFT"
+    REL_POWER = "Relative Power"
+    REL_DB = "Log Power (Relative dB)"
+
+
+class SpectralOutput(OptionsEnum):
+    """The expected spectral contents."""
+
+    FULL = "Full Spectrum"
+    POSITIVE = "Positive Frequencies"
+    NEGATIVE = "Negative Frequencies"
+
+
+class SpectrumSettings(ez.Settings):
+    """
+    Settings for :obj:`Spectrum.
+    See :obj:`spectrum` for a description of the parameters.
+    """
+
+    axis: str | None = None
+    """
+    The name of the axis on which to calculate the spectrum.
+      Note: The axis must have an .axes entry of type LinearAxis, not CoordinateAxis.
+    """
+
+    # n: int | None = None # n parameter for fft
+
+    out_axis: str | None = "freq"
+    """The name of the new axis. Defaults to "freq". If none; don't change dim name"""
+
+    window: WindowFunction = WindowFunction.HAMMING
+    """The :obj:`WindowFunction` to apply to the data slice prior to calculating the spectrum."""
+
+    transform: SpectralTransform = SpectralTransform.REL_DB
+    """The :obj:`SpectralTransform` to apply to the spectral magnitude."""
+
+    output: SpectralOutput = SpectralOutput.POSITIVE
+    """The :obj:`SpectralOutput` format."""
+
+    norm: str | None = "forward"
+    """
+    Normalization mode. Default "forward" is best used when the inverse transform is not needed,
+      for example when the goal is to get spectral power. Use "backward" (equivalent to None) to not
+      scale the spectrum which is useful when the spectra will be manipulated and possibly inverse-transformed.
+      See numpy.fft.fft for details.
+    """
+
+    do_fftshift: bool = True
+    """
+    Whether to apply fftshift to the output. Default is True.
+      This value is ignored unless output is SpectralOutput.FULL.
+    """
+
+    nfft: int | None = None
+    """
+    The number of points to use for the FFT. If None, the length of the input data is used.
+    """
+
+
+@processor_state
+class SpectrumState:
+    f_sl: slice | None = None
+    # I would prefer `slice(None)` as f_sl default but this fails because it is mutable.
+    freq_axis: AxisArray.LinearAxis | None = None
+    fftfun: typing.Callable | None = None
+    f_transform: typing.Callable | None = None
+    new_dims: list[str] | None = None
+    window: npt.NDArray | None = None
+
+
+class SpectrumTransformer(BaseStatefulTransformer[SpectrumSettings, AxisArray, AxisArray, SpectrumState]):
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[0]
+        ax_idx = message.get_axis_idx(axis)
+        ax_info = message.axes[axis]
+        targ_len = message.data.shape[ax_idx]
+        return hash((targ_len, message.data.ndim, message.data.dtype.kind, ax_idx, ax_info.gain))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis = self.settings.axis or message.dims[0]
+        ax_idx = message.get_axis_idx(axis)
+        ax_info = message.axes[axis]
+        targ_len = message.data.shape[ax_idx]
+        nfft = self.settings.nfft or targ_len
+
+        # Pre-calculate windowing
+        window = WINDOWS[self.settings.window](targ_len)
+        window = window.reshape(
+            [1] * ax_idx
+            + [
+                len(window),
+            ]
+            + [1] * (message.data.ndim - 1 - ax_idx)
+        )
+        if self.settings.transform != SpectralTransform.RAW_COMPLEX and not (
+            self.settings.transform == SpectralTransform.REAL or self.settings.transform == SpectralTransform.IMAG
+        ):
+            scale = np.sum(window**2.0) * ax_info.gain
+
+        if self.settings.window != WindowFunction.NONE:
+            self.state.window = window
+
+        # Pre-calculate frequencies and select our fft function.
+        b_complex = message.data.dtype.kind == "c"
+        self.state.f_sl = slice(None)
+        if (not b_complex) and self.settings.output == SpectralOutput.POSITIVE:
+            # If input is not complex and desired output is SpectralOutput.POSITIVE, we can save some computation
+            #  by using rfft and rfftfreq.
+            self.state.fftfun = partial(np.fft.rfft, n=nfft, axis=ax_idx, norm=self.settings.norm)
+            freqs = np.fft.rfftfreq(nfft, d=ax_info.gain * targ_len / nfft)
+        else:
+            self.state.fftfun = partial(np.fft.fft, n=nfft, axis=ax_idx, norm=self.settings.norm)
+            freqs = np.fft.fftfreq(nfft, d=ax_info.gain * targ_len / nfft)
+            if self.settings.output == SpectralOutput.POSITIVE:
+                self.state.f_sl = slice(None, nfft // 2 + 1 - (nfft % 2))
+            elif self.settings.output == SpectralOutput.NEGATIVE:
+                freqs = np.fft.fftshift(freqs, axes=-1)
+                self.state.f_sl = slice(None, nfft // 2 + 1)
+            elif self.settings.do_fftshift and self.settings.output == SpectralOutput.FULL:
+                freqs = np.fft.fftshift(freqs, axes=-1)
+            freqs = freqs[self.state.f_sl]
+        freqs = freqs.tolist()  # To please type checking
+        self.state.freq_axis = AxisArray.LinearAxis(unit="Hz", gain=freqs[1] - freqs[0], offset=freqs[0])
+        self.state.new_dims = (
+            message.dims[:ax_idx]
+            + [
+                self.settings.out_axis or axis,
+            ]
+            + message.dims[ax_idx + 1 :]
+        )
+
+        def f_transform(x):
+            return x
+
+        if self.settings.transform != SpectralTransform.RAW_COMPLEX:
+            if self.settings.transform == SpectralTransform.REAL:
+
+                def f_transform(x):
+                    return x.real
+            elif self.settings.transform == SpectralTransform.IMAG:
+
+                def f_transform(x):
+                    return x.imag
+            else:
+
+                def f1(x):
+                    return (np.abs(x) ** 2.0) / scale
+
+                if self.settings.transform == SpectralTransform.REL_DB:
+
+                    def f_transform(x):
+                        return 10 * np.log10(f1(x))
+                else:
+                    f_transform = f1
+        self.state.f_transform = f_transform
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[0]
+        ax_idx = message.get_axis_idx(axis)
+        targ_len = message.data.shape[ax_idx]
+
+        new_axes = {k: v for k, v in message.axes.items() if k not in [self.settings.out_axis, axis]}
+        new_axes[self.settings.out_axis or axis] = self.state.freq_axis
+
+        if self.state.window is not None:
+            win_dat = message.data * self.state.window
+        else:
+            win_dat = message.data
+        spec = self.state.fftfun(
+            win_dat,
+            n=self.settings.nfft or targ_len,
+            axis=ax_idx,
+            norm=self.settings.norm,
+        )
+        # Note: norm="forward" equivalent to `/ nfft`
+        if (
+            self.settings.do_fftshift and self.settings.output == SpectralOutput.FULL
+        ) or self.settings.output == SpectralOutput.NEGATIVE:
+            spec = np.fft.fftshift(spec, axes=ax_idx)
+        spec = self.state.f_transform(spec)
+        spec = slice_along_axis(spec, self.state.f_sl, ax_idx)
+
+        msg_out = replace(message, data=spec, dims=self.state.new_dims, axes=new_axes)
+        return msg_out
+
+
+class Spectrum(BaseTransformerUnit[SpectrumSettings, AxisArray, AxisArray, SpectrumTransformer]):
+    SETTINGS = SpectrumSettings
+
+
+def spectrum(
+    axis: str | None = None,
+    out_axis: str | None = "freq",
+    window: WindowFunction = WindowFunction.HANNING,
+    transform: SpectralTransform = SpectralTransform.REL_DB,
+    output: SpectralOutput = SpectralOutput.POSITIVE,
+    norm: str | None = "forward",
+    do_fftshift: bool = True,
+    nfft: int | None = None,
+) -> SpectrumTransformer:
+    """
+    Calculate a spectrum on a data slice.
+
+    Returns:
+        A :obj:`SpectrumTransformer` object that expects an :obj:`AxisArray` via `.(axis_array)` (__call__)
+        containing continuous data and returns an :obj:`AxisArray` with data of spectral magnitudes or powers.
+    """
+    return SpectrumTransformer(
+        SpectrumSettings(
+            axis=axis,
+            out_axis=out_axis,
+            window=window,
+            transform=transform,
+            output=output,
+            norm=norm,
+            do_fftshift=do_fftshift,
+            nfft=nfft,
+        )
+    )

ezmsg/sigproc/transpose.py

@@ -0,0 +1,134 @@
+"""
+Transpose or permute array dimensions.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+    Memory layout optimization (C/F order) only applies to NumPy arrays.
+"""
+
+from types import EllipsisType
+
+import ezmsg.core as ez
+import numpy as np
+from array_api_compat import get_namespace, is_numpy_array
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    replace,
+)
+
+
+class TransposeSettings(ez.Settings):
+    """
+    Settings for :obj:`Transpose` node.
+
+    Fields:
+      axes:
+    """
+
+    axes: tuple[int | str | EllipsisType, ...] | None = None
+    order: str | None = None
+
+
+@processor_state
+class TransposeState:
+    axes_ints: tuple[int, ...] | None = None
+
+
+class TransposeTransformer(BaseStatefulTransformer[TransposeSettings, AxisArray, AxisArray, TransposeState]):
+    """
+    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(tuple(message.dims))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        if self.settings.axes is None:
+            self._state.axes_ints = None
+        else:
+            ell_ix = [ix for ix, ax in enumerate(self.settings.axes) if ax is Ellipsis]
+            if len(ell_ix) > 1:
+                raise ValueError("Only one Ellipsis is allowed in axes.")
+            ell_ix = ell_ix[0] if len(ell_ix) == 1 else len(message.dims)
+            prefix = []
+            for ax in self.settings.axes[:ell_ix]:
+                if isinstance(ax, int):
+                    prefix.append(ax)
+                else:
+                    if ax not in message.dims:
+                        raise ValueError(f"Axis {ax} not found in message dims.")
+                    prefix.append(message.dims.index(ax))
+            suffix = []
+            for ax in self.settings.axes[ell_ix + 1 :]:
+                if isinstance(ax, int):
+                    suffix.append(ax)
+                else:
+                    if ax not in message.dims:
+                        raise ValueError(f"Axis {ax} not found in message dims.")
+                    suffix.append(message.dims.index(ax))
+            ells = [_ for _ in range(message.data.ndim) if _ not in prefix and _ not in suffix]
+            re_ix = tuple(prefix + ells + suffix)
+            if re_ix == tuple(range(message.data.ndim)):
+                self._state.axes_ints = None
+            else:
+                self._state.axes_ints = re_ix
+        if self.settings.order is not None and self.settings.order.upper()[0] not in [
+            "C",
+            "F",
+        ]:
+            raise ValueError("order must be 'C' or 'F'.")
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        if self.settings.axes is None and self.settings.order is None:
+            # Passthrough
+            return message
+        return super().__call__(message)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        if self.state.axes_ints is None:
+            # No transpose required
+            if self.settings.order is None:
+                # No memory relayout required
+                # Note: We should not be able to reach here because it should be shortcutted at passthrough.
+                msg_out = message
+            else:
+                # Memory layout optimization only applies to numpy arrays
+                if is_numpy_array(message.data):
+                    msg_out = replace(
+                        message,
+                        data=np.require(message.data, requirements=self.settings.order.upper()[0]),
+                    )
+                else:
+                    msg_out = message
+        else:
+            dims_out = [message.dims[ix] for ix in self.state.axes_ints]
+            data_out = xp.permute_dims(message.data, axes=self.state.axes_ints)
+            if self.settings.order is not None and is_numpy_array(data_out):
+                # Memory layout optimization only applies to numpy arrays
+                data_out = np.require(data_out, requirements=self.settings.order.upper()[0])
+            msg_out = replace(
+                message,
+                data=data_out,
+                dims=dims_out,
+            )
+        return msg_out
+
+
+class Transpose(BaseTransformerUnit[TransposeSettings, AxisArray, AxisArray, TransposeTransformer]):
+    SETTINGS = TransposeSettings
+
+
+def transpose(
+    axes: tuple[int | str | EllipsisType, ...] | None = None, order: str | None = None
+) -> TransposeTransformer:
+    return TransposeTransformer(TransposeSettings(axes=axes, order=order))

ezmsg/sigproc/util/asio.py

@@ -0,0 +1,25 @@
+"""
+Backwards-compatible re-exports from ezmsg.baseproc.util.asio.
+
+New code should import directly from ezmsg.baseproc instead.
+"""
+
+import warnings
+
+warnings.warn(
+    "Importing from 'ezmsg.sigproc.util.asio' is deprecated. Please import from 'ezmsg.baseproc.util.asio' instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+from ezmsg.baseproc.util.asio import (  # noqa: E402
+    CoroutineExecutionError,
+    SyncToAsyncGeneratorWrapper,
+    run_coroutine_sync,
+)
+
+__all__ = [
+    "CoroutineExecutionError",
+    "SyncToAsyncGeneratorWrapper",
+    "run_coroutine_sync",
+]