ezmsg-sigproc 1.8.2__py3-none-any.whl → 2.1.0__py3-none-any.whl

ezmsg/sigproc/spectrogram.py

@@ -1,95 +1,97 @@
-import typing
-
+from typing import Generator
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.generator import consumer, compose
 from ezmsg.util.messages.modify import modify_axis
 
-from .window import windowing, Anchor
-from .spectrum import spectrum, WindowFunction, SpectralTransform, SpectralOutput
-from .base import GenAxisArray
-
-
-@consumer
-def spectrogram(
-    window_dur: float | None = None,
-    window_shift: float | None = None,
-    window_anchor: str | Anchor = Anchor.BEGINNING,
-    window: WindowFunction = WindowFunction.HANNING,
-    transform: SpectralTransform = SpectralTransform.REL_DB,
-    output: SpectralOutput = SpectralOutput.POSITIVE,
-) -> typing.Generator[AxisArray | None, AxisArray, None]:
-    """
-    Calculate a spectrogram on streaming data.
-
-    Chains :obj:`ezmsg.sigproc.window.windowing` to apply a moving window on the data,
-    :obj:`ezmsg.sigproc.spectrum.spectrum` to calculate spectra for each window,
-    and finally :obj:`ezmsg.util.messages.modify.modify_axis` to convert the win axis back to time axis.
-
-    Args:
-        window_dur: See :obj:`ezmsg.sigproc.window.windowing`
-        window_shift: See :obj:`ezmsg.sigproc.window.windowing`
-        window_anchor: See :obj:`ezmsg.sigproc.window.windowing`
-        window: See :obj:`ezmsg.sigproc.spectrum.spectrum`
-        transform: See :obj:`ezmsg.sigproc.spectrum.spectrum`
-        output: See :obj:`ezmsg.sigproc.spectrum.spectrum`
-
-    Returns:
-        A primed generator object that expects an :obj:`AxisArray` via `.send(axis_array)`
-        with continuous data in its .data payload, and yields an :obj:`AxisArray` of time-frequency power values.
-    """
-
-    pipeline = compose(
-        windowing(
-            axis="time",
-            newaxis="win",
-            window_dur=window_dur,
-            window_shift=window_shift,
-            zero_pad_until="shift" if window_shift is not None else "input",
-            anchor=window_anchor,
-        ),
-        spectrum(axis="time", window=window, transform=transform, output=output),
-        modify_axis(name_map={"win": "time"}),
-    )
-
-    # State variables
-    msg_out: AxisArray | None = None
-
-    while True:
-        msg_in: AxisArray = yield msg_out
-        msg_out = pipeline(msg_in)
+from .window import Anchor, WindowTransformer
+from .spectrum import (
+    WindowFunction,
+    SpectralTransform,
+    SpectralOutput,
+    SpectrumTransformer,
+)
+from .base import (
+    CompositeProcessor,
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+)
 
 
 class SpectrogramSettings(ez.Settings):
     """
-    Settings for :obj:`Spectrogram`.
-    See :obj:`spectrogram` for a description of the parameters.
+    Settings for :obj:`SpectrogramTransformer`.
     """
 
-    window_dur: float | None = None  # window duration in seconds
+    window_dur: float | None = None
+    """window duration in seconds."""
+
     window_shift: float | None = None
     """"window step in seconds. If None, window_shift == window_dur"""
+
     window_anchor: str | Anchor = Anchor.BEGINNING
+    """See :obj"`WindowTransformer`"""
 
-    # See SpectrumSettings for details of following settings:
     window: WindowFunction = WindowFunction.HAMMING
-    transform: SpectralTransform = SpectralTransform.REL_DB
-    output: SpectralOutput = SpectralOutput.POSITIVE
-
+    """The :obj:`WindowFunction` to apply to the data slice prior to calculating the spectrum."""
 
-class Spectrogram(GenAxisArray):
-    """
-    Unit for :obj:`spectrogram`.
-    """
+    transform: SpectralTransform = SpectralTransform.REL_DB
+    """The :obj:`SpectralTransform` to apply to the spectral magnitude."""
 
+    output: SpectralOutput = SpectralOutput.POSITIVE
+    """The :obj:`SpectralOutput` format."""
+
+
+class SpectrogramTransformer(
+    CompositeProcessor[SpectrogramSettings, AxisArray, AxisArray]
+):
+    @staticmethod
+    def _initialize_processors(
+        settings: SpectrogramSettings,
+    ) -> dict[str, BaseStatefulProcessor | Generator[AxisArray, AxisArray, None]]:
+        return {
+            "windowing": WindowTransformer(
+                axis="time",
+                newaxis="win",
+                window_dur=settings.window_dur,
+                window_shift=settings.window_shift,
+                zero_pad_until="shift"
+                if settings.window_shift is not None
+                else "input",
+                anchor=settings.window_anchor,
+            ),
+            "spectrum": SpectrumTransformer(
+                axis="time",
+                window=settings.window,
+                transform=settings.transform,
+                output=settings.output,
+            ),
+            "modify_axis": modify_axis(name_map={"win": "time"}),
+        }
+
+
+class Spectrogram(
+    BaseTransformerUnit[
+        SpectrogramSettings, AxisArray, AxisArray, SpectrogramTransformer
+    ]
+):
     SETTINGS = SpectrogramSettings
 
