ezmsg-sigproc 1.7.0__tar.gz → 1.7.1__tar.gz

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/PKG-INFO

@@ -1,14 +1,16 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 1.7.0
+Version: 1.7.1
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>
-License: MIT
+License-Expression: MIT
+License-File: LICENSE.txt
 Requires-Python: >=3.10.15
 Requires-Dist: ezmsg>=3.6.0
 Requires-Dist: numpy>=1.26.0
 Requires-Dist: pywavelets>=1.6.0
 Requires-Dist: scipy>=1.13.1
+Requires-Dist: sparse>=0.15.4
 Provides-Extra: test
 Requires-Dist: flake8>=7.1.1; extra == 'test'
 Requires-Dist: frozendict>=2.4.4; extra == 'test'

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "numpy>=1.26.0",
     "pywavelets>=1.6.0",
     "scipy>=1.13.1",
+    "sparse>=0.15.4",
 ]
 
 [project.optional-dependencies]

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/__version__.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.7.0'
-__version_tuple__ = version_tuple = (1, 7, 0)
+__version__ = version = '1.7.1'
+__version_tuple__ = version_tuple = (1, 7, 1)

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/bandpower.py

@@ -35,6 +35,7 @@ def bandpower(
     f_spec = spectrogram(
         window_dur=spectrogram_settings.window_dur,
         window_shift=spectrogram_settings.window_shift,
+        window_anchor=spectrogram_settings.window_anchor,
         window=spectrogram_settings.window,
         transform=spectrogram_settings.transform,
         output=spectrogram_settings.output,

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/base.py

@@ -6,6 +6,8 @@ import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.generator import GenState
 
+from .util.profile import profile_subpub
+
 
 class GenAxisArray(ez.Unit):
     STATE = GenState
@@ -28,6 +30,7 @@ class GenAxisArray(ez.Unit):
 
     @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:
         try:
             ret = self.STATE.gen.send(message)

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/filterbank.py

@@ -214,7 +214,8 @@ def filterbank(
                 pad = np.zeros(dest_arr.shape[:-1] + (n_dest - dest_arr.shape[-1],))
                 dest_arr = np.concatenate(dest_arr, pad, axis=-1)
             dest_arr.fill(0)
-            # TODO: Parallelize this loop.
+            # Note: I tried several alternatives to this loop; all were slower than this.
+            #  numba.jit; stride_tricks + np.einsum; threading. Latter might be better with Python 3.13.
             for k_ix, k in enumerate(kernels):
                 n_out = in_dat.shape[-1] + k.shape[-1] - 1
                 dest_arr[..., k_ix, :n_out] = np.apply_along_axis(

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/sampler.py

@@ -14,6 +14,8 @@ from ezmsg.util.messages.axisarray import (
 from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
+from .util.profile import profile_subpub
+
 
 @dataclass(unsafe_hash=True)
 class SampleTriggerMessage:
@@ -284,6 +286,7 @@ class Sampler(ez.Unit):
 
     @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
     @ez.publisher(OUTPUT_SAMPLE)
+    @profile_subpub(trace_oldest=False)
     async def on_signal(self, msg: AxisArray) -> typing.AsyncGenerator:
         pub_samples = self.STATE.gen.send(msg)
         for sample in pub_samples:

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/signalinjector.py

@@ -6,6 +6,8 @@ from ezmsg.util.messages.util import replace
 import numpy as np
 import numpy.typing as npt
 
+from .util.profile import profile_subpub
+
 
 class SignalInjectorSettings(ez.Settings):
     time_dim: str = "time"  # Input signal needs a time dimension with units in sec.
@@ -50,6 +52,7 @@ class SignalInjector(ez.Unit):
 
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
+    @profile_subpub(trace_oldest=False)
     async def inject(self, msg: AxisArray) -> typing.AsyncGenerator:
         if self.STATE.cur_shape != msg.shape:
             self.STATE.cur_shape = msg.shape

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/spectrogram.py

@@ -5,7 +5,7 @@ from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.generator import consumer, compose
 from ezmsg.util.messages.modify import modify_axis
 
-from .window import windowing
+from .window import windowing, Anchor
 from .spectrum import spectrum, WindowFunction, SpectralTransform, SpectralOutput
 from .base import GenAxisArray
 
@@ -14,6 +14,7 @@ from .base import GenAxisArray
 def spectrogram(
     window_dur: float | None = None,
     window_shift: float | None = None,
+    window_anchor: str | Anchor = Anchor.BEGINNING,
     window: WindowFunction = WindowFunction.HANNING,
     transform: SpectralTransform = SpectralTransform.REL_DB,
     output: SpectralOutput = SpectralOutput.POSITIVE,
@@ -28,6 +29,7 @@ def spectrogram(
     Args:
         window_dur: See :obj:`ezmsg.sigproc.window.windowing`
         window_shift: See :obj:`ezmsg.sigproc.window.windowing`
+        window_anchor: See :obj:`ezmsg.sigproc.window.windowing`
         window: See :obj:`ezmsg.sigproc.spectrum.spectrum`
         transform: See :obj:`ezmsg.sigproc.spectrum.spectrum`
         output: See :obj:`ezmsg.sigproc.spectrum.spectrum`
@@ -44,6 +46,7 @@ def spectrogram(
             window_dur=window_dur,
             window_shift=window_shift,
             zero_pad_until="shift" if window_shift is not None else "input",
+            anchor=window_anchor,
         ),
         spectrum(axis="time", window=window, transform=transform, output=output),
         modify_axis(name_map={"win": "time"}),
@@ -66,6 +69,7 @@ class SpectrogramSettings(ez.Settings):
     window_dur: float | None = None  # window duration in seconds
     window_shift: float | None = None
     """"window step in seconds. If None, window_shift == window_dur"""
+    window_anchor: str | Anchor = Anchor.BEGINNING
 
     # See SpectrumSettings for details of following settings:
     window: WindowFunction = WindowFunction.HAMMING
@@ -84,6 +88,7 @@ class Spectrogram(GenAxisArray):
         self.STATE.gen = spectrogram(
             window_dur=self.SETTINGS.window_dur,
             window_shift=self.SETTINGS.window_shift,
+            window_anchor=self.SETTINGS.window_anchor,
             window=self.SETTINGS.window,
             transform=self.SETTINGS.transform,
             output=self.SETTINGS.output,

ezmsg_sigproc-1.7.1/src/ezmsg/sigproc/util/profile.py

@@ -0,0 +1,131 @@
+import functools
+import logging
+import os
+from pathlib import Path
+import time
+import typing
+
+import ezmsg.core as ez
+
+
+def get_logger_path() -> Path:
+    # Retrieve the logfile name from the environment variable
+    logfile = os.environ.get("EZMSG_PROFILE", None)
+
+    # Determine the log file path, defaulting to "ezprofiler.log" if not set
+    logpath = Path(logfile or "ezprofiler.log")
+
+    # If the log path is not absolute, prepend it with the user's home directory and ".ezmsg/profile"
+    if not logpath.is_absolute():
+        logpath = Path.home() / ".ezmsg" / "profile" / logpath
+
+    return logpath
+
+
+def _setup_logger(append: bool = False) -> logging.Logger:
+    logpath = get_logger_path()
+    logpath.parent.mkdir(parents=True, exist_ok=True)
+
+    if not append:
+        # Remove the file if it exists
+        logpath.unlink(missing_ok=True)
+
+    # Create a logger with the name "ezprofile"
+    _logger = logging.getLogger("ezprofile")
+
+    # Set the logger's level to EZMSG_LOGLEVEL env var value if it exists, otherwise INFO
+    _logger.setLevel(os.environ.get("EZMSG_LOGLEVEL", "INFO").upper())
+
+    # Create a file handler to write log messages to the log file
+    fh = logging.FileHandler(logpath)
+    fh.setLevel(logging.DEBUG)  # Set the file handler log level to DEBUG
+
+    # Add the file handler to the logger
+    _logger.addHandler(fh)
+
+    # Add the first row without formatting.
+    _logger.debug(",".join(["Time", "Source", "Topic", "SampleTime", "PerfCounter", "Elapsed"]))
+
+    # Set the log message format
+    formatter = logging.Formatter(
+        "%(asctime)s,%(message)s",
+        datefmt="%Y-%m-%dT%H:%M:%S%z"
+    )
+    fh.setFormatter(formatter)
+
+    return _logger
+
+
+logger = _setup_logger(append=True)
+
+
+def _process_obj(obj, trace_oldest: bool = True):
+    samp_time = None
+    if hasattr(obj, "axes") and ("time" in obj.axes or "win" in obj.axes):
+        axis = "win" if "win" in obj.axes else "time"
+        ax = obj.get_axis(axis)
+        len = obj.data.shape[obj.get_axis_idx(axis)]
+        if len > 0:
+            idx = 0 if trace_oldest else (len - 1)
+            if hasattr(ax, "data"):
+                samp_time = ax.data[idx]
+            else:
+                samp_time = ax.value(idx)
+            if ax == "win" and "time" in obj.axes:
+                if hasattr(obj.axes["time"], "data"):
+                    samp_time += obj.axes["time"].data[idx]
+                else:
+                    samp_time += obj.axes["time"].value(idx)
+    return samp_time
+
+
+def profile_method(trace_oldest: bool = True):
+    """
+    Decorator to profile a method by logging its execution time and other details.
+
+    Args:
+        trace_oldest (bool): If True, trace the oldest sample time; otherwise, trace the newest.
+
+    Returns:
+        Callable: The decorated function with profiling.
+    """
+    def profiling_decorator(func: typing.Callable):
+        @functools.wraps(func)
+        def wrapped_func(caller, *args, **kwargs):
+            start = time.perf_counter()
+            res = func(caller, *args, **kwargs)
+            stop = time.perf_counter()
+            source = '.'.join((caller.__class__.__module__, caller.__class__.__name__))
+            topic = f"{caller.address}"
+            samp_time = _process_obj(res, trace_oldest=trace_oldest)
+            logger.debug(",".join([source, topic, f"{samp_time}", f"{stop}", f"{(stop - start) * 1e3:0.4f}"]))
+            return res
+        return wrapped_func if logger.level == logging.DEBUG else func
+    return profiling_decorator
+
+
+def profile_subpub(trace_oldest: bool = True):
+    """
+    Decorator to profile a subscriber-publisher method in an ezmsg Unit
+    by logging its execution time and other details.
+
+    Args:
+        trace_oldest (bool): If True, trace the oldest sample time; otherwise, trace the newest.
+
+    Returns:
+        Callable: The decorated async task with profiling.
+    """
+    def profiling_decorator(func: typing.Callable):
+        @functools.wraps(func)
+        async def wrapped_task(unit: ez.Unit, msg: typing.Any = None) -> None:
+            source = '.'.join((unit.__class__.__module__, unit.__class__.__name__))
+            topic = f"{unit.address}"
+            start = time.perf_counter()
+            async for stream, obj in func(unit, msg):
+                stop = time.perf_counter()
+                samp_time = _process_obj(obj, trace_oldest=trace_oldest)
+                logger.debug(",".join([source, topic, f"{samp_time}", f"{stop}", f"{(stop - start) * 1e3:0.4f}"]))
+                start = stop
+                yield stream, obj
+        return wrapped_task if logger.level == logging.DEBUG else func
+    return profiling_decorator

ezmsg_sigproc-1.7.1/src/ezmsg/sigproc/util/sparse.py

@@ -0,0 +1,29 @@
+import sparse
+
+
+def sliding_win_oneaxis(
+    s: sparse.SparseArray, nwin: int, axis: int, step: int = 1
+) -> sparse.SparseArray:
+    """
+    Like `ezmsg.util.messages.axisarray.sliding_win_oneaxis` but for sparse arrays.
+
+    Args:
+        s: The input sparse array.
+        nwin: The size of the sliding window.
+        axis: The axis along which the sliding window will be applied.
+        step: The size of the step between windows. If > 1, the strided window will be sliced with `slice_along_axis`.
+
+    Returns:
+
+    """
+    if -s.ndim <= axis < 0:
+        axis = s.ndim + axis
+    targ_slices = [slice(_, _ + nwin) for _ in range(0, s.shape[axis] - nwin + 1, step)]
+    s = s.reshape(s.shape[:axis] + (1,) + s.shape[axis:])
+    full_slices = (slice(None),) * s.ndim
+    full_slices = [
+        full_slices[: axis + 1] + (sl,) + full_slices[axis + 2 :] for sl in targ_slices
+    ]
+    result = sparse.concatenate([s[_] for _ in full_slices], axis=axis)
+    # TODO: Profile this approach vs modifying coords only.
+    return result

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/wavelets.py

@@ -171,11 +171,11 @@ class CWTSettings(ez.Settings):
     Settings for :obj:`CWT`
     See :obj:`cwt` for argument details.
     """
-
-    scales: list | tuple | npt.NDArray
+    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
 
 
 class CWT(GenAxisArray):
@@ -187,8 +187,9 @@ class CWT(GenAxisArray):
 
     def construct_generator(self):
         self.STATE.gen = cwt(
-            scales=self.SETTINGS.scales,
+            frequencies=self.SETTINGS.frequencies,
             wavelet=self.SETTINGS.wavelet,
             min_phase=self.SETTINGS.min_phase,
             axis=self.SETTINGS.axis,
+            scales=self.SETTINGS.scales,
         )

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/src/ezmsg/sigproc/window.py

@@ -1,9 +1,11 @@
+import enum
 import traceback
 import typing
 
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
+import sparse
 from ezmsg.util.messages.axisarray import (
     AxisArray,
     slice_along_axis,
@@ -13,6 +15,13 @@ from ezmsg.util.messages.axisarray import (
 from ezmsg.util.generator import consumer
 
 from .base import GenAxisArray
+from .util.sparse import sliding_win_oneaxis as sparse_sliding_win_oneaxis
+
+
+class Anchor(enum.Enum):
+    BEGINNING = "beginning"
+    END = "end"
+    MIDDLE = "middle"
 
 
 @consumer
@@ -22,6 +31,7 @@ def windowing(
     window_dur: float | None = None,
     window_shift: float | None = None,
     zero_pad_until: str = "input",
+    anchor: str | Anchor = Anchor.BEGINNING,
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Apply a sliding window along the specified axis to input streaming data.
@@ -49,6 +59,8 @@ def windowing(
             - "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`.
 
     Returns:
         A primed generator that accepts an :obj:`AxisArray` via `.send(axis_array)`
@@ -69,10 +81,15 @@ def windowing(
             "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."
         )
+    try:
+        anchor = Anchor(anchor)
+    except ValueError:
+        raise ValueError(f"Invalid anchor: {anchor}. Valid anchor are: {', '.join([e.value for e in Anchor])}")
+
     msg_out = AxisArray(np.array([]), dims=[""])
 
     # State variables
-    buffer: npt.NDArray | None = None
+    buffer: npt.NDArray | sparse.SparseArray | None = None
     window_samples: int | None = None
     window_shift_samples: int | None = None
     # Number of incoming samples to ignore. Only relevant when shift > window.:
@@ -83,6 +100,8 @@ def windowing(
     out_dims: list[str] | None = None
 
     check_inputs = {"samp_shape": None, "fs": None, "key": None}
+    concat_fun = np.concatenate
+    sliding_win_fun = sliding_win_oneaxis
 
     while True:
         msg_in: AxisArray = yield msg_out
@@ -116,6 +135,15 @@ def windowing(
             check_inputs["fs"] = fs
             check_inputs["key"] = msg_in.key
 
+            if isinstance(msg_in.data, sparse.SparseArray):
+                concat_fun = sparse.concatenate
+                sliding_win_fun = sparse_sliding_win_oneaxis
+                zero_fun = sparse.zeros
+            else:
+                concat_fun = np.concatenate
+                sliding_win_fun = sliding_win_oneaxis  # Requires updated signature in ezmsg dev.
+                zero_fun = np.zeros
+
             window_samples = int(window_dur * fs)
             if not b_1to1:
                 window_shift_samples = int(window_shift * fs)
@@ -131,7 +159,7 @@ def windowing(
                 + (n_zero,)
                 + msg_in.data.shape[axis_idx + 1 :]
             )
-            buffer = np.zeros(init_buffer_shape, dtype=msg_in.data.dtype)
+            buffer = zero_fun(init_buffer_shape, dtype=msg_in.data.dtype)
 
         # Add new data to buffer.
         # Currently, we concatenate the new time samples and clip the output.
@@ -140,21 +168,21 @@ def windowing(
         # is generally faster than np.roll and slicing anyway, but this could still
         # be a performance bottleneck for large memory arrays.
         # A circular buffer might be faster.
-        buffer = np.concatenate((buffer, msg_in.data), axis=axis_idx)
+        buffer = concat_fun((buffer, msg_in.data), axis=axis_idx)
 
         # 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[-msg_in.data.shape[axis_idx]]
+        buffer_tvec = np.arange(buffer.shape[axis_idx]).astype(float)
+        # Adjust so first _new_ sample at index 0.
+        buffer_tvec -= buffer_tvec[-msg_in.data.shape[axis_idx]]
         # Convert form indices to 'units' (probably seconds).
-        buffer_offset *= axis_info.gain
-        buffer_offset += axis_info.offset
+        buffer_tvec *= axis_info.gain
+        buffer_tvec += 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, slice(n_skip, None), axis_idx)
-                buffer_offset = buffer_offset[n_skip:]
+                buffer_tvec = buffer_tvec[n_skip:]
                 shift_deficit -= n_skip
 
         # Prepare reusable parts of output
@@ -171,8 +199,12 @@ def windowing(
         out_axes = {k: v for k, v in msg_in.axes.items() if k not in [newaxis, axis]}
 
         # Update targeted (windowed) axis so that its offset is relative to the new axis
-        # TODO: If we have `anchor_newest=True` then offset should be -win_dur
-        out_axes[axis] = replace(axis_info, offset=0.0)
+        if anchor == Anchor.BEGINNING:
+            out_axes[axis] = replace(axis_info, offset=0.0)
+        elif anchor == Anchor.END:
+            out_axes[axis] = replace(axis_info, offset=-window_dur)
+        elif anchor == Anchor.MIDDLE:
+            out_axes[axis] = replace(axis_info, offset=-window_dur / 2)
 
         # How we update .data and .axes[newaxis] depends on the windowing mode.
         if b_1to1:
@@ -181,19 +213,14 @@ def windowing(
             out_dat = buffer.reshape(
                 buffer.shape[:axis_idx] + (1,) + buffer.shape[axis_idx:]
             )
-            out_newaxis = replace(out_newaxis, offset=buffer_offset[-window_samples])
+            win_offset = buffer_tvec[-window_samples]
         elif buffer.shape[axis_idx] >= window_samples:
             # Deterministic window shifts.
-            # Note: After https://github.com/ezmsg-org/ezmsg/pull/152, add `window_shift_samples` as the last arg
-            #  to `sliding_win_oneaxis` and remove the call to `slice_along_axis`.
-            out_dat = sliding_win_oneaxis(buffer, window_samples, axis_idx)
-            out_dat = slice_along_axis(
-                out_dat, slice(None, None, window_shift_samples), axis_idx
-            )
-            offset_view = sliding_win_oneaxis(buffer_offset, window_samples, 0)[
+            out_dat = sliding_win_fun(buffer, window_samples, axis_idx, step=window_shift_samples)
+            offset_view = sliding_win_oneaxis(buffer_tvec, window_samples, 0)[
                 ::window_shift_samples
             ]
-            out_newaxis = replace(out_newaxis, offset=offset_view[0, 0])
+            win_offset = offset_view[0, 0]
 
             # Drop expired beginning of buffer and update shift_deficit
             multi_shift = window_shift_samples * out_dat.shape[axis_idx]
@@ -208,7 +235,13 @@ def windowing(
             )
             out_dat = np.zeros(empty_data_shape, dtype=msg_in.data.dtype)
             # out_newaxis will have first timestamp in input... but mostly meaningless because output is size-zero.
-            out_newaxis = replace(out_newaxis, offset=axis_info.offset)
+            win_offset = axis_info.offset
+
+        if anchor == Anchor.END:
+            win_offset += window_dur
+        elif anchor == Anchor.MIDDLE:
+            win_offset += window_dur / 2
+        out_newaxis = replace(out_newaxis, offset=win_offset)
 
         msg_out = replace(
             msg_in, data=out_dat, dims=out_dims, axes={**out_axes, newaxis: out_newaxis}
@@ -221,6 +254,7 @@ class WindowSettings(ez.Settings):
     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):
@@ -243,6 +277,7 @@ class Window(GenAxisArray):
             window_dur=self.SETTINGS.window_dur,
             window_shift=self.SETTINGS.window_shift,
             zero_pad_until=self.SETTINGS.zero_pad_until,
+            anchor=self.SETTINGS.anchor,
         )
 
     @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
@@ -261,10 +296,7 @@ class Window(GenAxisArray):
                     # We need to split out_msg into multiple yields, dropping newaxis.
                     axis_idx = out_msg.get_axis_idx("win")
                     win_axis = out_msg.axes["win"]
-                    offsets = (
-                        np.arange(out_msg.data.shape[axis_idx]) * win_axis.gain
-                        + win_axis.offset
-                    )
+                    offsets = win_axis.value(np.arange(out_msg.data.shape[axis_idx]))
                     for msg_ix in range(out_msg.data.shape[axis_idx]):
                         # Need to drop 'win' and replace self.SETTINGS.axis from axes.
                         _out_axes = {

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/tests/test_window.py

@@ -7,6 +7,7 @@ import numpy as np
 import numpy.typing as npt
 from numpy.lib.stride_tricks import sliding_window_view
 from frozendict import frozendict
+import sparse
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messagegate import MessageGate, MessageGateSettings
@@ -27,6 +28,7 @@ def calculate_expected_results(
     fs,
     win_shift,
     zero_pad,
+    anchor,
     msg_block_size,
     shift_len,
     win_len,
@@ -70,6 +72,12 @@ def calculate_expected_results(
     else:
         expected = expected[:, ::shift_len]
         tvec = tvec[::shift_len, 0]
+
+    if anchor == "middle":
+        tvec = tvec + win_len / (2 * fs)
+    elif anchor == "end":
+        tvec = tvec + win_len / fs
+
     # Transpose to put time_ax and win_ax in the correct locations.
     if win_ax == 0:
         expected = np.moveaxis(expected, 0, -1)
@@ -111,6 +119,7 @@ def test_window_gen_nodur():
 @pytest.mark.parametrize("win_shift", [None, 0.2, 1.0])
 @pytest.mark.parametrize("zero_pad", ["input", "shift", "none"])
 @pytest.mark.parametrize("fs", [10.0, 500.0])
+@pytest.mark.parametrize("anchor", ["beginning", "middle", "end"])
 @pytest.mark.parametrize("time_ax", [0, 1])
 def test_window_generator(
     msg_block_size: int,
@@ -119,6 +128,7 @@ def test_window_generator(
     win_shift: float | None,
     zero_pad: str,
     fs: float,
+    anchor: str,
     time_ax: int,
 ):
     nchans = 3
@@ -141,6 +151,7 @@ def test_window_generator(
         window_dur=win_dur,
         window_shift=win_shift,
         zero_pad_until=zero_pad,
+        anchor=anchor
     )
 
     # Create inputs and send them to the generator, collecting the results along the way.
@@ -211,6 +222,7 @@ def test_window_generator(
         fs,
         win_shift,
         zero_pad,
+        anchor,
         msg_block_size,
         shift_len,
         win_len,
@@ -225,6 +237,62 @@ def test_window_generator(
     assert np.allclose(offsets, tvec)
 
 
+@pytest.mark.parametrize("win_dur", [0.3, 1.0])
+@pytest.mark.parametrize("win_shift", [0.2, 1.0, None])
+@pytest.mark.parametrize("zero_pad", ["input", "shift", "none"])
+def test_sparse_window(
+    win_dur: float,
+    win_shift: float | None,
+    zero_pad: str,
+):
+    fs = 100.0
+    n_ch = 5
+    n_samps = 1_000
+    msg_len = 100
+    win_len = int(win_dur * fs)
+    rng = np.random.default_rng()
+    s = sparse.random((n_samps, n_ch), density=0.1, random_state=rng) > 0
+    in_msgs = [
+        AxisArray(
+            data=s[msg_ix * msg_len : (msg_ix + 1) * msg_len],
+            dims=["time", "ch"],
+            axes={
+                "time": AxisArray.Axis.TimeAxis(fs=fs, offset=msg_ix / fs),
+            },
+            key="test_sparse_window",
+        )
+        for msg_ix in range(10)
+    ]
+
+    proc = windowing(
+        axis="time",
+        newaxis="win",
+        window_dur=win_dur,
+        window_shift=win_shift,
+        zero_pad_until=zero_pad,
+    )
+    out_msgs = [proc.send(_) for _ in in_msgs]
+    nwins = 0
+    for om in out_msgs:
+        assert om.dims == ["win", "time", "ch"]
+        assert om.data.shape[1] == win_len
+        assert om.data.shape[2] == n_ch
+        nwins += om.data.shape[0]
+    if win_shift is None:
+        # 1:1 mode
+        assert nwins == len(out_msgs)
+    else:
+        shift_len = int(win_shift * fs)
+        prepended = 0
+        if zero_pad == "input":
+            prepended = max(0, win_len - msg_len)
+        elif zero_pad == "shift":
+            prepended = max(0, win_len - shift_len)
+        win_offsets = np.arange(n_samps + prepended)[::shift_len]
+        expected_nwins = np.sum(win_offsets <= (n_samps + prepended - win_len))
+        assert nwins == expected_nwins
+
+
 class WindowSystemSettings(ez.Settings):
     num_msgs: int
     counter_settings: CounterSettings
@@ -374,6 +442,7 @@ def test_window_system(
         fs,
         win_shift,
         zero_pad,
+        "beginning",
         msg_block_size,
         shift_len,
         win_len,

{ezmsg_sigproc-1.7.0 → ezmsg_sigproc-1.7.1}/uv.lock

@@ -115,13 +115,14 @@ wheels = [
 
 [[package]]
 name = "ezmsg-sigproc"
-version = "1.6.0"
+version = "1.6.1.dev3+ga96f0f9.d20241223"
 source = { editable = "." }
 dependencies = [
     { name = "ezmsg" },
     { name = "numpy" },
     { name = "pywavelets" },
     { name = "scipy" },
+    { name = "sparse" },
 ]
 
 [package.optional-dependencies]
@@ -150,6 +151,7 @@ requires-dist = [
     { name = "pytest-cov", marker = "extra == 'test'", specifier = ">=5.0.0" },
     { name = "pywavelets", specifier = ">=1.6.0" },
     { name = "scipy", specifier = ">=1.13.1" },
+    { name = "sparse", specifier = ">=0.15.4" },
 ]
 
 [package.metadata.requires-dev]
@@ -218,6 +220,29 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ef/a6/62565a6e1cf69e10f5727360368e451d4b7f58beeac6173dc9db836a5b46/iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374", size = 5892 },
 ]
 
+[[package]]
+name = "llvmlite"
+version = "0.43.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/9f/3d/f513755f285db51ab363a53e898b85562e950f79a2e6767a364530c2f645/llvmlite-0.43.0.tar.gz", hash = "sha256:ae2b5b5c3ef67354824fb75517c8db5fbe93bc02cd9671f3c62271626bc041d5", size = 157069 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/23/ff/6ca7e98998b573b4bd6566f15c35e5c8bea829663a6df0c7aa55ab559da9/llvmlite-0.43.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:a289af9a1687c6cf463478f0fa8e8aa3b6fb813317b0d70bf1ed0759eab6f761", size = 31064408 },
+    { url = "https://files.pythonhosted.org/packages/ca/5c/a27f9257f86f0cda3f764ff21d9f4217b9f6a0d45e7a39ecfa7905f524ce/llvmlite-0.43.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6d4fd101f571a31acb1559ae1af30f30b1dc4b3186669f92ad780e17c81e91bc", size = 28793153 },
+    { url = "https://files.pythonhosted.org/packages/7e/3c/4410f670ad0a911227ea2ecfcba9f672a77cf1924df5280c4562032ec32d/llvmlite-0.43.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7d434ec7e2ce3cc8f452d1cd9a28591745de022f931d67be688a737320dfcead", size = 42857276 },
+    { url = "https://files.pythonhosted.org/packages/c6/21/2ffbab5714e72f2483207b4a1de79b2eecd9debbf666ff4e7067bcc5c134/llvmlite-0.43.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6912a87782acdff6eb8bf01675ed01d60ca1f2551f8176a300a886f09e836a6a", size = 43871781 },
+    { url = "https://files.pythonhosted.org/packages/f2/26/b5478037c453554a61625ef1125f7e12bb1429ae11c6376f47beba9b0179/llvmlite-0.43.0-cp310-cp310-win_amd64.whl", hash = "sha256:14f0e4bf2fd2d9a75a3534111e8ebeb08eda2f33e9bdd6dfa13282afacdde0ed", size = 28123487 },
+    { url = "https://files.pythonhosted.org/packages/95/8c/de3276d773ab6ce3ad676df5fab5aac19696b2956319d65d7dd88fb10f19/llvmlite-0.43.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3e8d0618cb9bfe40ac38a9633f2493d4d4e9fcc2f438d39a4e854f39cc0f5f98", size = 31064409 },
+    { url = "https://files.pythonhosted.org/packages/ee/e1/38deed89ced4cf378c61e232265cfe933ccde56ae83c901aa68b477d14b1/llvmlite-0.43.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e0a9a1a39d4bf3517f2af9d23d479b4175ead205c592ceeb8b89af48a327ea57", size = 28793149 },
+    { url = "https://files.pythonhosted.org/packages/2f/b2/4429433eb2dc8379e2cb582502dca074c23837f8fd009907f78a24de4c25/llvmlite-0.43.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c1da416ab53e4f7f3bc8d4eeba36d801cc1894b9fbfbf2022b29b6bad34a7df2", size = 42857277 },
+    { url = "https://files.pythonhosted.org/packages/6b/99/5d00a7d671b1ba1751fc9f19d3b36f3300774c6eebe2bcdb5f6191763eb4/llvmlite-0.43.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:977525a1e5f4059316b183fb4fd34fa858c9eade31f165427a3977c95e3ee749", size = 43871781 },
+    { url = "https://files.pythonhosted.org/packages/20/ab/ed5ed3688c6ba4f0b8d789da19fd8e30a9cf7fc5852effe311bc5aefe73e/llvmlite-0.43.0-cp311-cp311-win_amd64.whl", hash = "sha256:d5bd550001d26450bd90777736c69d68c487d17bf371438f975229b2b8241a91", size = 28107433 },
+    { url = "https://files.pythonhosted.org/packages/0b/67/9443509e5d2b6d8587bae3ede5598fa8bd586b1c7701696663ea8af15b5b/llvmlite-0.43.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:f99b600aa7f65235a5a05d0b9a9f31150c390f31261f2a0ba678e26823ec38f7", size = 31064409 },
+    { url = "https://files.pythonhosted.org/packages/a2/9c/24139d3712d2d352e300c39c0e00d167472c08b3bd350c3c33d72c88ff8d/llvmlite-0.43.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:35d80d61d0cda2d767f72de99450766250560399edc309da16937b93d3b676e7", size = 28793145 },
+    { url = "https://files.pythonhosted.org/packages/bf/f1/4c205a48488e574ee9f6505d50e84370a978c90f08dab41a42d8f2c576b6/llvmlite-0.43.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:eccce86bba940bae0d8d48ed925f21dbb813519169246e2ab292b5092aba121f", size = 42857276 },
+    { url = "https://files.pythonhosted.org/packages/00/5f/323c4d56e8401c50185fd0e875fcf06b71bf825a863699be1eb10aa2a9cb/llvmlite-0.43.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:df6509e1507ca0760787a199d19439cc887bfd82226f5af746d6977bd9f66844", size = 43871781 },
+    { url = "https://files.pythonhosted.org/packages/c6/94/dea10e263655ce78d777e78d904903faae39d1fc440762be4a9dc46bed49/llvmlite-0.43.0-cp312-cp312-win_amd64.whl", hash = "sha256:7a2872ee80dcf6b5dbdc838763d26554c2a18aa833d31a2635bff16aafefb9c9", size = 28107442 },
+]
+
 [[package]]
 name = "mccabe"
 version = "0.7.0"
@@ -236,6 +261,33 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d2/1d/1b658dbd2b9fa9c4c9f32accbfc0205d532c8c6194dc0f2a4c0428e7128a/nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9", size = 22314 },
 ]
 
+[[package]]
+name = "numba"
+version = "0.60.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "llvmlite" },
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/3c/93/2849300a9184775ba274aba6f82f303343669b0592b7bb0849ea713dabb0/numba-0.60.0.tar.gz", hash = "sha256:5df6158e5584eece5fc83294b949fd30b9f1125df7708862205217e068aabf16", size = 2702171 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f7/cf/baa13a7e3556d73d9e38021e6d6aa4aeb30d8b94545aa8b70d0f24a1ccc4/numba-0.60.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5d761de835cd38fb400d2c26bb103a2726f548dc30368853121d66201672e651", size = 2647627 },
+    { url = "https://files.pythonhosted.org/packages/ac/ba/4b57fa498564457c3cc9fc9e570a6b08e6086c74220f24baaf04e54b995f/numba-0.60.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:159e618ef213fba758837f9837fb402bbe65326e60ba0633dbe6c7f274d42c1b", size = 2650322 },
+    { url = "https://files.pythonhosted.org/packages/28/98/7ea97ee75870a54f938a8c70f7e0be4495ba5349c5f9db09d467c4a5d5b7/numba-0.60.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1527dc578b95c7c4ff248792ec33d097ba6bef9eda466c948b68dfc995c25781", size = 3407390 },
+    { url = "https://files.pythonhosted.org/packages/79/58/cb4ac5b8f7ec64200460aef1fed88258fb872ceef504ab1f989d2ff0f684/numba-0.60.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fe0b28abb8d70f8160798f4de9d486143200f34458d34c4a214114e445d7124e", size = 3699694 },
+    { url = "https://files.pythonhosted.org/packages/1c/b0/c61a93ca947d12233ff45de506ddbf52af3f752066a0b8be4d27426e16da/numba-0.60.0-cp310-cp310-win_amd64.whl", hash = "sha256:19407ced081d7e2e4b8d8c36aa57b7452e0283871c296e12d798852bc7d7f198", size = 2687030 },
+    { url = "https://files.pythonhosted.org/packages/98/ad/df18d492a8f00d29a30db307904b9b296e37507034eedb523876f3a2e13e/numba-0.60.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a17b70fc9e380ee29c42717e8cc0bfaa5556c416d94f9aa96ba13acb41bdece8", size = 2647254 },
+    { url = "https://files.pythonhosted.org/packages/9a/51/a4dc2c01ce7a850b8e56ff6d5381d047a5daea83d12bad08aa071d34b2ee/numba-0.60.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3fb02b344a2a80efa6f677aa5c40cd5dd452e1b35f8d1c2af0dfd9ada9978e4b", size = 2649970 },
+    { url = "https://files.pythonhosted.org/packages/f9/4c/8889ac94c0b33dca80bed11564b8c6d9ea14d7f094e674c58e5c5b05859b/numba-0.60.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5f4fde652ea604ea3c86508a3fb31556a6157b2c76c8b51b1d45eb40c8598703", size = 3412492 },
+    { url = "https://files.pythonhosted.org/packages/57/03/2b4245b05b71c0cee667e6a0b51606dfa7f4157c9093d71c6b208385a611/numba-0.60.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4142d7ac0210cc86432b818338a2bc368dc773a2f5cf1e32ff7c5b378bd63ee8", size = 3705018 },
+    { url = "https://files.pythonhosted.org/packages/79/89/2d924ca60dbf949f18a6fec223a2445f5f428d9a5f97a6b29c2122319015/numba-0.60.0-cp311-cp311-win_amd64.whl", hash = "sha256:cac02c041e9b5bc8cf8f2034ff6f0dbafccd1ae9590dc146b3a02a45e53af4e2", size = 2686920 },
+    { url = "https://files.pythonhosted.org/packages/eb/5c/b5ec752c475e78a6c3676b67c514220dbde2725896bbb0b6ec6ea54b2738/numba-0.60.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:d7da4098db31182fc5ffe4bc42c6f24cd7d1cb8a14b59fd755bfee32e34b8404", size = 2647866 },
+    { url = "https://files.pythonhosted.org/packages/65/42/39559664b2e7c15689a638c2a38b3b74c6e69a04e2b3019b9f7742479188/numba-0.60.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38d6ea4c1f56417076ecf8fc327c831ae793282e0ff51080c5094cb726507b1c", size = 2650208 },
+    { url = "https://files.pythonhosted.org/packages/67/88/c4459ccc05674ef02119abf2888ccd3e2fed12a323f52255f4982fc95876/numba-0.60.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:62908d29fb6a3229c242e981ca27e32a6e606cc253fc9e8faeb0e48760de241e", size = 3466946 },
+    { url = "https://files.pythonhosted.org/packages/8b/41/ac11cf33524def12aa5bd698226ae196a1185831c05ed29dc0c56eaa308b/numba-0.60.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0ebaa91538e996f708f1ab30ef4d3ddc344b64b5227b67a57aa74f401bb68b9d", size = 3761463 },
+    { url = "https://files.pythonhosted.org/packages/ca/bd/0fe29fcd1b6a8de479a4ed25c6e56470e467e3611c079d55869ceef2b6d1/numba-0.60.0-cp312-cp312-win_amd64.whl", hash = "sha256:f75262e8fe7fa96db1dca93d53a194a38c46da28b112b8a4aca168f0df860347", size = 2707588 },
+]
+
 [[package]]
 name = "numpy"
 version = "2.0.2"
@@ -510,6 +562,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/3e/df/963384e90733e08eac978cd103c34df181d1fec424de383cdc443f418dd4/scipy-1.13.1-cp312-cp312-win_amd64.whl", hash = "sha256:cdd7dacfb95fea358916410ec61bbc20440f7860333aee6d882bb8046264e949", size = 45910955 },
 ]
 
+[[package]]
+name = "sparse"
+version = "0.15.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numba" },
+    { name = "numpy" },
+    { name = "scipy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/26/6a/a1d00d932597c00509d333d9cde6f10d6c85470a3701455b0c48fc45886b/sparse-0.15.4.tar.gz", hash = "sha256:d4b1c57d24ff0f64f2fd5b5a95b49b7fb84ed207a26d7d58ce2764dcc5c72b84", size = 359721 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1a/f2/8d5bc8cc6b822feac1cd671dac6fb0249a5202ad15ee9549d1a61d4141b5/sparse-0.15.4-py2.py3-none-any.whl", hash = "sha256:76ec76fee2aee82a84eb97155dd530a9644e3b1fdea2406bc4b454698b36d938", size = 237311 },
+]
+
 [[package]]
 name = "tomli"
 version = "2.1.0"