ezmsg-sigproc 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

ezmsg/sigproc/synth.py

@@ -1,31 +1,62 @@
 import asyncio
-from dataclasses import replace
+import time
+from dataclasses import dataclass, replace, field
 
 import ezmsg.core as ez
 import numpy as np
 
-from .messages import TSMessage
-from .butterworthfilter import (
-    ButterworthFilter, 
-    ButterworthFilterSettings, 
-    ButterworthFilterDesign
-)
+from ezmsg.util.messages.axisarray import AxisArray
+
+from .butterworthfilter import ButterworthFilter, ButterworthFilterSettings
 
 from typing import Optional, AsyncGenerator, Union
 
 
+class ClockSettings(ez.Settings):
+    # Message dispatch rate (Hz), or None (fast as possible)
+    dispatch_rate: Optional[float]
+
+
+class ClockState(ez.State):
+    cur_settings: ClockSettings
+
+
+class Clock(ez.Unit):
+    SETTINGS: ClockSettings
+    STATE: ClockState
+
+    INPUT_SETTINGS = ez.InputStream(ClockSettings)
+    OUTPUT_CLOCK = ez.OutputStream(ez.Flag)
+
+    def initialize(self) -> None:
+        self.STATE.cur_settings = self.SETTINGS
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: ClockSettings) -> None:
+        self.STATE.cur_settings = msg
+
+    @ez.publisher(OUTPUT_CLOCK)
+    async def generate(self) -> AsyncGenerator:
+        while True:
+            if self.STATE.cur_settings.dispatch_rate is not None:
+                await asyncio.sleep(1.0 / self.STATE.cur_settings.dispatch_rate)
+            yield self.OUTPUT_CLOCK, ez.Flag
+
+
 class CounterSettings(ez.Settings):
     """
+    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)
     """
+
     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 synthesize
+    n_ch: int = 1  # Number of channels to synthesize
 
-    # Message dispatch rate (Hz), 'realtime', or None (fast as possible)
+    # Message dispatch rate (Hz), 'realtime', 'ext_clock', or None (fast as possible)
     dispatch_rate: Optional[Union[float, str]] = None
 
     # If set to an integer, counter will rollover
@@ -33,77 +64,117 @@ class CounterSettings(ez.Settings):
 
 
 class CounterState(ez.State):
+    cur_settings: CounterSettings
     samp: int = 0  # current sample counter
+    clock_event: asyncio.Event
 
 
 class Counter(ez.Unit):
-    """ Generates monotonically increasing counter """
+    """Generates monotonically increasing counter"""
 
     SETTINGS: CounterSettings
     STATE: CounterState
 
-    OUTPUT_SIGNAL = ez.OutputStream(TSMessage)
+    INPUT_CLOCK = ez.InputStream(ez.Flag)
+    INPUT_SETTINGS = ez.InputStream(CounterSettings)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    def initialize(self) -> None:
+        self.STATE.clock_event = asyncio.Event()
+        self.STATE.clock_event.clear()
+        self.validate_settings(self.SETTINGS)
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: CounterSettings) -> None:
+        self.validate_settings(msg)
+
+    def validate_settings(self, settings: CounterSettings) -> None:
+        if isinstance(
+            settings.dispatch_rate, str
+        ) and self.SETTINGS.dispatch_rate not in ["realtime", "ext_clock"]:
+            raise ValueError(f"Unknown dispatch_rate: {self.SETTINGS.dispatch_rate}")
+
+        self.STATE.cur_settings = settings
+
+    @ez.subscriber(INPUT_CLOCK)
+    async def on_clock(self, _: ez.Flag):
+        self.STATE.clock_event.set()
 
     @ez.publisher(OUTPUT_SIGNAL)
     async def publish(self) -> AsyncGenerator:
+        while True:
+            block_dur = self.STATE.cur_settings.n_time / self.STATE.cur_settings.fs
 
-        block_samp = np.arange(self.SETTINGS.n_time)[:, np.newaxis]
-        block_dur = self.SETTINGS.n_time / self.SETTINGS.fs
+            dispatch_rate = self.STATE.cur_settings.dispatch_rate
+            if dispatch_rate is not None:
+                if isinstance(dispatch_rate, str):
+                    if dispatch_rate == "realtime":
+                        await asyncio.sleep(block_dur)
+                    elif dispatch_rate == "ext_clock":
+                        await self.STATE.clock_event.wait()
+                        self.STATE.clock_event.clear()
+                else:
+                    await asyncio.sleep(1.0 / dispatch_rate)
 
