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

ezmsg/sigproc/synth.py

@@ -1,15 +1,42 @@
 import asyncio
-import time
+from collections import deque
 from dataclasses import dataclass, replace, field
+import time
+from typing import Optional, Generator, AsyncGenerator, Union
 
-import ezmsg.core as ez
 import numpy as np
-
+import ezmsg.core as ez
+from ezmsg.util.generator import consumer, GenAxisArray
 from ezmsg.util.messages.axisarray import AxisArray
 
 from .butterworthfilter import ButterworthFilter, ButterworthFilterSettings
 
-from typing import Optional, AsyncGenerator, Union
+
+# CLOCK -- generate events at a specified rate #
+def clock(
+    dispatch_rate: Optional[float]
+) -> Generator[ez.Flag, None, None]:
+    n_dispatch = -1
+    t_0 = time.time()
+    while True:
+        if dispatch_rate is not None:
+            n_dispatch += 1
+            t_next = t_0 + n_dispatch / dispatch_rate
+            time.sleep(max(0, t_next - time.time()))
+        yield ez.Flag()
+
+
+async def aclock(
+    dispatch_rate: Optional[float]
+) -> AsyncGenerator[ez.Flag, None]:
+    t_0 = time.time()
+    n_dispatch = -1
+    while True:
+        if dispatch_rate is not None:
+            n_dispatch += 1
+            t_next = t_0 + n_dispatch / dispatch_rate
+            await asyncio.sleep(t_next - time.time())
+        yield ez.Flag()
 
 
 class ClockSettings(ez.Settings):
@@ -19,6 +46,7 @@ class ClockSettings(ez.Settings):
 
 class ClockState(ez.State):
     cur_settings: ClockSettings
+    gen: AsyncGenerator
 
 
 class Clock(ez.Unit):
@@ -30,17 +58,95 @@ class Clock(ez.Unit):
 
     def initialize(self) -> None:
         self.STATE.cur_settings = self.SETTINGS
+        self.construct_generator()
+
+    def construct_generator(self):
+        self.STATE.gen = aclock(self.STATE.cur_settings.dispatch_rate)
 
     @ez.subscriber(INPUT_SETTINGS)
     async def on_settings(self, msg: ClockSettings) -> None:
         self.STATE.cur_settings = msg
+        self.construct_generator()
 
     @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
+            out = await self.STATE.gen.__anext__()
+            if out:
+                yield self.OUTPUT_CLOCK, out
+
+
+# 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.
+    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]:
+
+    # TODO: Adapt this to use ezmsg.util.rate?
+
+    counter_start: int = 0  # next sample's first value
+
+    b_realtime = False
+    b_manual_dispatch = False
+    b_ext_clock = False
+    if dispatch_rate is not None:
+        if isinstance(dispatch_rate, str):
+            if dispatch_rate.lower() == "realtime":
+                b_realtime = True
+            elif dispatch_rate.lower() == "ext_clock":
+                b_ext_clock = True
+        else:
+            b_manual_dispatch = True
+
+    n_sent: int = 0  # It is convenient to know how many samples we have sent.
+    clock_zero: float = time.time()  # time associated with first sample
+
+    while True:
+        # 1. Sleep, if necessary, until we are at the end of the current block
+        if b_realtime:
+            n_next = n_sent + n_time
+            t_next = clock_zero + n_next / fs
+            await asyncio.sleep(t_next - time.time())
+        elif b_manual_dispatch:
+            n_disp_next = 1 + n_sent / n_time
+            t_disp_next = clock_zero + n_disp_next / dispatch_rate
+            await asyncio.sleep(t_disp_next - time.time())
+
+        # 2. Prepare counter data.
+        block_samp = np.arange(counter_start, counter_start + n_time)[:, np.newaxis]
+        if mod is not None:
+            block_samp %= mod
+        block_samp = np.tile(block_samp, (1, n_ch))
+
+        # 3. Prepare offset - the time associated with block_samp[0]
+        if b_realtime:
+            offset = t_next - n_time / fs
+        elif b_ext_clock:
+            offset = time.time()
+        else:
+            # Purely synthetic.
+            offset = n_sent / fs
+            # offset += clock_zero  # ??
+
+        # 4. yield output
+        yield AxisArray(
+            block_samp,
+            dims=["time", "ch"],
+            axes={"time": AxisArray.Axis.TimeAxis(fs=fs, offset=offset)},
+        )
+
+        # 5. Update state for next iteration (after next yield)
+        counter_start = block_samp[-1, 0] + 1  # do not % mod
+        n_sent += n_time
 
 
 class CounterSettings(ez.Settings):
