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

ezmsg/sigproc/synth.py

@@ -1,21 +1,27 @@
 import asyncio
-from collections import deque
-from dataclasses import dataclass, replace, field
+from dataclasses import replace, field
 import time
 from typing import Optional, Generator, AsyncGenerator, Union
 
 import numpy as np
 import ezmsg.core as ez
-from ezmsg.util.generator import consumer, GenAxisArray
+from ezmsg.util.generator import consumer
 from ezmsg.util.messages.axisarray import AxisArray
 
 from .butterworthfilter import ButterworthFilter, ButterworthFilterSettings
+from .base import GenAxisArray
 
 
-# CLOCK -- generate events at a specified rate #
-def clock(
-    dispatch_rate: Optional[float]
-) -> Generator[ez.Flag, None, None]:
+def clock(dispatch_rate: Optional[float]) -> Generator[ez.Flag, None, None]:
+    """
+    Construct a generator that yields events at a specified rate.
+
+    Args:
+        dispatch_rate: event rate in seconds.
+
+    Returns:
+        A generator object that yields :obj:`ez.Flag` events at a specified rate.
+    """
     n_dispatch = -1
     t_0 = time.time()
     while True:
@@ -26,9 +32,13 @@ def clock(
         yield ez.Flag()
 
 
-async def aclock(
-    dispatch_rate: Optional[float]
-) -> AsyncGenerator[ez.Flag, None]:
+async def aclock(dispatch_rate: Optional[float]) -> AsyncGenerator[ez.Flag, None]:
+    """
+    ``asyncio`` version of :obj:`clock`.
+
+    Returns:
+        asynchronous generator object. Must use `anext` or `async for`.
+    """
     t_0 = time.time()
     n_dispatch = -1
     while True:
@@ -40,6 +50,8 @@ async def aclock(
 
 
 class ClockSettings(ez.Settings):
+    """Settings for :obj:`Clock`. See :obj:`clock` for parameter description."""
+
     # Message dispatch rate (Hz), or None (fast as possible)
     dispatch_rate: Optional[float]
 
@@ -50,13 +62,15 @@ class ClockState(ez.State):
 
 
 class Clock(ez.Unit):
-    SETTINGS: ClockSettings
-    STATE: ClockState
+    """Unit for :obj:`clock`."""
+
+    SETTINGS = ClockSettings
+    STATE = ClockState
 
     INPUT_SETTINGS = ez.InputStream(ClockSettings)
     OUTPUT_CLOCK = ez.OutputStream(ez.Flag)
 
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         self.STATE.cur_settings = self.SETTINGS
         self.construct_generator()
 
@@ -78,18 +92,33 @@ class Clock(ez.Unit):
 
 # COUNTER - Generate incrementing integer. fs and dispatch_rate parameters combine to give many options. #
 async def acounter(
-    n_time: int,  # Number of samples to output per block
-    fs: Optional[float],  # Sampling rate of signal output in Hz
-    n_ch: int = 1,  # Number of channels to synthesize
-
-    # Message dispatch rate (Hz), 'realtime' or None (fast as possible)
-    #  Note: if dispatch_rate is a float then time offsets will be synthetic and the
-    #  system will run faster or slower than wall clock time.
+    n_time: int,
+    fs: Optional[float],
+    n_ch: int = 1,
     dispatch_rate: Optional[Union[float, str]] = None,
-
-    # If set to an integer, counter will rollover at this number.
     mod: Optional[int] = None,
 ) -> AsyncGenerator[AxisArray, None]:
+    """
+    Construct an asynchronous generator to generate AxisArray objects at a specified rate
+    and with the specified sampling rate.
+
+    NOTE: This module uses asyncio.sleep to delay appropriately in realtime mode.
+    This method of sleeping/yielding execution priority has quirky behavior with
+    sub-millisecond sleep periods which may result in unexpected behavior (e.g.
+    fs = 2000, n_time = 1, realtime = True -- may result in ~1400 msgs/sec)
+
+    Args:
+        n_time: Number of samples to output per block.
+        fs: Sampling rate of signal output in Hz.
+        n_ch: Number of channels to synthesize
+        dispatch_rate: Message dispatch rate (Hz), 'realtime' or None (fast as possible)
+            Note: if dispatch_rate is a float then time offsets will be synthetic and the
+            system will run faster or slower than wall clock time.
+        mod: If set to an integer, counter will rollover at this number.
+
+    Returns:
+        An asynchronous generator.
+    """
 
     # TODO: Adapt this to use ezmsg.util.rate?
 
@@ -150,12 +179,10 @@ async def acounter(
 
 
 class CounterSettings(ez.Settings):
+    # TODO: Adapt this to use ezmsg.util.rate?
     """
-    TODO: Adapt this to use ezmsg.util.rate?
-    NOTE: This module uses asyncio.sleep to delay appropriately in realtime mode.
-    This method of sleeping/yielding execution priority has quirky behavior with
-    sub-millisecond sleep periods which may result in unexpected behavior (e.g.
-    fs = 2000, n_time = 1, realtime = True -- may result in ~1400 msgs/sec)
+    Settings for :obj:`Counter`.
+    See :obj:`acounter` for a description of the parameters.
     """
 
     n_time: int  # Number of samples to output per block
@@ -178,10 +205,10 @@ class CounterState(ez.State):
 
 
 class Counter(ez.Unit):
-    """Generates monotonically increasing counter"""
+    """Generates monotonically increasing counter. Unit for :obj:`acounter`."""
 
-    SETTINGS: CounterSettings
-    STATE: CounterState
+    SETTINGS = CounterSettings
+    STATE = CounterState
 
     INPUT_CLOCK = ez.InputStream(ez.Flag)
     INPUT_SETTINGS = ez.InputStream(CounterSettings)
@@ -209,27 +236,26 @@ class Counter(ez.Unit):
             self.STATE.cur_settings.fs,
             n_ch=self.STATE.cur_settings.n_ch,
             dispatch_rate=self.STATE.cur_settings.dispatch_rate,
-            mod=self.STATE.cur_settings.mod
+            mod=self.STATE.cur_settings.mod,
         )
         self.STATE.new_generator.set()
-    
+
     @ez.subscriber(INPUT_CLOCK)
     @ez.publisher(OUTPUT_SIGNAL)
     async def on_clock(self, clock: ez.Flag):
-        if self.STATE.cur_settings.dispatch_rate == 'ext_clock':
+        if self.STATE.cur_settings.dispatch_rate == "ext_clock":
             out = await self.STATE.gen.__anext__()
             yield self.OUTPUT_SIGNAL, out
 
     @ez.publisher(OUTPUT_SIGNAL)
     async def run_generator(self) -> AsyncGenerator:
         while True:
-    
             await self.STATE.new_generator.wait()
             self.STATE.new_generator.clear()
-                
-            if self.STATE.cur_settings.dispatch_rate == 'ext_clock':
+
+            if self.STATE.cur_settings.dispatch_rate == "ext_clock":
                 continue
-            
+
             while not self.STATE.new_generator.is_set():
                 out = await self.STATE.gen.__anext__()
                 yield self.OUTPUT_SIGNAL, out
@@ -238,29 +264,46 @@ class Counter(ez.Unit):
 @consumer
 def sin(
     axis: Optional[str] = "time",
-    freq: float = 1.0,  # Oscillation frequency in Hz
-    amp: float = 1.0,  # Amplitude
-    phase: float = 0.0,  # Phase offset (in radians)
+    freq: float = 1.0,
+    amp: float = 1.0,
+    phase: float = 0.0,
 ) -> Generator[AxisArray, AxisArray, None]:
-    axis_arr_in = AxisArray(np.array([]), dims=[""])
-    axis_arr_out = AxisArray(np.array([]), dims=[""])
+    """
+    Construct a generator of sinusoidal waveforms in AxisArray objects.
+
+    Args:
+        axis: The name of the axis over which the sinusoid passes.
+        freq: The frequency of the sinusoid, in Hz.
+        amp: The amplitude of the sinusoid.
+        phase: The initial phase of the sinusoid, in radians.
+
+    Returns:
+        A primed generator that expects .send(axis_array) of sample counts
+        and yields an AxisArray of sinusoids.
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
 
     ang_freq = 2.0 * np.pi * freq
 
     while True:
-        axis_arr_in = yield axis_arr_out
-        # axis_arr_in is expected to be sample counts
+        msg_in: AxisArray = yield msg_out
+        # msg_in is expected to be sample counts
 
         axis_name = axis
         if axis_name is None:
-            axis_name = axis_arr_in.dims[0]
+            axis_name = msg_in.dims[0]
 
-        w = (ang_freq * axis_arr_in.get_axis(axis_name).gain) * axis_arr_in.data
+        w = (ang_freq * msg_in.get_axis(axis_name).gain) * msg_in.data
         out_data = amp * np.sin(w + phase)
-        axis_arr_out = replace(axis_arr_in, data=out_data)
+        msg_out = replace(msg_in, data=out_data)
 
 
 class SinGeneratorSettings(ez.Settings):
+    """
+    Settings for :obj:`SinGenerator`.
+    See :obj:`sin` for parameter descriptions.
+    """
+
     time_axis: Optional[str] = "time"
     freq: float = 1.0  # Oscillation frequency in Hz
     amp: float = 1.0  # Amplitude
@@ -268,30 +311,55 @@ class SinGeneratorSettings(ez.Settings):
 
 
 class SinGenerator(GenAxisArray):
-    SETTINGS: SinGeneratorSettings
+    """
+    Unit for :obj:`sin`.
+    """
+
+    SETTINGS = SinGeneratorSettings
 
     def construct_generator(self):
         self.STATE.gen = sin(
             axis=self.SETTINGS.time_axis,
             freq=self.SETTINGS.freq,
             amp=self.SETTINGS.amp,
-            phase=self.SETTINGS.phase
+            phase=self.SETTINGS.phase,
         )
 
 
 class OscillatorSettings(ez.Settings):
-    n_time: int  # Number of samples to output per block
-    fs: float  # Sampling rate of signal output in Hz
-    n_ch: int = 1  # Number of channels to output per block
-    dispatch_rate: Optional[Union[float, str]] = None  # (Hz) | 'realtime' | 'ext_clock'
-    freq: float = 1.0  # Oscillation frequency in Hz
-    amp: float = 1.0  # Amplitude
-    phase: float = 0.0  # Phase offset (in radians)
-    sync: bool = False  # Adjust `freq` to sync with sampling rate
+    """Settings for :obj:`Oscillator`"""
+
+    n_time: int
+    """Number of samples to output per block."""
+
+    fs: float
+    """Sampling rate of signal output in Hz"""
+
+    n_ch: int = 1
+    """Number of channels to output per block"""
+
+    dispatch_rate: Optional[Union[float, str]] = None
+    """(Hz) | 'realtime' | 'ext_clock'"""
+
+    freq: float = 1.0
+    """Oscillation frequency in Hz"""
+
+    amp: float = 1.0
+    """Amplitude"""
+
+    phase: float = 0.0
+    """Phase offset (in radians)"""
+
+    sync: bool = False
+    """Adjust `freq` to sync with sampling rate"""
 
 
 class Oscillator(ez.Collection):
-    SETTINGS: OscillatorSettings
+    """
+    :obj:`Collection that chains :obj:`Counter` and :obj:`SinGenerator`.
+    """
+
+    SETTINGS = OscillatorSettings
 
     INPUT_CLOCK = ez.InputStream(ez.Flag)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
@@ -334,11 +402,18 @@ class Oscillator(ez.Collection):
 
 class RandomGeneratorSettings(ez.Settings):
     loc: float = 0.0
+    """loc argument for :obj:`numpy.random.normal`"""
+
     scale: float = 1.0
+    """scale argument for :obj:`numpy.random.normal`"""
 
 
 class RandomGenerator(ez.Unit):
-    SETTINGS: RandomGeneratorSettings
+    """
+    Replaces input data with random data and yields the result.
+    """
+
+    SETTINGS = RandomGeneratorSettings
 
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
@@ -354,12 +429,15 @@ class RandomGenerator(ez.Unit):
 
 
 class NoiseSettings(ez.Settings):
+    """
+    See :obj:`CounterSettings` and :obj:`RandomGeneratorSettings`.
+    """
+
     n_time: int  # Number of samples to output per block
     fs: float  # Sampling rate of signal output in Hz
     n_ch: int = 1  # Number of channels to output
-    dispatch_rate: Optional[
-        Union[float, str]
-    ] = None  # (Hz), 'realtime', or 'ext_clock'
+    dispatch_rate: Optional[Union[float, str]] = None
+    """(Hz), 'realtime', or 'ext_clock'"""
     loc: float = 0.0  # DC offset
     scale: float = 1.0  # Scale (in standard deviations)
 
@@ -368,7 +446,11 @@ WhiteNoiseSettings = NoiseSettings
 
 
 class WhiteNoise(ez.Collection):
-    SETTINGS: NoiseSettings
+    """
+    A :obj:`Collection` that chains a :obj:`Counter` and :obj:`RandomGenerator`.
+    """
+
+    SETTINGS = NoiseSettings
 
     INPUT_CLOCK = ez.InputStream(ez.Flag)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
@@ -403,7 +485,11 @@ PinkNoiseSettings = NoiseSettings
 
 
 class PinkNoise(ez.Collection):
-    SETTINGS: PinkNoiseSettings
+    """
+    A :obj:`Collection` that chains :obj:`WhiteNoise` and :obj:`ButterworthFilter`.
+    """
+
+    SETTINGS = PinkNoiseSettings
 
     INPUT_CLOCK = ez.InputStream(ez.Flag)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
@@ -415,7 +501,9 @@ class PinkNoise(ez.Collection):
         self.WHITE_NOISE.apply_settings(self.SETTINGS)
         self.FILTER.apply_settings(
             ButterworthFilterSettings(
-                axis="time", order=1, cutoff=self.SETTINGS.fs * 0.01  # Hz
+                axis="time",
+                order=1,
+                cutoff=self.SETTINGS.fs * 0.01,  # Hz
             )
         )
 
@@ -435,7 +523,7 @@ class AddState(ez.State):
 class Add(ez.Unit):
     """Add two signals together.  Assumes compatible/similar axes/dimensions."""
 
-    STATE: AddState
+    STATE = AddState
 
     INPUT_SIGNAL_A = ez.InputStream(AxisArray)
     INPUT_SIGNAL_B = ez.InputStream(AxisArray)
@@ -459,6 +547,8 @@ class Add(ez.Unit):
 
 
 class EEGSynthSettings(ez.Settings):
+    """See :obj:`OscillatorSettings`."""
+
     fs: float = 500.0  # Hz
     n_time: int = 100
     alpha_freq: float = 10.5  # Hz
@@ -466,7 +556,12 @@ class EEGSynthSettings(ez.Settings):
 
 
 class EEGSynth(ez.Collection):
-    SETTINGS: EEGSynthSettings
+    """
+    A :obj:`Collection` that chains a :obj:`Clock` to both :obj:`PinkNoise`
+    and :obj:`Oscillator`, then :obj:`Add` s the result.
+    """
+
+    SETTINGS = EEGSynthSettings
 
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 

ezmsg/sigproc/wavelets.py

@@ -0,0 +1,167 @@
+from dataclasses import replace
+import typing
+
+import numpy as np
+import numpy.typing as npt
+import pywt
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
+
+from .base import GenAxisArray
+from .filterbank import filterbank, FilterbankMode, MinPhaseMode
+
+
+@consumer
+def cwt(
+    scales: typing.Union[list, tuple, npt.NDArray],
+    wavelet: typing.Union[str, pywt.ContinuousWavelet, pywt.Wavelet],
+    min_phase: MinPhaseMode = MinPhaseMode.NONE,
+    axis: str = "time",
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Build a generator to perform a continuous wavelet transform on sent AxisArray messages.
+    The function is equivalent to the `pywt.cwt` function, but is designed to work with streaming data.
+
+    Args:
+        scales: The wavelet scales to use.
+        wavelet: Wavelet object or name of wavelet to use.
+        min_phase: See filterbank MinPhaseMode for details.
+        axis: The target axis for operation. Note that this will be moved to the -1th dimension
+          because fft and matrix multiplication is much faster on the last axis.
+
+    Returns:
+        A Generator object that expects `.send(axis_array)` of continuous data
+    """
+    msg_out: typing.Optional[AxisArray] = None
+
+    # Check parameters
+    scales = np.array(scales)
+    assert np.all(scales > 0), "Scales must be positive."
+    assert scales.ndim == 1, "Scales must be a 1D list, tuple, or array."
+    if not isinstance(wavelet, (pywt.ContinuousWavelet, pywt.Wavelet)):
+        wavelet = pywt.DiscreteContinuousWavelet(wavelet)
+    precision = 10
+
+    # State variables
+    neg_rt_scales = -np.sqrt(scales)[:, None]
+    int_psi, wave_xvec = pywt.integrate_wavelet(wavelet, precision=precision)
+    int_psi = np.conj(int_psi) if wavelet.complex_cwt else int_psi
+    template: typing.Optional[AxisArray] = None
+    fbgen: typing.Optional[typing.Generator[AxisArray, AxisArray, None]] = None
+    last_conv_samp: typing.Optional[npt.NDArray] = None
+
+    # Reset if input changed
+    check_input = {
+        "kind": None,  # Need to recalc kernels at same complexity as input
+        "gain": None,  # Need to recalc freqs
+        "shape": None,  # Need to recalc template and buffer
+        "key": None,  # Buffer obsolete
+    }
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+        ax_idx = msg_in.get_axis_idx(axis)
+        in_shape = msg_in.data.shape[:ax_idx] + msg_in.data.shape[ax_idx + 1 :]
+
+        b_reset = msg_in.data.dtype.kind != check_input["kind"]
+        b_reset = b_reset or msg_in.axes[axis].gain != check_input["gain"]
+        b_reset = b_reset or in_shape != check_input["shape"]
+        b_reset = b_reset or msg_in.key != check_input["key"]
+        b_reset = b_reset and msg_in.data.size > 0
+        if b_reset:
+            check_input["kind"] = msg_in.data.dtype.kind
+            check_input["gain"] = msg_in.axes[axis].gain
+            check_input["shape"] = in_shape
+            check_input["key"] = msg_in.key
+
+            # convert int_psi, wave_xvec to the same precision as the data
+            dt_data = msg_in.data.dtype  # _check_dtype(msg_in.data)
+            dt_cplx = np.result_type(dt_data, np.complex64)
+            dt_psi = dt_cplx if int_psi.dtype.kind == "c" else dt_data
+            int_psi = np.asarray(int_psi, dtype=dt_psi)
+            # TODO: Currently int_psi cannot be made non-complex once it is complex.
+
+            # Calculate waves for each scale
+            wave_xvec = np.asarray(wave_xvec, dtype=msg_in.data.real.dtype)
+            wave_range = wave_xvec[-1] - wave_xvec[0]
+            step = wave_xvec[1] - wave_xvec[0]
+            int_psi_scales = []
+            for scale in scales:
+                reix = (np.arange(scale * wave_range + 1) / (scale * step)).astype(int)
+                if reix[-1] >= int_psi.size:
+                    reix = np.extract(reix < int_psi.size, reix)
+                int_psi_scales.append(int_psi[reix][::-1])
+
+            # CONV is probably best because we often get huge kernels.
+            fbgen = filterbank(
+                int_psi_scales, mode=FilterbankMode.CONV, min_phase=min_phase, axis=axis
+            )
+
+            freqs = (
+                pywt.scale2frequency(wavelet, scales, precision)
+                / msg_in.axes[axis].gain
+            )
+            fstep = (freqs[1] - freqs[0]) if len(freqs) > 1 else 1.0
+            # Create output template
+            dummy_shape = in_shape + (len(scales), 0)
+            template = AxisArray(
+                np.zeros(
+                    dummy_shape, dtype=dt_cplx if wavelet.complex_cwt else dt_data
+                ),
+                dims=msg_in.dims[:ax_idx] + msg_in.dims[ax_idx + 1 :] + ["freq", axis],
+                axes={
+                    **msg_in.axes,
+                    "freq": AxisArray.Axis("Hz", offset=freqs[0], gain=fstep),
+                },
+            )
+            last_conv_samp = np.zeros(
+                dummy_shape[:-1] + (1,), dtype=template.data.dtype
+            )
+
+        conv_msg = fbgen.send(msg_in)
+
+        # Prepend with last_conv_samp before doing diff
+        dat = np.concatenate((last_conv_samp, conv_msg.data), axis=-1)
+        coef = neg_rt_scales * np.diff(dat, axis=-1)
+        # Store last_conv_samp for next iteration.
+        last_conv_samp = conv_msg.data[..., -1:]
+
+        if template.data.dtype.kind != "c":
+            coef = coef.real
+
+        # pywt.cwt slices off the beginning and end of the result where the convolution overran. We don't have
+        #  that luxury when streaming.
+        # d = (coef.shape[-1] - msg_in.data.shape[ax_idx]) / 2.
+        # coef = coef[..., math.floor(d):-math.ceil(d)]
+        msg_out = replace(
+            template, data=coef, axes={**template.axes, axis: msg_in.axes[axis]}
+        )
+
+
+class CWTSettings(ez.Settings):
+    """
+    Settings for :obj:`CWT`
+    See :obj:`cwt` for argument details.
+    """
+
+    scales: typing.Union[list, tuple, npt.NDArray]
+    wavelet: typing.Union[str, pywt.ContinuousWavelet, pywt.Wavelet]
+    min_phase: MinPhaseMode = MinPhaseMode.NONE
+    axis: str = "time"
+
+
+class CWT(GenAxisArray):
+    """
+    :obj:`Unit` for :obj:`common_rereference`.
+    """
+
+    SETTINGS = CWTSettings
+
+    def construct_generator(self):
+        self.STATE.gen = cwt(
+            scales=self.SETTINGS.scales,
+            wavelet=self.SETTINGS.wavelet,
+            min_phase=self.SETTINGS.min_phase,
+            axis=self.SETTINGS.axis,
+        )