-        while True:
+            block_samp = np.arange(self.STATE.cur_settings.n_time)[:, np.newaxis]
 
             t_samp = block_samp + self.STATE.samp
             self.STATE.samp = t_samp[-1] + 1
 
-            if self.SETTINGS.mod is not None:
-                t_samp %= self.SETTINGS.mod
-                self.STATE.samp %= self.SETTINGS.mod
+            if self.STATE.cur_settings.mod is not None:
+                t_samp %= self.STATE.cur_settings.mod
+                self.STATE.samp %= self.STATE.cur_settings.mod
+
+            t_samp = np.tile(t_samp, (1, self.STATE.cur_settings.n_ch))
 
-            t_samp = np.tile( t_samp, ( 1, self.SETTINGS.n_ch ) )
+            offset_adj = self.STATE.cur_settings.n_time / self.STATE.cur_settings.fs
 
-            yield (
-                self.OUTPUT_SIGNAL,
-                TSMessage(
-                    t_samp,
-                    fs=self.SETTINGS.fs
-                )
+            out = AxisArray(
+                t_samp,
+                dims=["time", "ch"],
+                axes=dict(
+                    time=AxisArray.Axis.TimeAxis(
+                        fs=self.STATE.cur_settings.fs, offset=time.time() - offset_adj
+                    )
+                ),
             )
 
-            if self.SETTINGS.dispatch_rate is not None:
-                await asyncio.sleep(
-                    block_dur if self.SETTINGS.dispatch_rate == 'realtime'
-                    else (1.0 / self.SETTINGS.dispatch_rate)
-                )
+            yield self.OUTPUT_SIGNAL, out
 
 
 class SinGeneratorSettings(ez.Settings):
+    time_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)
 
 
 class SinGeneratorState(ez.State):
-    ang_freq: Optional[float] = None  # pre-calculated angular frequency in radians
+    ang_freq: float  # pre-calculated angular frequency in radians
 
 
 class SinGenerator(ez.Unit):
     SETTINGS: SinGeneratorSettings
     STATE: SinGeneratorState
 
-    INPUT_SIGNAL = ez.InputStream(TSMessage)
-    OUTPUT_SIGNAL = ez.OutputStream(TSMessage)
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     def initialize(self) -> None:
         self.STATE.ang_freq = 2.0 * np.pi * self.SETTINGS.freq
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
-    async def generate(self, msg: TSMessage) -> AsyncGenerator:
+    async def generate(self, msg: AxisArray) -> AsyncGenerator:
         """
         msg is assumed to be a monotonically increasing counter ..
         .. or at least a counter with an intelligently chosen modulus
         """
-        t_sec = msg.data / msg.fs
+        axis_name = self.SETTINGS.time_axis
+        if axis_name is None:
+            axis_name = msg.dims[0]
+        fs = 1.0 / msg.get_axis(axis_name).gain
+        t_sec = msg.data / fs
         w = self.STATE.ang_freq * t_sec
         out_data = self.SETTINGS.amp * np.sin(w + self.SETTINGS.phase)
         yield (self.OUTPUT_SIGNAL, replace(msg, data=out_data))
@@ -112,8 +183,8 @@ class SinGenerator(ez.Unit):
 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)
+    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)
@@ -121,16 +192,15 @@ class OscillatorSettings(ez.Settings):
 
 
 class Oscillator(ez.Collection):
-
     SETTINGS: OscillatorSettings
 
-    OUTPUT_SIGNAL = ez.OutputStream(TSMessage)
+    INPUT_CLOCK = ez.InputStream(ez.Flag)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     COUNTER = Counter()
     SIN = SinGenerator()
 
     def configure(self) -> None:
-
         # Calculate synchronous settings if necessary
         freq = self.SETTINGS.freq
         mod = None
@@ -145,54 +215,72 @@ class Oscillator(ez.Collection):
                 fs=self.SETTINGS.fs,
                 n_ch=self.SETTINGS.n_ch,
                 dispatch_rate=self.SETTINGS.dispatch_rate,
