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

ezmsg/sigproc/slicer.py

@@ -0,0 +1,98 @@
+from dataclasses import replace
+import typing
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.generator import consumer, GenAxisArray
+
+
+"""
+Slicer:Select a subset of data along a particular axis.
+"""
+
+
+def parse_slice(s: str) -> typing.Tuple[typing.Union[slice, int], ...]:
+    """
+    Parses a string representation of a slice and returns a tuple of slice objects.
+    * "" -> slice(None, None, None)  (take all)
+    * ":" -> slice(None, None, None)
+    * '"none"` (case-insensitive) -> slice(None, None, None)
+    * "{start}:{stop}" or {start}:{stop}:{step} -> slice(start, stop, step)
+    * "5" (or any integer) -> (5,). Take only that item.
+        applying this to a ndarray or AxisArray will drop the dimension.
+    * A comma-separated list of the above -> a tuple of slices | ints
+
+    Args:
+        s (str): The string representation of the slice.
+
+    Returns:
+        tuple[slice | int, ...]: A tuple of slice objects and/or ints.
+    """
+    if s.lower() in ["", ":", "none"]:
+        return (slice(None),)
+    if "," not in s:
+        parts = [part.strip() for part in s.split(":")]
+        if len(parts) == 1:
+            return (int(parts[0]),)
+        return (slice(*(int(part.strip()) if part else None for part in parts)),)
+    l = [parse_slice(_) for _ in s.split(",")]
+    return tuple([item for sublist in l for item in sublist])
+
+
+@consumer
+def slicer(
+    selection: str = "", axis: typing.Optional[str] = None
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = AxisArray(np.array([]), dims=[""])
+    _slice = None
+    b_change_dims = False
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+
+        if axis is None:
+            axis = axis_arr_in.dims[-1]
+        axis_idx = axis_arr_in.get_axis_idx(axis)
+
+        if _slice is None:
+            _slices = parse_slice(selection)
+            if len(_slices) == 1:
+                _slice = _slices[0]
+                b_change_dims = isinstance(_slice, int)
+            else:
+                # Multiple slices, but this cannot be done in a single step, so we convert the slices
+                #  to a discontinuous set of integer indexes.
+                indices = np.arange(axis_arr_in.data.shape[axis_idx])
+                indices = np.hstack([indices[_] for _ in _slices])
+                _slice = np.s_[indices]
+
+        if b_change_dims:
+            out_dims = [_ for dim_ix, _ in enumerate(axis_arr_in.dims) if dim_ix != axis_idx]
+            out_axes = axis_arr_in.axes.copy()
+            out_axes.pop(axis, None)
+        else:
+            out_dims = axis_arr_in.dims
+            out_axes = axis_arr_in.axes
+
+        axis_arr_out = replace(
+            axis_arr_in,
+            dims=out_dims,
+            axes=out_axes,
+            data=slice_along_axis(axis_arr_in.data, _slice, axis_idx),
+        )
+
+
+class SlicerSettings(ez.Settings):
+    selection: str = ""
+    axis: typing.Optional[str] = None
+
+
+class Slicer(GenAxisArray):
+    SETTINGS: SlicerSettings
+
+    def construct_generator(self):
+        self.STATE.gen = slicer(
+            selection=self.SETTINGS.selection, axis=self.SETTINGS.axis
+        )

ezmsg/sigproc/spectral.py

@@ -1,132 +1,9 @@
-import enum
-
-from dataclasses import replace
-
-import numpy as np
-import ezmsg.core as ez
-
-from ezmsg.util.messages.axisarray import AxisArray
-
-from typing import Optional, AsyncGenerator
-
-
-class OptionsEnum(enum.Enum):
-    @classmethod
-    def options(cls):
-        return list(map(lambda c: c.value, cls))
-
-
-class WindowFunction(OptionsEnum):
-    NONE = "None (Rectangular)"
-    HAMMING = "Hamming"
-    HANNING = "Hanning"
-    BARTLETT = "Bartlett"
-    BLACKMAN = "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):
-    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):
-    FULL = "Full Spectrum"
-    POSITIVE = "Positive Frequencies"
-    NEGATIVE = "Negative Frequencies"
-
-
-class SpectrumSettings(ez.Settings):
-    axis: Optional[str] = None
-    # n: Optional[int] = None # n parameter for fft
-    out_axis: Optional[str] = "freq"  # If none; don't change dim name
-    window: WindowFunction = WindowFunction.HAMMING
-    transform: SpectralTransform = SpectralTransform.REL_DB
-    output: SpectralOutput = SpectralOutput.POSITIVE
-
-
-class SpectrumState(ez.State):
-    cur_settings: SpectrumSettings
-
-
-class Spectrum(ez.Unit):
-    SETTINGS: SpectrumSettings
-    STATE: SpectrumState
-
-    INPUT_SETTINGS = ez.InputStream(SpectrumSettings)
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
-
-    def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
-
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: SpectrumSettings):
-        self.STATE.cur_settings = msg
-
-    @ez.subscriber(INPUT_SIGNAL)
-    @ez.publisher(OUTPUT_SIGNAL)
-    async def on_data(self, message: AxisArray) -> AsyncGenerator:
-        axis_name = self.STATE.cur_settings.axis
-        if axis_name is None:
-            axis_name = message.dims[0]
-        axis_idx = message.get_axis_idx(axis_name)
-        axis = message.get_axis(axis_name)
-
-        spectrum = np.moveaxis(message.data, axis_idx, -1)
-
-        n_time = message.data.shape[axis_idx]
-        window = WINDOWS[self.STATE.cur_settings.window](n_time)
-
-        spectrum = np.fft.fft(spectrum * window) / n_time
-        spectrum = np.fft.fftshift(spectrum, axes=-1)
-        freqs = np.fft.fftshift(np.fft.fftfreq(n_time, d=axis.gain), axes=-1)
-
-        if self.STATE.cur_settings.transform != SpectralTransform.RAW_COMPLEX:
-            if self.STATE.cur_settings.transform == SpectralTransform.REAL:
-                spectrum = spectrum.real
-            elif self.STATE.cur_settings.transform == SpectralTransform.IMAG:
-                spectrum = spectrum.imag
-            else:
-                scale = np.sum(window**2.0) * axis.gain
-                spectrum = (2.0 * (np.abs(spectrum) ** 2.0)) / scale
-
-                if self.STATE.cur_settings.transform == SpectralTransform.REL_DB:
-                    spectrum = 10 * np.log10(spectrum)
-
-        axis_offset = freqs[0]
-        if self.STATE.cur_settings.output == SpectralOutput.POSITIVE:
-            axis_offset = freqs[n_time // 2]
-            spectrum = spectrum[..., n_time // 2 :]
-        elif self.STATE.cur_settings.output == SpectralOutput.NEGATIVE:
-            spectrum = spectrum[..., : n_time // 2]
-
-        spectrum = np.moveaxis(spectrum, axis_idx, -1)
-
-        out_axis = self.SETTINGS.out_axis
-        if out_axis is None:
-            out_axis = axis_name
-
-        freq_axis = AxisArray.Axis(
-            unit="Hz", gain=1.0 / (axis.gain * n_time), offset=axis_offset
-        )
-        new_axes = {**message.axes, **{out_axis: freq_axis}}
-
-        new_dims = [d for d in message.dims]
-        if self.SETTINGS.out_axis is not None:
-            new_dims[axis_idx] = self.SETTINGS.out_axis
-
-        out_msg = replace(message, data=spectrum, dims=new_dims, axes=new_axes)
-
-        yield self.OUTPUT_SIGNAL, out_msg
+from .spectrum import (
+    OptionsEnum,
+    WindowFunction,
+    SpectralTransform,
+    SpectralOutput,
+    SpectrumSettings,
+    SpectrumState,
+    Spectrum
+)

ezmsg/sigproc/spectrogram.py

@@ -0,0 +1,68 @@
+import typing
+
+import numpy as np
+
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer, GenAxisArray  # , compose
+from ezmsg.util.messages.modify import modify_axis
+from ezmsg.sigproc.window import windowing
+from ezmsg.sigproc.spectrum import (
+    spectrum,
+    WindowFunction, SpectralTransform, SpectralOutput
+)
+
+
+@consumer
+def spectrogram(
+    window_dur: typing.Optional[float] = None,
+    window_shift: typing.Optional[float] = None,
+    window: WindowFunction = WindowFunction.HANNING,
+    transform: SpectralTransform = SpectralTransform.REL_DB,
+    output: SpectralOutput = SpectralOutput.POSITIVE
+) -> typing.Generator[typing.Optional[AxisArray], AxisArray, None]:
+
+    # We cannot use `compose` because `windowing` returns a list of axisarray objects,
+    #  even though the length is always exactly 1 for the settings used here.
+    # pipeline = compose(
+    f_win = windowing(axis="time", newaxis="step", window_dur=window_dur, window_shift=window_shift)
+    f_spec = spectrum(axis="time", window=window, transform=transform, output=output)
+    f_modify = modify_axis(name_map={"step": "time"})
+    # )
+
+    # State variables
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out: typing.Optional[AxisArray] = None
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+
+        # axis_arr_out = pipeline(axis_arr_in)
+        axis_arr_out = None
+        wins = f_win.send(axis_arr_in)
+        if len(wins):
+            specs = f_spec.send(wins[0])
+            if specs is not None:
+                axis_arr_out = f_modify.send(specs)
+
+
+class SpectrogramSettings(ez.Settings):
+    window_dur: typing.Optional[float] = None  # window duration in seconds
+    window_shift: typing.Optional[float] = None  # window step in seconds. If None, window_shift == window_dur
+    # See SpectrumSettings for details of following settings:
+    window: WindowFunction = WindowFunction.HAMMING
+    transform: SpectralTransform = SpectralTransform.REL_DB
+    output: SpectralOutput = SpectralOutput.POSITIVE
+
+
+class Spectrogram(GenAxisArray):
+    SETTINGS: SpectrogramSettings
+
+    def construct_generator(self):
+        self.STATE.gen = spectrogram(
+            window_dur=self.SETTINGS.window_dur,
+            window_shift=self.SETTINGS.window_shift,
+            window=self.SETTINGS.window,
+            transform=self.SETTINGS.transform,
+            output=self.SETTINGS.output
+        )

ezmsg/sigproc/spectrum.py

@@ -0,0 +1,158 @@
+from dataclasses import replace
+import enum
+from typing import Optional, Generator, AsyncGenerator
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.generator import consumer, GenAxisArray
+
+
+class OptionsEnum(enum.Enum):
+    @classmethod
+    def options(cls):
+        return list(map(lambda c: c.value, cls))
+
+
+class WindowFunction(OptionsEnum):
+    NONE = "None (Rectangular)"
+    HAMMING = "Hamming"
+    HANNING = "Hanning"
+    BARTLETT = "Bartlett"
+    BLACKMAN = "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):
+    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):
+    FULL = "Full Spectrum"
+    POSITIVE = "Positive Frequencies"
+    NEGATIVE = "Negative Frequencies"
+
+
+@consumer
+def spectrum(
+    axis: Optional[str] = None,
+    out_axis: Optional[str] = "freq",
+    window: WindowFunction = WindowFunction.HANNING,
+    transform: SpectralTransform = SpectralTransform.REL_DB,
+    output: SpectralOutput = SpectralOutput.POSITIVE
+) -> Generator[AxisArray, AxisArray, None]:
+
+    # State variables
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = AxisArray(np.array([]), dims=[""])
+
+    axis_name = axis
+    axis_idx = None
+    n_time = None
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+
+        if axis_name is None:
+            axis_name = axis_arr_in.dims[0]
+
+        # Initial setup
+        if n_time is None or axis_idx is None or axis_arr_in.data.shape[axis_idx] != n_time:
+            axis_idx = axis_arr_in.get_axis_idx(axis_name)
+            _axis = axis_arr_in.get_axis(axis_name)
+            n_time = axis_arr_in.data.shape[axis_idx]
+            freqs = np.fft.fftshift(np.fft.fftfreq(n_time, d=_axis.gain), axes=-1)
+            window = WINDOWS[window](n_time)
+            window = window.reshape([1] * axis_idx + [len(window),] + [1] * (axis_arr_in.data.ndim-2))
+            if (transform != SpectralTransform.RAW_COMPLEX and
+                    not (transform == SpectralTransform.REAL or transform == SpectralTransform.IMAG)):
+                scale = np.sum(window ** 2.0) * _axis.gain
+            axis_offset = freqs[0]
+            if output == SpectralOutput.POSITIVE:
+                axis_offset = freqs[n_time // 2]
+            freq_axis = AxisArray.Axis(
+                unit="Hz", gain=1.0 / (_axis.gain * n_time), offset=axis_offset
+            )
+            if out_axis is None:
+                out_axis = axis_name
+            new_dims = axis_arr_in.dims[:axis_idx] + [out_axis, ] + axis_arr_in.dims[axis_idx + 1:]
+
+            f_transform = lambda x: x
+            if transform != SpectralTransform.RAW_COMPLEX:
+                if transform == SpectralTransform.REAL:
+                    f_transform = lambda x: x.real
+                elif transform == SpectralTransform.IMAG:
+                    f_transform = lambda x: x.imag
+                else:
+                    f1 = lambda x: (2.0 * (np.abs(x) ** 2.0)) / scale
+                    if transform == SpectralTransform.REL_DB:
+                        f_transform = lambda x: 10 * np.log10(f1(x))
+                    else:
+                        f_transform = f1
+
+        new_axes = {**axis_arr_in.axes, **{out_axis: freq_axis}}
+        if out_axis != axis_name:
+            new_axes.pop(axis_name, None)
+
+        spec = np.fft.fft(axis_arr_in.data * window, axis=axis_idx) / n_time
+        spec = np.fft.fftshift(spec, axes=axis_idx)
+        spec = f_transform(spec)
+
+        if output == SpectralOutput.POSITIVE:
+            spec = slice_along_axis(spec, slice(n_time // 2, None), axis_idx)
+
+        elif output == SpectralOutput.NEGATIVE:
+            spec = slice_along_axis(spec, slice(None, n_time // 2), axis_idx)
+
+        axis_arr_out = replace(axis_arr_in, data=spec, dims=new_dims, axes=new_axes)
+
+
+class SpectrumSettings(ez.Settings):
+    axis: Optional[str] = None
+    # n: Optional[int] = None # n parameter for fft
+    out_axis: Optional[str] = "freq"  # If none; don't change dim name
+    window: WindowFunction = WindowFunction.HAMMING
+    transform: SpectralTransform = SpectralTransform.REL_DB
+    output: SpectralOutput = SpectralOutput.POSITIVE
+
+
+class SpectrumState(ez.State):
+    gen: Generator
+    cur_settings: SpectrumSettings
+
+
+class Spectrum(GenAxisArray):
+    SETTINGS: SpectrumSettings
+    STATE: SpectrumState
+
+    INPUT_SETTINGS = ez.InputStream(SpectrumSettings)
+
+    def initialize(self) -> None:
+        self.STATE.cur_settings = self.SETTINGS
+        super().initialize()
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: SpectrumSettings):
+        self.STATE.cur_settings = msg
+        self.construct_generator()
+
+    def construct_generator(self):
+        self.STATE.gen = spectrum(
+            axis=self.STATE.cur_settings.axis,
+            out_axis=self.STATE.cur_settings.out_axis,
+            window=self.STATE.cur_settings.window,
+            transform=self.STATE.cur_settings.transform,
+            output=self.STATE.cur_settings.output
+        )