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

ezmsg/sigproc/wavelets.py

@@ -0,0 +1,187 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+import pywt
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .filterbank import FilterbankMode, MinPhaseMode, filterbank
+
+
+class CWTSettings(ez.Settings):
+    """
+    Settings for :obj:`CWT`
+    See :obj:`cwt` for argument details.
+    """
+
+    frequencies: list | tuple | npt.NDArray | None
+    wavelet: str | pywt.ContinuousWavelet | pywt.Wavelet
+    min_phase: MinPhaseMode = MinPhaseMode.NONE
+    axis: str = "time"
+    scales: list | tuple | npt.NDArray | None = None
+
+
+@processor_state
+class CWTState:
+    neg_rt_scales: npt.NDArray | None = None
+    int_psi_scales: list[npt.NDArray] | None = None
+    template: AxisArray | None = None
+    fbgen: typing.Generator[AxisArray, AxisArray, None] | None = None
+    last_conv_samp: npt.NDArray | None = None
+
+
+class CWTTransformer(BaseStatefulTransformer[CWTSettings, AxisArray, AxisArray, CWTState]):
+    def _hash_message(self, message: AxisArray) -> int:
+        ax_idx = message.get_axis_idx(self.settings.axis)
+        in_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
+        return hash(
+            (
+                message.data.dtype.kind,
+                message.axes[self.settings.axis].gain,
+                in_shape,
+                message.key,
+            )
+        )
+
+    def _reset_state(self, message: AxisArray) -> None:
+        precision = 10
+
+        # Process wavelet
+        wavelet = (
+            self.settings.wavelet
+            if isinstance(self.settings.wavelet, (pywt.ContinuousWavelet, pywt.Wavelet))
+            else pywt.DiscreteContinuousWavelet(self.settings.wavelet)
+        )
+        # Process wavelet integration
+        int_psi, wave_xvec = pywt.integrate_wavelet(wavelet, precision=precision)
+        int_psi = np.conj(int_psi) if wavelet.complex_cwt else int_psi
+
+        # Calculate scales and frequencies
+        if self.settings.frequencies is not None:
+            frequencies = np.sort(np.array(self.settings.frequencies))
+            scales = pywt.frequency2scale(
+                wavelet,
+                frequencies * message.axes[self.settings.axis].gain,
+                precision=precision,
+            )
+        else:
+            scales = np.sort(self.settings.scales)[::-1]
+
+        self._state.neg_rt_scales = -np.sqrt(scales)[:, None]
+
+        # Convert to appropriate dtype
+        dt_data = message.data.dtype
+        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)
+        # Note: Currently int_psi cannot be made non-complex once it is complex.
+
+        # Calculate waves for each scale
+        wave_xvec = np.asarray(wave_xvec, dtype=message.data.real.dtype)
+        wave_range = wave_xvec[-1] - wave_xvec[0]
+        step = wave_xvec[1] - wave_xvec[0]
+        self._state.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)
+            self._state.int_psi_scales.append(int_psi[reix][::-1])
+
+        # Setup filterbank generator
+        self._state.fbgen = filterbank(
+            self._state.int_psi_scales,
+            mode=FilterbankMode.CONV,
+            min_phase=self.settings.min_phase,
+            axis=self.settings.axis,
+        )
+
+        # Create output template
+        ax_idx = message.get_axis_idx(self.settings.axis)
+        in_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
+        freqs = pywt.scale2frequency(wavelet, scales, precision) / message.axes[self.settings.axis].gain
+        dummy_shape = in_shape + (len(scales), 0)
+        self._state.template = AxisArray(
+            np.zeros(dummy_shape, dtype=dt_cplx if wavelet.complex_cwt else dt_data),
+            dims=message.dims[:ax_idx] + message.dims[ax_idx + 1 :] + ["freq", self.settings.axis],
+            axes={
+                **message.axes,
+                "freq": AxisArray.CoordinateAxis(unit="Hz", data=freqs, dims=["freq"]),
+            },
+            key=message.key,
+        )
+        self._state.last_conv_samp = np.zeros(dummy_shape[:-1] + (1,), dtype=self._state.template.data.dtype)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        conv_msg = self._state.fbgen.send(message)
+
+        # Prepend with last_conv_samp before doing diff
+        dat = np.concatenate((self._state.last_conv_samp, conv_msg.data), axis=-1)
+        coef = self._state.neg_rt_scales * np.diff(dat, axis=-1)
+        # Store last_conv_samp for next iteration
+        self._state.last_conv_samp = conv_msg.data[..., -1:]
+
+        if self._state.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)]
+        return replace(
+            self._state.template,
+            data=coef,
+            axes={
+                **self._state.template.axes,
+                self.settings.axis: message.axes[self.settings.axis],
+            },
+        )
+
+
+class CWT(BaseTransformerUnit[CWTSettings, AxisArray, AxisArray, CWTTransformer]):
+    SETTINGS = CWTSettings
+
+
+def cwt(
+    frequencies: list | tuple | npt.NDArray | None,
+    wavelet: str | pywt.ContinuousWavelet | pywt.Wavelet,
+    min_phase: MinPhaseMode = MinPhaseMode.NONE,
+    axis: str = "time",
+    scales: list | tuple | npt.NDArray | None = None,
+) -> CWTTransformer:
+    """
+    Perform a continuous wavelet transform.
+    The function is equivalent to the :obj:`pywt.cwt` function, but is designed to work with streaming data.
+
+    Args:
+        frequencies: The wavelet frequencies to use in Hz. If `None` provided then the scales will be used.
+          Note: frequencies will be sorted from smallest to largest.
+        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.
+          This axis must be in the msg.axes and it must be of type AxisArray.LinearAxis.
+        scales: The scales to use. If None, the scales will be calculated from the frequencies.
+          Note: Scales will be sorted from largest to smallest.
+          Note: Use of scales is deprecated in favor of frequencies. Convert scales to frequencies using
+            `pywt.scale2frequency(wavelet, scales, precision=10) * fs` where fs is the sampling frequency.
+
+    Returns:
+        A primed Generator object that expects an :obj:`AxisArray` via `.send(axis_array)` of continuous data
+        and yields an :obj:`AxisArray` with a continuous wavelet transform in its data.
+    """
+    return CWTTransformer(
+        CWTSettings(
+            frequencies=frequencies,
+            wavelet=wavelet,
+            min_phase=min_phase,
+            axis=axis,
+            scales=scales,
+        )
+    )