-                mod=mod
+                mod=mod,
             )
         )
 
         self.SIN.apply_settings(
             SinGeneratorSettings(
-                freq=freq,
-                amp=self.SETTINGS.amp,
-                phase=self.SETTINGS.phase
+                freq=freq, amp=self.SETTINGS.amp, phase=self.SETTINGS.phase
             )
         )
 
     def network(self) -> ez.NetworkDefinition:
         return (
+            (self.INPUT_CLOCK, self.COUNTER.INPUT_CLOCK),
             (self.COUNTER.OUTPUT_SIGNAL, self.SIN.INPUT_SIGNAL),
-            (self.SIN.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL)
+            (self.SIN.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL),
         )
 
 
+class RandomGeneratorSettings(ez.Settings):
+    loc: float = 0.0
+    scale: float = 1.0
+
+
 class RandomGenerator(ez.Unit):
-    INPUT_SIGNAL = ez.InputStream(TSMessage)
-    OUTPUT_SIGNAL = ez.OutputStream(TSMessage)
+    SETTINGS: RandomGeneratorSettings
+
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
-    async def generate(self, msg: TSMessage) -> AsyncGenerator:
-        random_data = np.random.normal(size=msg.shape)
-        yield (self.OUTPUT_SIGNAL, replace(msg, data=random_data))
+    async def generate(self, msg: AxisArray) -> AsyncGenerator:
+        random_data = np.random.normal(
+            size=msg.shape, loc=self.SETTINGS.loc, scale=self.SETTINGS.scale
+        )
+
+        yield self.OUTPUT_SIGNAL, replace(msg, data=random_data)
 
 
 class NoiseSettings(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
-    dispatch_rate: Optional[Union[float, str]] = None  # (Hz)
+    n_ch: int = 1  # Number of channels to output
+    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)
+
 
 WhiteNoiseSettings = NoiseSettings
 
-class WhiteNoise(ez.Collection):
 
+class WhiteNoise(ez.Collection):
     SETTINGS: NoiseSettings
 
-    OUTPUT_SIGNAL = ez.OutputStream(TSMessage)
+    INPUT_CLOCK = ez.InputStream(ez.Flag)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     COUNTER = Counter()
     RANDOM = RandomGenerator()
 
     def configure(self) -> None:
+        self.RANDOM.apply_settings(
+            RandomGeneratorSettings(loc=self.SETTINGS.loc, scale=self.SETTINGS.scale)
+        )
 
         self.COUNTER.apply_settings(
             CounterSettings(
@@ -200,37 +288,124 @@ class WhiteNoise(ez.Collection):
                 fs=self.SETTINGS.fs,
                 n_ch=self.SETTINGS.n_ch,
                 dispatch_rate=self.SETTINGS.dispatch_rate,
-                mod=None
+                mod=None,
             )
         )
 
     def network(self) -> ez.NetworkDefinition:
         return (
+            (self.INPUT_CLOCK, self.COUNTER.INPUT_CLOCK),
             (self.COUNTER.OUTPUT_SIGNAL, self.RANDOM.INPUT_SIGNAL),
-            (self.RANDOM.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL)
+            (self.RANDOM.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL),
         )
 
-class PinkNoise(ez.Collection):
 
-    SETTINGS: WhiteNoiseSettings
+PinkNoiseSettings = NoiseSettings
 
-    OUTPUT_SIGNAL = ez.OutputStream(TSMessage)
+
+class PinkNoise(ez.Collection):
+    SETTINGS: PinkNoiseSettings
+
+    INPUT_CLOCK = ez.InputStream(ez.Flag)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     WHITE_NOISE = WhiteNoise()
     FILTER = ButterworthFilter()
 
     def configure(self) -> None:
-
-        self.WHITE_NOISE.apply_settings( self.SETTINGS )
+        self.WHITE_NOISE.apply_settings(self.SETTINGS)
         self.FILTER.apply_settings(
             ButterworthFilterSettings(
-                order = 1,
-                cutoff = self.SETTINGS.fs * 0.01 # Hz
+                axis="time", order=1, cutoff=self.SETTINGS.fs * 0.01  # Hz
             )
         )
 
     def network(self) -> ez.NetworkDefinition:
         return (
+            (self.INPUT_CLOCK, self.WHITE_NOISE.INPUT_CLOCK),
             (self.WHITE_NOISE.OUTPUT_SIGNAL, self.FILTER.INPUT_SIGNAL),
-            (self.FILTER.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL)
-        )
+            (self.FILTER.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL),
+        )
+
+
+class AddState(ez.State):
+    queue_a: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue)
+    queue_b: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue)
+
+
+class Add(ez.Unit):
+    """Add two signals together.  Assumes compatible/similar axes/dimensions."""
+
+    STATE: AddState
+
+    INPUT_SIGNAL_A = ez.InputStream(AxisArray)
+    INPUT_SIGNAL_B = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    @ez.subscriber(INPUT_SIGNAL_A)
+    async def on_a(self, msg: AxisArray) -> None:
+        self.STATE.queue_a.put_nowait(msg)
+
+    @ez.subscriber(INPUT_SIGNAL_B)
+    async def on_b(self, msg: AxisArray) -> None:
+        self.STATE.queue_b.put_nowait(msg)
+
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def output(self) -> AsyncGenerator:
+        while True:
+            a = await self.STATE.queue_a.get()
+            b = await self.STATE.queue_b.get()
+
+            yield (self.OUTPUT_SIGNAL, replace(a, data=a.data + b.data))
+
+
+class EEGSynthSettings(ez.Settings):
+    fs: float = 500.0  # Hz
+    n_time: int = 100
+    alpha_freq: float = 10.5  # Hz
+    n_ch: int = 8
+
+
+class EEGSynth(ez.Collection):
+    SETTINGS: EEGSynthSettings
+
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    CLOCK = Clock()
+    NOISE = PinkNoise()
+    OSC = Oscillator()
+    ADD = Add()
+
+    def configure(self) -> None:
+        self.CLOCK.apply_settings(
+            ClockSettings(dispatch_rate=self.SETTINGS.fs / self.SETTINGS.n_time)
+        )
+
+        self.OSC.apply_settings(
+            OscillatorSettings(
+                n_time=self.SETTINGS.n_time,
+                fs=self.SETTINGS.fs,
+                n_ch=self.SETTINGS.n_ch,
+                dispatch_rate="ext_clock",
+                freq=self.SETTINGS.alpha_freq,
+            )
+        )
+
+        self.NOISE.apply_settings(
+            PinkNoiseSettings(
+                n_time=self.SETTINGS.n_time,
+                fs=self.SETTINGS.fs,
+                n_ch=self.SETTINGS.n_ch,
+                dispatch_rate="ext_clock",
+                scale=5.0,
+            )
+        )
+
+    def network(self) -> ez.NetworkDefinition:
+        return (
+            (self.CLOCK.OUTPUT_CLOCK, self.OSC.INPUT_CLOCK),
+            (self.CLOCK.OUTPUT_CLOCK, self.NOISE.INPUT_CLOCK),
+            (self.OSC.OUTPUT_SIGNAL, self.ADD.INPUT_SIGNAL_A),
+            (self.NOISE.OUTPUT_SIGNAL, self.ADD.INPUT_SIGNAL_B),
+            (self.ADD.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL),
+        )

ezmsg/sigproc/window.py

@@ -1,88 +1,88 @@
-from dataclasses import dataclass, replace
+from dataclasses import replace
 
 import ezmsg.core as ez
 import numpy as np
+import numpy.typing as npt
 
-from ezmsg.util.messages import view2d, shape2d
-from ezmsg.sigproc.messages import TSMessage
+from ezmsg.util.messages.axisarray import AxisArray
 
-from typing import (
-    AsyncGenerator,
-    Optional,
-    Tuple
-)
-
-@dataclass( frozen = True )
-class WindowSettingsMessage:
-    window_dur: Optional[float] = None  # Sec.  If "None" -- passthrough; window_shift is ignored.
-    window_shift: Optional[float] = None  # Sec.  If "None", activate "1:1 mode"
+from typing import AsyncGenerator, Optional, Tuple, List
 
 
-class WindowSettings(WindowSettingsMessage, ez.Settings):
-    ...
+class WindowSettings(ez.Settings):
+    axis: Optional[str] = None
+    newaxis: Optional[
+        str
+    ] = None  # Optional new axis for output.  If "None" - no new axes on output
+    window_dur: Optional[
+        float
+    ] = None  # Sec.  If "None" -- passthrough; window_shift is ignored.
+    window_shift: Optional[float] = None  # Sec.  If "None", activate "1:1 mode"
 
 
 class WindowState(ez.State):