-    def construct_generator(self):
-        self.STATE.gen = spectrogram(
-            window_dur=self.SETTINGS.window_dur,
-            window_shift=self.SETTINGS.window_shift,
-            window_anchor=self.SETTINGS.window_anchor,
-            window=self.SETTINGS.window,
-            transform=self.SETTINGS.transform,
-            output=self.SETTINGS.output,
+
+def spectrogram(
+    window_dur: float | None = None,
+    window_shift: float | None = None,
+    window_anchor: str | Anchor = Anchor.BEGINNING,
+    window: WindowFunction = WindowFunction.HAMMING,
+    transform: SpectralTransform = SpectralTransform.REL_DB,
+    output: SpectralOutput = SpectralOutput.POSITIVE,
+) -> SpectrogramTransformer:
+    return SpectrogramTransformer(
+        SpectrogramSettings(
+            window_dur=window_dur,
+            window_shift=window_shift,
+            window_anchor=window_anchor,
+            window=window,
+            transform=transform,
+            output=output,
         )
+    )

ezmsg/sigproc/spectrum.py

@@ -3,15 +3,19 @@ from functools import partial
 import typing
 
 import numpy as np
+import numpy.typing as npt
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import (
     AxisArray,
     slice_along_axis,
     replace,
 )
-from ezmsg.util.generator import consumer
 
-from .base import GenAxisArray
+from .base import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 
 
 class OptionsEnum(enum.Enum):
@@ -66,198 +70,229 @@ class SpectralOutput(OptionsEnum):
     NEGATIVE = "Negative Frequencies"
 
 
-@consumer
-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,
-) -> typing.Generator[AxisArray, AxisArray, None]:
+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.
     """
-    Calculate a spectrum on a data slice.
 
-    Args:
-        axis: The name of the axis on which to calculate the spectrum.
-            Note: The axis must have an .axes entry of type LinearAxis, not CoordinateAxis.
-        out_axis: The name of the new axis. Defaults to "freq".
-        window: The :obj:`WindowFunction` to apply to the data slice prior to calculating the spectrum.
-        transform: The :obj:`SpectralTransform` to apply to the spectral magnitude.
-        output: The :obj:`SpectralOutput` format.
-        norm: 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: Whether to apply fftshift to the output. Default is True. This value is ignored unless
-          output is SpectralOutput.FULL.
-        nfft: The number of points to use for the FFT. If None, the length of the input data is used.
+    # n: int | None = None # n parameter for fft
 
-    Returns:
-        A primed generator object that expects an :obj:`AxisArray` via `.send(axis_array)` containing continuous data
-        and yields an :obj:`AxisArray` with data of spectral magnitudes or powers.
+    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
     """