ezmsg/sigproc/window.py

@@ -1,144 +1,328 @@
-from dataclasses import replace
+import enum
+import traceback
+import typing
 
 import ezmsg.core as ez
-import numpy as np
 import numpy.typing as npt
+import sparse
+from array_api_compat import get_namespace, is_pydata_sparse_namespace
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    replace,
+    slice_along_axis,
+    sliding_win_oneaxis,
+)
 
-from ezmsg.util.messages.axisarray import AxisArray
+from .util.profile import profile_subpub
+from .util.sparse import sliding_win_oneaxis as sparse_sliding_win_oneaxis
 
-from typing import AsyncGenerator, Optional, Tuple, List
+
+class Anchor(enum.Enum):
+    BEGINNING = "beginning"
+    END = "end"
+    MIDDLE = "middle"
 
 
 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"
+    axis: str | None = None
+    newaxis: str | None = None  # new axis for output. No new axes if None
+    window_dur: float | None = None  # Sec. passthrough if None
+    window_shift: float | None = None  # Sec. Use "1:1 mode" if None
+    zero_pad_until: str = "full"  # "full", "shift", "input", "none"
+    anchor: str | Anchor = Anchor.BEGINNING
 
 
-class WindowState(ez.State):
-    cur_settings: WindowSettings
+@processor_state
+class WindowState:
+    buffer: npt.NDArray | sparse.SparseArray | None = None
 
-    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
+    window_samples: int | None = None
 
+    window_shift_samples: int | None = None
 