-    cur_settings: Optional[ WindowSettingsMessage ] = None
+    cur_settings: WindowSettings
 
     samp_shape: Optional[Tuple[int, ...]] = None  # Shape of individual sample
     out_fs: Optional[float] = None
-    buffer: Optional[np.ndarray] = None
+    buffer: Optional[npt.NDArray] = None
     window_samples: Optional[int] = None
     window_shift_samples: Optional[int] = None
 
 
 class Window(ez.Unit):
-
     STATE: WindowState
     SETTINGS: WindowSettings
 
-    INPUT_SIGNAL = ez.InputStream(TSMessage)
-    OUTPUT_SIGNAL = ez.OutputStream(TSMessage)
-    INPUT_SETTINGS = ez.InputStream(WindowSettingsMessage)
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    INPUT_SETTINGS = ez.InputStream(WindowSettings)
 
     def initialize(self) -> None:
         self.STATE.cur_settings = self.SETTINGS
 
     @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: WindowSettingsMessage) -> None:
+    async def on_settings(self, msg: WindowSettings) -> None:
         self.STATE.cur_settings = msg
-        self.STATE.out_fs = None # This should trigger a reallocation
+        self.STATE.out_fs = None  # This should trigger a reallocation
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
-    async def on_signal(self, msg: TSMessage) -> AsyncGenerator:
-
+    async def on_signal(self, msg: AxisArray) -> AsyncGenerator:
         if self.STATE.cur_settings.window_dur is None:
             yield self.OUTPUT_SIGNAL, msg
             return
 
+        axis_name = self.STATE.cur_settings.axis
+        if axis_name is None:
+            axis_name = msg.dims[0]
+        axis_idx = msg.get_axis_idx(axis_name)
+        axis = msg.get_axis(axis_name)
+        fs = 1.0 / axis.gain
+
         # Create a view of data with time axis at dim 0
-        time_view = np.moveaxis(msg.data, msg.time_dim, 0)
+        time_view = np.moveaxis(msg.data, axis_idx, 0)
         samp_shape = time_view.shape[1:]
 