@@ -57,6 +163,8 @@ class CounterSettings(ez.Settings):
     n_ch: int = 1  # Number of channels to synthesize
 
     # Message dispatch rate (Hz), 'realtime', 'ext_clock', 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.
     dispatch_rate: Optional[Union[float, str]] = None
 
     # If set to an integer, counter will rollover
@@ -64,9 +172,9 @@ class CounterSettings(ez.Settings):
 
 
 class CounterState(ez.State):
+    gen: AsyncGenerator[AxisArray, Optional[ez.Flag]]
     cur_settings: CounterSettings
-    samp: int = 0  # current sample counter
-    clock_event: asyncio.Event
+    new_generator: asyncio.Event
 
 
 class Counter(ez.Unit):
@@ -79,9 +187,8 @@ class Counter(ez.Unit):
     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()
+    async def initialize(self) -> None:
+        self.STATE.new_generator = asyncio.Event()
         self.validate_settings(self.SETTINGS)
 
     @ez.subscriber(INPUT_SETTINGS)
@@ -93,53 +200,64 @@ class Counter(ez.Unit):
             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
-
+        self.construct_generator()
+
+    def construct_generator(self):
+        self.STATE.gen = acounter(
+            self.STATE.cur_settings.n_time,
+            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
+        )
+        self.STATE.new_generator.set()
+    
     @ez.subscriber(INPUT_CLOCK)
-    async def on_clock(self, _: ez.Flag):
-        self.STATE.clock_event.set()
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def on_clock(self, clock: ez.Flag):
+        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 publish(self) -> AsyncGenerator:
+    async def run_generator(self) -> AsyncGenerator:
         while True:
-            block_dur = self.STATE.cur_settings.n_time / self.STATE.cur_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)
-
-            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.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))
-
-            offset_adj = self.STATE.cur_settings.n_time / self.STATE.cur_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
-                    )
-                ),
-            )
+    
+            await self.STATE.new_generator.wait()
+            self.STATE.new_generator.clear()
+                
+            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
+
+
+@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)
+) -> Generator[AxisArray, AxisArray, None]:
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_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
+
+        axis_name = axis
+        if axis_name is None:
+            axis_name = axis_arr_in.dims[0]
 
-            yield self.OUTPUT_SIGNAL, out
+        w = (ang_freq * axis_arr_in.get_axis(axis_name).gain) * axis_arr_in.data
+        out_data = amp * np.sin(w + phase)
+        axis_arr_out = replace(axis_arr_in, data=out_data)
 
 
 class SinGeneratorSettings(ez.Settings):
@@ -149,35 +267,16 @@ class SinGeneratorSettings(ez.Settings):
     phase: float = 0.0  # Phase offset (in radians)
 
 
-class SinGeneratorState(ez.State):
-    ang_freq: float  # pre-calculated angular frequency in radians
-
-
-class SinGenerator(ez.Unit):
+class SinGenerator(GenAxisArray):
     SETTINGS: SinGeneratorSettings
-    STATE: SinGeneratorState
-
-    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: AxisArray) -> AsyncGenerator:
-        """
-        msg is assumed to be a monotonically increasing counter ..
-        .. or at least a counter with an intelligently chosen modulus
-        """
-        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))
+    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
+        )
 
 
 class OscillatorSettings(ez.Settings):

ezmsg/sigproc/window.py

@@ -1,33 +1,207 @@
 from dataclasses import replace