-class Window(ez.Unit):
-    STATE: WindowState
-    SETTINGS: WindowSettings
+    shift_deficit: int = 0
+    """ Number of incoming samples to ignore. Only relevant when shift > window."""
 
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
-    INPUT_SETTINGS = ez.InputStream(WindowSettings)
+    newaxis_warned: bool = False
 
-    def initialize(self) -> None:
-        self.STATE.cur_settings = self.SETTINGS
+    out_newaxis: AxisArray.LinearAxis | None = None
 
-    @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
+    out_dims: list[str] | None = None
 
-    @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
+
+class WindowTransformer(BaseStatefulTransformer[WindowSettings, AxisArray, AxisArray, WindowState]):
+    """
+    Apply a sliding window along the specified axis to input streaming data.
+    The `windowing` method is perhaps the most useful and versatile method in ezmsg.sigproc, but its parameterization
+    can be difficult. Please read the argument descriptions carefully.
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        """
+
+        Args:
+            axis: The axis along which to segment windows.
+                If None, defaults to the first dimension of the first seen AxisArray.
+                Note: The windowed axis must be an AxisArray.LinearAxis, not an AxisArray.CoordinateAxis.
+            newaxis: New axis on which windows are delimited, immediately
+                preceding the target windowed axis. The data length along newaxis may be 0 if
+                this most recent push did not provide enough data for a new window.
+                If window_shift is None then the newaxis length will always be 1.
+            window_dur: The duration of the window in seconds.
+                If None, the function acts as a passthrough and all other parameters are ignored.
+            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.
+            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.
+            anchor: Determines the entry in `axis` that gets assigned `0`, which references the
+                value in `newaxis`. Can be of class :obj:`Anchor` or a string representation of an :obj:`Anchor`.
+        """
+        super().__init__(*args, **kwargs)
+
+        # Sanity-check settings
+        # if self.settings.newaxis is None:
+        #     ez.logger.warning("`newaxis=None` will be replaced with `newaxis='win'`.")
+        #     object.__setattr__(self.settings, "newaxis", "win")
+        if self.settings.window_shift is None and self.settings.zero_pad_until != "input":
+            ez.logger.warning(
+                "`zero_pad_until` must be 'input' if `window_shift` is None. "
+                f"Ignoring received argument value: {self.settings.zero_pad_until}"
+            )
+            object.__setattr__(self.settings, "zero_pad_until", "input")
+        elif self.settings.window_shift is not None and self.settings.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 `zero_pad_until='shift'` when `window_shift` is float-valued."
+            )
+        try:
+            object.__setattr__(self.settings, "anchor", Anchor(self.settings.anchor))
+        except ValueError:
+            raise ValueError(
+                f"Invalid anchor: {self.settings.anchor}. Valid anchor are: {', '.join([e.value for e in Anchor])}"
             )
 
-        # 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
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        axis_info = message.get_axis(axis)
+        fs = 1.0 / axis_info.gain
+        samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+
+        return hash(samp_shape + (fs, message.key))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        _newaxis = self.settings.newaxis or "win"
+        if not self._state.newaxis_warned and _newaxis in message.dims:
+            ez.logger.warning(f"newaxis {_newaxis} present in input dims. Using {_newaxis}_win instead")
+            self._state.newaxis_warned = True
+            self.settings.newaxis = f"{_newaxis}_win"
+
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        axis_info = message.get_axis(axis)
+        fs = 1.0 / axis_info.gain
+
+        xp = get_namespace(message.data)
+
+        self._state.window_samples = int(self.settings.window_dur * fs)
+        if self.settings.window_shift is not None:
+            # If window_shift is None, we are in "1:1 mode" and window_shift_samples is not used.
+            self._state.window_shift_samples = int(self.settings.window_shift * fs)
+        if self.settings.zero_pad_until == "none":
+            req_samples = self._state.window_samples
+        elif self.settings.zero_pad_until == "shift" and self.settings.window_shift is not None:
+            req_samples = self._state.window_shift_samples
+        else:  # i.e. zero_pad_until == "input"
+            req_samples = message.data.shape[axis_idx]
+        n_zero = max(0, self._state.window_samples - req_samples)
+        init_buffer_shape = message.data.shape[:axis_idx] + (n_zero,) + message.data.shape[axis_idx + 1 :]
+        self._state.buffer = xp.zeros(init_buffer_shape, dtype=message.data.dtype)
+
+        # Prepare reusable parts of output
+        if self._state.out_newaxis is None:
+            self._state.out_dims = list(message.dims[:axis_idx]) + [_newaxis] + list(message.dims[axis_idx:])
+            self._state.out_newaxis = replace(
+                axis_info,
+                gain=0.0 if self.settings.window_shift is None else axis_info.gain * self._state.window_shift_samples,
+                offset=0.0,  # offset modified per-msg below
+            )
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        if self.settings.window_dur is None:
+            # Shortcut for no windowing
+            return message
+        return super().__call__(message)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        axis_info = message.get_axis(axis)
+
+        xp = get_namespace(message.data)
+
+        # Add new data to buffer.
+        # Currently, we concatenate the new time samples and clip the output.
+        # np.roll is not preferred as it 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)
+        # A circular buffer might be faster.
+        self._state.buffer = xp.concatenate((self._state.buffer, message.data), axis=axis_idx)
+
+        # Create a vector of buffer timestamps to track axis `offset` in output(s)
+        buffer_t0 = 0.0
+        buffer_tlen = self._state.buffer.shape[axis_idx]
+
+        # Adjust so first _new_ sample at index 0.
+        buffer_t0 -= self._state.buffer.shape[axis_idx] - message.data.shape[axis_idx]
+
+        # Convert form indices to 'units' (probably seconds).
+        buffer_t0 *= axis_info.gain
+        buffer_t0 += axis_info.offset
+
+        if self.settings.window_shift is not None and self._state.shift_deficit > 0:
+            n_skip = min(self._state.buffer.shape[axis_idx], self._state.shift_deficit)
+            if n_skip > 0:
+                self._state.buffer = slice_along_axis(self._state.buffer, slice(n_skip, None), axis_idx)
+                buffer_t0 += n_skip * axis_info.gain
+                buffer_tlen -= n_skip
+                self._state.shift_deficit -= n_skip
 
-        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
+        # Generate outputs.
+        # Preliminary copy of axes without the axes that we are modifying.
+        _newaxis = self.settings.newaxis or "win"
+        out_axes = {k: v for k, v in message.axes.items() if k not in [_newaxis, axis]}
 
-        outputs: List[Tuple[npt.NDArray, float]] = []
+        # Update targeted (windowed) axis so that its offset is relative to the new axis
+        if self.settings.anchor == Anchor.BEGINNING:
+            out_axes[axis] = replace(axis_info, offset=0.0)
+        elif self.settings.anchor == Anchor.END:
+            out_axes[axis] = replace(axis_info, offset=-self.settings.window_dur)
+        elif self.settings.anchor == Anchor.MIDDLE:
+            out_axes[axis] = replace(axis_info, offset=-self.settings.window_dur / 2)
 
-        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]))
+        # How we update .data and .axes[newaxis] depends on the windowing mode.
+        if self.settings.window_shift is None:
+            # one-to-one mode -- Each send yields exactly one window containing only the most recent samples.
+            self._state.buffer = slice_along_axis(
+                self._state.buffer, slice(-self._state.window_samples, None), axis_idx
+            )
+            out_dat = self._state.buffer.reshape(
+                self._state.buffer.shape[:axis_idx] + (1,) + self._state.buffer.shape[axis_idx:]
+            )
+            win_offset = buffer_t0 + axis_info.gain * (buffer_tlen - self._state.window_samples)
+        elif self._state.buffer.shape[axis_idx] >= self._state.window_samples:
+            # Deterministic window shifts.
+            sliding_win_fun = sparse_sliding_win_oneaxis if is_pydata_sparse_namespace(xp) else sliding_win_oneaxis
+            out_dat = sliding_win_fun(
+                self._state.buffer,
+                self._state.window_samples,
+                axis_idx,
+                step=self._state.window_shift_samples,
+            )
+            win_offset = buffer_t0
 
+            # Drop expired beginning of buffer and update shift_deficit
+            multi_shift = self._state.window_shift_samples * out_dat.shape[axis_idx]
+            self._state.shift_deficit = max(0, multi_shift - self._state.buffer.shape[axis_idx])
+            self._state.buffer = slice_along_axis(self._state.buffer, slice(multi_shift, None), axis_idx)
         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)
+            # Not enough data to make a new window. Return empty data.
+            empty_data_shape = (
+                message.data.shape[:axis_idx] + (0, self._state.window_samples) + message.data.shape[axis_idx + 1 :]
+            )
+            out_dat = xp.zeros(empty_data_shape, dtype=message.data.dtype)
+            # out_newaxis will have first timestamp in input... but mostly meaningless because output is size-zero.
+            win_offset = axis_info.offset
+
+        if self.settings.anchor == Anchor.END:
+            win_offset += self.settings.window_dur
+        elif self.settings.anchor == Anchor.MIDDLE:
+            win_offset += self.settings.window_dur / 2
+        self._state.out_newaxis = replace(self._state.out_newaxis, offset=win_offset)
+
+        msg_out = replace(
+            message,
+            data=out_dat,
+            dims=self._state.out_dims,
+            axes={**out_axes, _newaxis: self._state.out_newaxis},
+        )
+        return msg_out
+
+
+class Window(BaseTransformerUnit[WindowSettings, AxisArray, AxisArray, WindowTransformer]):
+    SETTINGS = WindowSettings
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(OUTPUT_SIGNAL)
+    @profile_subpub(trace_oldest=False)
+    async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
+        """
+        override superclass on_signal so we can opt to yield once or multiple times after dropping the win axis.
+        """
+        # TODO: The transfomer overwrites settings.newaxis from None to "win",
+        #  then we no longer know if the user wants to trim out the newaxis from the unit.
+        xp = get_namespace(message.data)
+        try:
+            ret = self.processor(message)
+            if ret.data.size > 0:
+                if self.SETTINGS.newaxis is not None or self.SETTINGS.window_dur is None:
+                    # Multi-win mode or pass-through mode.
+                    yield self.OUTPUT_SIGNAL, ret
                 else:
-                    yield self.OUTPUT_SIGNAL, replace(msg, data=out_view)
+                    # We need to split out_msg into multiple yields, dropping newaxis.
+                    axis_idx = ret.get_axis_idx("win")
+                    win_axis = ret.axes["win"]
+                    offsets = win_axis.value(xp.asarray(range(ret.data.shape[axis_idx])))
+                    for msg_ix in range(ret.data.shape[axis_idx]):
+                        # Need to drop 'win' and replace self.SETTINGS.axis from axes.
+                        _out_axes = {
+                            **{k: v for k, v in ret.axes.items() if k not in ["win", self.SETTINGS.axis]},
+                            self.SETTINGS.axis: replace(ret.axes[self.SETTINGS.axis], offset=offsets[msg_ix]),
+                        }
+                        _ret = replace(
+                            ret,
+                            data=slice_along_axis(ret.data, msg_ix, axis_idx),
+                            dims=ret.dims[:axis_idx] + ret.dims[axis_idx + 1 :],
+                            axes=_out_axes,
+                        )
+                        yield self.OUTPUT_SIGNAL, _ret
+
+        except Exception:
+            ez.logger.info(traceback.format_exc())
+
+
+def windowing(
+    axis: str | None = None,
+    newaxis: str | None = None,
+    window_dur: float | None = None,
+    window_shift: float | None = None,
+    zero_pad_until: str = "full",
+    anchor: str | Anchor = Anchor.BEGINNING,
+) -> WindowTransformer:
+    return WindowTransformer(
+        WindowSettings(
+            axis=axis,
+            newaxis=newaxis,
+            window_dur=window_dur,
+            window_shift=window_shift,
+            zero_pad_until=zero_pad_until,
+            anchor=anchor,
+        )
+    )