-        if (self.STATE.samp_shape != samp_shape) or (self.STATE.out_fs != msg.fs):
-            # Pre(re?)allocate window data
-            self.STATE.samp_shape = samp_shape
-            self.STATE.out_fs = msg.fs
-            self.STATE.window_samples = int(
-                self.STATE.cur_settings.window_dur * self.STATE.out_fs
+        # Pre(re?)allocate buffer
+        window_samples = int(self.STATE.cur_settings.window_dur * fs)
+        if (
+            (self.STATE.samp_shape != samp_shape)
+            or (self.STATE.out_fs != fs)
+            or self.STATE.buffer is None
+        ):
+            self.STATE.buffer = np.zeros(tuple([window_samples] + list(samp_shape)))
+
+        self.STATE.window_samples = window_samples
+        self.STATE.samp_shape = samp_shape
+        self.STATE.out_fs = fs
+
+        self.STATE.window_shift_samples = None
+        if self.STATE.cur_settings.window_shift is not None:
+            self.STATE.window_shift_samples = int(
+                fs * self.STATE.cur_settings.window_shift
             )
 
-            if self.STATE.cur_settings.window_shift is not None:
-                self.STATE.window_shift_samples = int(
-                    self.STATE.out_fs * self.STATE.cur_settings.window_shift
-                )
-
-            if self.STATE.buffer is None:
-                self.STATE.buffer = np.zeros(tuple([self.STATE.window_samples] + list(self.STATE.samp_shape)))
-            elif self.STATE.buffer.shape[0] > self.STATE.window_samples:
-                self.STATE.buffer = self.STATE.buffer[:self.STATE.window_samples, ...]
-            elif self.STATE.buffer.shape[0] < self.STATE.window_samples:
-                extra_samples = self.STATE.window_samples - self.STATE.buffer.shape[0]
-                extra_samples = np.zeros(tuple([extra_samples] + list(self.STATE.samp_shape)))
-                self.STATE.buffer = np.concatenate((extra_samples, self.STATE.buffer), axis=0)
-
         # Currently we just concatenate the new time samples and clip the output
         # np.roll actually returns a copy, and there's no way to construct a
         # rolling view of the data.  In current numpy implementations, np.concatenate
@@ -90,23 +90,55 @@ class Window(ez.Unit):
         # be a performance bottleneck for large memory arrays.
         self.STATE.buffer = np.concatenate((self.STATE.buffer, time_view), axis=0)
 
-        if self.STATE.window_shift_samples is None:  # one-to-one mode
-
-            self.STATE.buffer = self.STATE.buffer[-self.STATE.window_samples:, ...]
+        buffer_offset = np.arange(self.STATE.buffer.shape[0] + time_view.shape[0])
+        buffer_offset -= self.STATE.buffer.shape[0] + 1
+        buffer_offset = (buffer_offset * axis.gain) + axis.offset
 
-            # Finally, move time axis back into location before yielding
-            out_view = np.moveaxis(self.STATE.buffer, 0, msg.time_dim)
-            yield (self.OUTPUT_SIGNAL, replace(msg, data=out_view))
+        outputs: List[Tuple[npt.NDArray, float]] = []
 
-        else:  # slightly more complicated window shifting
+        if self.STATE.window_shift_samples is None:  # one-to-one mode
+            self.STATE.buffer = self.STATE.buffer[-self.STATE.window_samples :, ...]
+            buffer_offset = buffer_offset[-self.STATE.window_samples :]
+            outputs.append((self.STATE.buffer, buffer_offset[0]))
 
+        else:
             yieldable_size = self.STATE.window_samples + self.STATE.window_shift_samples
             while self.STATE.buffer.shape[0] >= yieldable_size:
+                outputs.append(
+                    (
+                        self.STATE.buffer[: self.STATE.window_samples, ...],
+                        buffer_offset[0],
+                    )
+                )
+                self.STATE.buffer = self.STATE.buffer[
+                    self.STATE.window_shift_samples :, ...
+                ]
+                buffer_offset = buffer_offset[self.STATE.window_shift_samples :]
+
+        for out_view, offset in outputs:
+            out_view = np.moveaxis(out_view, 0, axis_idx)
+
+            if (
+                self.STATE.cur_settings.newaxis is not None
+                and self.STATE.cur_settings.newaxis != self.STATE.cur_settings.axis
+            ):
+                new_gain = 0.0
+                if self.STATE.window_shift_samples is not None:
+                    new_gain = axis.gain * self.STATE.window_shift_samples
+
+                out_axis = replace(axis, unit=axis.unit, gain=new_gain, offset=offset)
+                out_axes = {**msg.axes, **{self.STATE.cur_settings.newaxis: out_axis}}
+                out_dims = [self.STATE.cur_settings.newaxis] + msg.dims
+                out_view = out_view[np.newaxis, ...]
+
+                yield self.OUTPUT_SIGNAL, replace(
+                    msg, data=out_view, dims=out_dims, axes=out_axes
+                )
 
-                # Yield if possible
-                out_view = self.STATE.buffer[:self.STATE.window_samples, ...]
-                out_view = np.moveaxis(out_view, 0, msg.time_dim)
-                yield (self.OUTPUT_SIGNAL, replace(msg, data=out_view))
-
-                # Shift window
-                self.STATE.buffer = self.STATE.buffer[self.STATE.window_shift_samples:, ...]
+            else:
+                if axis_name in msg.axes:
+                    out_axes = msg.axes
+                    out_axes[axis_name] = replace(axis, offset=offset)
+                    yield self.OUTPUT_SIGNAL, replace(msg, data=out_view, axes=out_axes)
+                else:
+                    yield self.OUTPUT_SIGNAL, replace(msg, data=out_view)

{ezmsg_sigproc-1.1.1.dist-info → ezmsg_sigproc-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ezmsg-sigproc
-Version: 1.1.1
+Version: 1.2.0
 Summary: Timeseries signal processing implementations in ezmsg
 Home-page: https://github.com/iscoe/ezmsg
 Author: Griffin Milsap
@@ -10,7 +10,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
-Requires-Dist: ezmsg (>=2.1.0)
+Requires-Dist: ezmsg (>=3.3.0)
 Requires-Dist: numpy (>=1.19.5)
 Requires-Dist: scipy (>=1.6.3)
 Provides-Extra: test