+import traceback
+from typing import AsyncGenerator, Optional, Tuple, List, Generator
 
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
 
-from ezmsg.util.messages.axisarray import AxisArray
-
-from typing import AsyncGenerator, Optional, Tuple, List
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis, sliding_win_oneaxis
+from ezmsg.util.generator import consumer
+
+
+@consumer
+def windowing(
+    axis: Optional[str] = None,
+    newaxis: Optional[str] = None,
+    window_dur: Optional[float] = None,
+    window_shift: Optional[float] = None,
+    zero_pad_until: str = "input"
+) -> Generator[AxisArray, List[AxisArray], None]:
+    """
+    Window function that generates windows of data from an input `AxisArray`.
+    :param axis: The axis along which to segment windows.
+        If None, defaults to the first dimension of the first seen AxisArray.
+    :param newaxis: Optional new axis for the output. If None, no new axes will be added.
+        If a string, windows will be stacked in a new axis with key `newaxis`, immediately preceding the windowed axis.
+    :param window_dur: The duration of the window in seconds.
+        If None, the function acts as a passthrough and all other parameters are ignored.
+    :param window_shift: The shift of the window in seconds.
+        If None (default), windowing operates in "1:1 mode", where each input yields exactly one most-recent window.
+    :param zero_pad_until: Determines how the function initializes the buffer.
+        Can be one of "input" (default), "full", "shift", or "none". If `window_shift` is None then this field is
+        ignored and "input" is always used.
+        "input" (default) initializes the buffer with the input then prepends with zeros to the window size.
+            The first input will always yield at least one output.
+        "shift" fills the buffer until `window_shift`.
+            No outputs will be yielded until at least `window_shift` data has been seen.
+        "none" does not pad the buffer. No outputs will be yielded until at least `window_dur` data has been seen.
+    :return:
+        A (primed) generator that accepts .send(an AxisArray object) and yields a list of windowed
+        AxisArray objects. The list will always be length-1 if `newaxis` is not None or `window_shift` is None.
+    """
+    # TODO: The return should be an AxisArray. i.e., always add a new axis. The Unit can do a multi-yield-per-pub
+    #  if the parameterization does not expect a newaxis.
+
+    if window_shift is None and zero_pad_until != "input":
+        ez.logger.warning("`zero_pad_until` must be 'input' if `window_shift` is None. "
+                          f"Ignoring received argument value: {zero_pad_until}")
+        zero_pad_until = "input"
+    elif window_shift is not None and zero_pad_until == "input":
+        ez.logger.warning("windowing is non-deterministic with `zero_pad_until='input'` as it depends on the size "
+                          "of the first input. We recommend using 'shift' when `window_shift` is float-valued.")
+    axis_arr_in = AxisArray(np.array([]), dims=[""])
+    axis_arr_out = [AxisArray(np.array([]), dims=[""])]
+
+    # State variables
+    prev_samp_shape: Optional[Tuple[int, ...]] = None
+    prev_fs: Optional[float] = None
+    buffer: Optional[npt.NDArray] = None
+    window_samples: Optional[int] = None
+    window_shift_samples: Optional[int] = None
+    shift_deficit: int = 0  # Number of incoming samples to ignore. Only relevant when shift > window.
+    newaxis_warn_flag: bool = False
+    mod_ax: Optional[str] = None  # The key of the modified axis in the output's .axes
+    out_template: Optional[AxisArray] = None  # Template for building return values.
+
+    while True:
+        axis_arr_in = yield axis_arr_out
+
+        if window_dur is None:
+            axis_arr_out = [axis_arr_in]
+            continue
+
+        if axis is None:
+            axis = axis_arr_in.dims[0]
+        axis_idx = axis_arr_in.get_axis_idx(axis)
+        axis_info = axis_arr_in.get_axis(axis)
+        fs = 1.0 / axis_info.gain
+
+        if (not newaxis_warn_flag) and newaxis is not None and newaxis in axis_arr_in.dims:
+            ez.logger.warning(f"newaxis {newaxis} present in input dims and will be ignored.")
+            newaxis_warn_flag = True
+        b_newaxis = newaxis is not None and newaxis not in axis_arr_in.dims
+
+        samp_shape = axis_arr_in.data.shape[:axis_idx] + axis_arr_in.data.shape[axis_idx + 1:]
+        window_samples = int(window_dur * fs)
+        b_1to1 = window_shift is None
+        if not b_1to1:
+            window_shift_samples = int(window_shift * fs)
+
+        # If buffer unset or input stats changed, create a new buffer
+        if buffer is None or samp_shape != prev_samp_shape or fs != prev_fs:
+            if zero_pad_until == "none":
+                req_samples = window_samples
+            elif zero_pad_until == "shift" and not b_1to1:
+                req_samples = window_shift_samples
+            else:  # i.e. zero_pad_until == "input"
+                req_samples = axis_arr_in.data.shape[axis_idx]
+            n_zero = max(0, window_samples - req_samples)
+            buffer_shape = axis_arr_in.data.shape[:axis_idx] + (n_zero,) + axis_arr_in.data.shape[axis_idx + 1:]
+            buffer = np.zeros(buffer_shape)
+            prev_samp_shape = samp_shape
+            prev_fs = fs
+
+        # Add new data to buffer.
+        # 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
+        # is generally faster than np.roll and slicing anyway, but this could still
+        # be a performance bottleneck for large memory arrays.
+        buffer = np.concatenate((buffer, axis_arr_in.data), axis=axis_idx)
+        # Note: if we ever move to using a circular buffer without copies then we need to create copies somewhere,
+        #  because currently the outputs are merely views into the buffer.
+
+        # Create a vector of buffer timestamps to track axis `offset` in output(s)
+        buffer_offset = np.arange(buffer.shape[axis_idx]).astype(float)
+        # Adjust so first _new_ sample at index 0
+        buffer_offset -= buffer_offset[-axis_arr_in.data.shape[axis_idx]]
+        # Convert form indices to 'units' (probably seconds).
+        buffer_offset *= axis_info.gain
+        buffer_offset += axis_info.offset
+
+        if not b_1to1 and shift_deficit > 0:
+            n_skip = min(buffer.shape[axis_idx], shift_deficit)
+            if n_skip > 0:
+                buffer = slice_along_axis(buffer, np.s_[n_skip:], axis_idx)
+                buffer_offset = buffer_offset[n_skip:]
+                shift_deficit -= n_skip
+
+        # Prepare reusable parts of output
+        if out_template is None:
+            out_dims = axis_arr_in.dims
+            if newaxis is None:
+                out_axes = {
+                    **axis_arr_in.axes,
+                    axis: replace(axis_info, offset=0.0)  # offset modified below.
+                }
+                mod_ax = axis
+            else:
+                out_dims = out_dims[:axis_idx] + [newaxis] + out_dims[axis_idx:]
+                out_axes = {
+                    **axis_arr_in.axes,
+                    newaxis: AxisArray.Axis(
+                        unit=axis_info.unit,
+                        gain=0.0 if b_1to1 else axis_info.gain * window_shift_samples,
+                        offset=0.0  # offset modified below
+                    )
+                }
+                mod_ax = newaxis
+            out_template = replace(axis_arr_in, data=np.zeros([0 for _ in out_dims]), dims=out_dims)
+
+        # Generate outputs.
+        axis_arr_out: List[AxisArray] = []
+        if b_1to1:
+            # one-to-one mode -- Each send yields exactly one window containing only the most recent samples.
+            buffer = slice_along_axis(buffer, np.s_[-window_samples:], axis_idx)
+            axis_arr_out.append(replace(
+                    out_template,
+                    data=np.expand_dims(buffer, axis=axis_idx) if b_newaxis else buffer,
+                    axes={
+                        **out_axes,
+                        mod_ax: replace(out_axes[mod_ax], offset=buffer_offset[-window_samples])
+                    }
+            ))
+        elif buffer.shape[axis_idx] >= window_samples:
+            # Deterministic window shifts.
+            win_view = sliding_win_oneaxis(buffer, window_samples, axis_idx)
+            win_view = slice_along_axis(win_view, np.s_[::window_shift_samples], axis_idx)
+            offset_view = sliding_win_oneaxis(buffer_offset, window_samples, 0)[::window_shift_samples]
+            # Place in output
+            if b_newaxis:
+                axis_arr_out.append(replace(
+                    out_template,
+                    data=win_view,
+                    axes={**out_axes, mod_ax: replace(out_axes[mod_ax], offset=offset_view[0, 0])}
+                ))
+            else:
+                for win_ix in range(win_view.shape[axis_idx]):
+                    axis_arr_out.append(replace(
+                        out_template,
+                        data=slice_along_axis(win_view, win_ix, axis_idx),
+                        axes={
+                            **out_axes,
+                            mod_ax: replace(out_axes[mod_ax], offset=offset_view[win_ix, 0])
+                        }
+                    ))
+
+            # Drop expired beginning of buffer and update shift_deficit
+            multi_shift = window_shift_samples * win_view.shape[axis_idx]
+            shift_deficit = max(0, multi_shift - buffer.shape[axis_idx])
+            buffer = slice_along_axis(buffer, np.s_[multi_shift:], axis_idx)
 
 
 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"