-    msg_out = AxisArray(np.array([]), dims=[""])
+    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.
+    """
+
 
-    # State variables
-    apply_window = window != WindowFunction.NONE
-    do_fftshift &= output == SpectralOutput.FULL
-    f_sl = slice(None)
+@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)
+        )
 
-    # Reset if input changes substantially
-    check_input = {
-        "n_time": None,  # Need to recalc windows
-        "ndim": None,  # Input ndim changed: Need to recalc windows
-        "kind": None,  # Input dtype changed: Need to re-init fft funcs
-        "ax_idx": None,  # Axis index changed: Need to re-init fft funcs
-        "gain": None,  # Gain changed: Need to re-calc freqs
-        # "key": None  # There's no temporal continuity; we can ignore key changes
-    }
-
-    while True:
-        msg_in: AxisArray = yield msg_out
-
-        # Get signal properties
-        axis = axis or msg_in.dims[0]
-        ax_idx = msg_in.get_axis_idx(axis)
-        ax_info = msg_in.axes[axis]
-        targ_len = msg_in.data.shape[ax_idx]
-
-        # Check signal properties for change
-        b_reset = targ_len != check_input["n_time"]
-        b_reset = b_reset or msg_in.data.ndim != check_input["ndim"]
-        b_reset = b_reset or msg_in.data.dtype.kind != check_input["kind"]
-        b_reset = b_reset or ax_idx != check_input["ax_idx"]
-        b_reset = b_reset or ax_info.gain != check_input["gain"]
-        if b_reset:
-            check_input["n_time"] = targ_len
-            check_input["ndim"] = msg_in.data.ndim
-            check_input["kind"] = msg_in.data.dtype.kind
-            check_input["ax_idx"] = ax_idx
-            check_input["gain"] = ax_info.gain
-
-            nfft = nfft or targ_len
-
-            # Pre-calculate windowing
-            window = WINDOWS[window](targ_len)
-            window = window.reshape(
-                [1] * ax_idx
-                + [
-                    len(window),
-                ]
-                + [1] * (msg_in.data.ndim - 1 - ax_idx)
-            )
-            if transform != SpectralTransform.RAW_COMPLEX and not (
-                transform == SpectralTransform.REAL
-                or transform == SpectralTransform.IMAG
-            ):
-                scale = np.sum(window**2.0) * ax_info.gain
-
-            # Pre-calculate frequencies and select our fft function.
-            b_complex = msg_in.data.dtype.kind == "c"
-            if (not b_complex) and output == SpectralOutput.POSITIVE:
-                # If input is not complex and desired output is SpectralOutput.POSITIVE, we can save some computation
-                #  by using rfft and rfftfreq.
-                fftfun = partial(np.fft.rfft, n=nfft, axis=ax_idx, norm=norm)
-                freqs = np.fft.rfftfreq(nfft, d=ax_info.gain * targ_len / nfft)
-            else:
-                fftfun = partial(np.fft.fft, n=nfft, axis=ax_idx, norm=norm)
-                freqs = np.fft.fftfreq(nfft, d=ax_info.gain * targ_len / nfft)
-                if output == SpectralOutput.POSITIVE:
-                    f_sl = slice(None, nfft // 2 + 1 - (nfft % 2))
-                elif output == SpectralOutput.NEGATIVE:
-                    freqs = np.fft.fftshift(freqs, axes=-1)
-                    f_sl = slice(None, nfft // 2 + 1)
-                elif do_fftshift:  # and FULL
-                    freqs = np.fft.fftshift(freqs, axes=-1)
-                freqs = freqs[f_sl]
-            freqs = freqs.tolist()  # To please type checking
-            freq_axis = AxisArray.LinearAxis(
-                unit="Hz", gain=freqs[1] - freqs[0], offset=freqs[0]
+    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
             )
-            if out_axis is None:
-                out_axis = axis
-            new_dims = (
-                msg_in.dims[:ax_idx]
-                + [
-                    out_axis,
-                ]
-                + msg_in.dims[ax_idx + 1 :]
+            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
+        def f_transform(x):
+            return x
 
-            if transform != SpectralTransform.RAW_COMPLEX:
-                if transform == SpectralTransform.REAL:
+        if self.settings.transform != SpectralTransform.RAW_COMPLEX:
+            if self.settings.transform == SpectralTransform.REAL:
 
-                    def f_transform(x):
-                        return x.real
-                elif transform == SpectralTransform.IMAG:
+                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 transform == SpectralTransform.REL_DB:
+                def f_transform(x):
+                    return x.imag
+            else:
 
-                        def f_transform(x):
-                            return 10 * np.log10(f1(x))
-                    else:
-                        f_transform = f1
+                def f1(x):
+                    return (np.abs(x) ** 2.0) / scale
 
-        new_axes = {k: v for k, v in msg_in.axes.items() if k not in [out_axis, axis]}
-        new_axes[out_axis] = freq_axis
+                if self.settings.transform == SpectralTransform.REL_DB:
 
-        if apply_window:
-            win_dat = msg_in.data * window
+                    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 = msg_in.data
-        spec = fftfun(win_dat, n=nfft, axis=ax_idx, norm=norm)
+            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 do_fftshift or output == SpectralOutput.NEGATIVE:
+        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 = f_transform(spec)
-        spec = slice_along_axis(spec, f_sl, ax_idx)
-
-        msg_out = replace(msg_in, data=spec, dims=new_dims, axes=new_axes)
+        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 SpectrumSettings(ez.Settings):
-    """
-    Settings for :obj:`Spectrum.
-    See :obj:`spectrum` for a description of the parameters.
-    """
-
-    axis: str | None = None
-    # n: int | None = None # n parameter for fft
-    out_axis: str | None = "freq"  # If none; don't change dim name
-    window: WindowFunction = WindowFunction.HAMMING
-    transform: SpectralTransform = SpectralTransform.REL_DB
-    output: SpectralOutput = SpectralOutput.POSITIVE
-
-
-class Spectrum(GenAxisArray):
-    """Unit for :obj:`spectrum`"""
 
+class Spectrum(
+    BaseTransformerUnit[SpectrumSettings, AxisArray, AxisArray, SpectrumTransformer]
+):
     SETTINGS = SpectrumSettings
 
-    INPUT_SETTINGS = ez.InputStream(SpectrumSettings)
 
-    def construct_generator(self):
-        self.STATE.gen = spectrum(
-            axis=self.SETTINGS.axis,
-            out_axis=self.SETTINGS.out_axis,
-            window=self.SETTINGS.window,
-            transform=self.SETTINGS.transform,
-            output=self.SETTINGS.output,
+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,
         )
+    )