+    newaxis: Optional[str] = None  # new axis for output. No new axes if None
+    window_dur: Optional[float] = None  # Sec. passthrough if None
+    window_shift: Optional[float] = None  # Sec. Use "1:1 mode" if None
+    zero_pad_until: str = "full"  # "full", "shift", "input", "none"
 
 
 class WindowState(ez.State):
     cur_settings: WindowSettings
-
-    samp_shape: Optional[Tuple[int, ...]] = None  # Shape of individual sample
-    out_fs: Optional[float] = None
-    buffer: Optional[npt.NDArray] = None
-    window_samples: Optional[int] = None
-    window_shift_samples: Optional[int] = None
+    gen: Generator
 
 
 class Window(ez.Unit):
@@ -40,105 +214,33 @@ class Window(ez.Unit):
 
     def initialize(self) -> None:
         self.STATE.cur_settings = self.SETTINGS
+        self.construct_generator()
 
     @ez.subscriber(INPUT_SETTINGS)
     async def on_settings(self, msg: WindowSettings) -> None:
         self.STATE.cur_settings = msg
-        self.STATE.out_fs = None  # This should trigger a reallocation
+        self.construct_generator()
+
+    def construct_generator(self):
+        self.STATE.gen = windowing(
+            axis=self.STATE.cur_settings.axis,
+            newaxis=self.STATE.cur_settings.newaxis,
+            window_dur=self.STATE.cur_settings.window_dur,
+            window_shift=self.STATE.cur_settings.window_shift,
+            zero_pad_until=self.STATE.cur_settings.zero_pad_until
+        )
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
     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, axis_idx, 0)
-        samp_shape = time_view.shape[1:]
-
-        # 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
-            )
-
-        # 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
-        # is generally faster than np.roll and slicing anyway, but this could still
-        # be a performance bottleneck for large memory arrays.
-        self.STATE.buffer = np.concatenate((self.STATE.buffer, time_view), axis=0)
-
-        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
-
-        outputs: List[Tuple[npt.NDArray, float]] = []
-
-        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
-                )
-
-            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)
+        try:
+            # TODO: Refactor window generator so it always returns an axis array.
+            #  Then, if the configuration is such that a new "win" axis is not expected,
+            #  then iterate over the "win" axis -- dropping the "win" axis in the process.
+            out_msgs = self.STATE.gen.send(msg)
+            for out_msg in out_msgs:
+                yield self.OUTPUT_SIGNAL, out_msg
+        except (StopIteration, GeneratorExit):
+            ez.logger.debug(f"Window closed in {self.address}")
+        except Exception:
+            ez.logger.info(traceback.format_exc())