ezmsg-event 0.2__0-py3-none-any.whl

ezmsg/event/__init__.py

@@ -0,0 +1 @@
+from .__version__ import __version__ as __version__

ezmsg/event/__version__.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

ezmsg/event/message.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from numbers import Number
+
+
+@dataclass
+class EventMessage:
+    offset: float
+    """The temporal offset at which the event occurred. This is a float in seconds. The reference point is
+    unspecified and depends on the clock the application uses. Most applications will use time.time."""
+
+    ch_idx: int
+
+    sub_idx: int = 0
+    """The sub-index of the channel. For Blackrock multi-unit data: 0=unsorted, 1-5 sorted unit, >5=noise"""
+
+    value: Number = 1
+    """The value of the event. This can be any number, but is usually an integer, and is often 1 for spikes."""

ezmsg/event/peak.py

@@ -0,0 +1,329 @@
+"""
+Detects peaks in a signal.
+"""
+
+from dataclasses import replace
+import typing
+
+import ezmsg.core as ez
+from ezmsg.sigproc.base import GenAxisArray
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
+from ezmsg.sigproc.scaler import scaler_np
+import numpy as np
+import numpy.typing as npt
+import sparse
+
+from .message import EventMessage
+
+
+@consumer
+def threshold_crossing(
+    threshold: float = -3.5,
+    max_peak_dur: float = 0.002,
+    refrac_dur: float = 0.001,
+    align_on_peak: bool = False,
+    return_peak_val: bool = False,
+    auto_scale_tau: float = 0.0,
+) -> typing.Generator[
+    typing.Union[typing.List[EventMessage], AxisArray], AxisArray, None
+]:
+    """
+    Detect threshold crossing events.
+
+    Args:
+        threshold: the value the signal must cross before the peak is found.
+        max_peak_dur: The maximum duration of a peak in seconds.
+        refrac_dur: The minimum duration between peaks in seconds. If 0 (default), no refractory period is enforced.
+        align_on_peak: If False (default), the returned sample index indicates the first sample across threshold.
+              If True, the sample index indicates the sample with the largest deviation after threshold crossing.
+        return_peak_val: If True then the peak value is included in the EventMessage or sparse matrix payload.
+        auto_scale_tau: If > 0, the data will be passed through a standard scaler prior to thresholding.
+
+    Note: If either align_on_peak or return_peak_val are True then it is necessary to find the actual peak and not
+        just the threshold crossing. This will drastically increase the computational demand. It is recommended to
+        tune max_peak_dur to a minimal-yet-reasonable value to limit the search space.
+
+    Returns:
+        A primed generator object that yields a list of :obj:`EventMessage` objects for every
+        :obj:`AxisArray` it receives via `send`.
+    """
+
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    # Initialize state variables
+    sample_shape: typing.Optional[typing.Tuple[int, ...]] = None
+    fs: typing.Optional[float] = None
+    max_width: int = 0
+    min_width: int = 1  # Consider making this a parameter.
+    refrac_width: int = 0
+
+    scaler: typing.Optional[typing.Generator[AxisArray, AxisArray, None]] = None
+    # adaptive z-scoring.
+    # TODO: This sample-by-sample adaptation is probably overkill. ezmsg-sigproc should add chunk-wise scaler updating.
+
+    _overs: typing.Optional[npt.NDArray] = (
+        None  # (n_feats, <=max_width) int == -1 or +1
+    )
+    # Trailing buffer to track whether the previous sample(s) were past threshold.
+
+    _data: typing.Optional[npt.NDArray] = None  # (n_feats, <=max_width) in_dtype
+    # Trailing buffer in case peak spans sample chunks. Only used if align_on_peak or return_peak_val.
+
+    _data_raw: typing.Optional[npt.NDArray] = None  # (n_feats, <=max_width) in_dtype
+    # Only used if return_peak_val and scaler is not None
+
+    _elapsed: typing.Optional[npt.NDArray] = None  # (n_feats,) int
+    # Number of samples since last event. Used to enforce refractory period across iterations.
+    #
+    # _n_skip: typing.Optional[npt.NDArray] = None  # (n_feats,) int
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+
+        # Extract basic metadata from message
+        ax_idx = msg_in.get_axis_idx("time")
+        in_sample_shape = msg_in.data.shape[:ax_idx] + msg_in.data.shape[ax_idx + 1 :]
+        in_fs = 1 / msg_in.axes["time"].gain
+
+        # If metadata has changed substantially, then reset state variables
+        b_reset = sample_shape is None or sample_shape != in_sample_shape
+        b_reset = b_reset or fs != in_fs
+        if b_reset:
+            sample_shape = in_sample_shape
+            fs = in_fs
+            max_width = int(max_peak_dur * fs)
+            refrac_width = int(refrac_dur * fs)
+            if auto_scale_tau > 0:
+                scaler = scaler_np(time_constant=auto_scale_tau, axis="time")
+            _overs = None
+            _data = None
+            _data_raw = None
+            n_flat_feats = np.prod(sample_shape)
+            _elapsed = np.zeros((n_flat_feats,), dtype=int) + (refrac_width + 1)
+            # _n_skip = np.zeros((n_flat_feats,), dtype=int)
+            # TODO: Support > 2 dim output with pydata.sparse
+            other_dim = "*".join([_ for _ in msg_in.dims if _ != "time"])
+            out_axes = (
+                msg_in.axes.copy()
+                if msg_in.data.ndim == 2
+                else {"time": msg_in.axes["time"]}
+            )
+            template = AxisArray(
+                sparse.SparseArray((0, 0)), dims=[other_dim, "time"], axes=out_axes
+            )
+
+        # Optionally scale data
+        data_raw: typing.Optional[npt.NDArray] = None
+        if scaler is not None:
+            if return_peak_val:
+                data_raw = msg_in.data.copy()
+            msg_in = scaler.send(msg_in)
+
+        data = msg_in.data
+
+        # Put the time axis in the 0th dim.
+        if ax_idx != data.ndim - 1:
+            data = np.moveaxis(data, ax_idx, -1)
+            if data_raw is not None:
+                data_raw = np.moveaxis(data_raw, ax_idx, -1)
+
+        # Flatten the feature dimensions
+        if data.ndim > 2:
+            data = data.reshape((-1, data.shape[-1]))
+            if data_raw is not None:
+                data_raw = data_raw.reshape((-1, data_raw.shape[-1]))
+
+        # Check each sample relative to threshold.
+        overs = data >= threshold if threshold >= 0 else data <= threshold
+
+        # Prepend our variables with previous iteration's values. We always expect at least 1 sample to carry over.
+        overs = np.concatenate(
+            (_overs if _overs is not None else overs[..., :1], overs), axis=-1
+        )
+        # If we need to identify _where_ the peak was then we must prepend previous data values.
+        if align_on_peak or return_peak_val:
+            data = np.concatenate(
+                (_data if _data is not None else data[..., :1], data), axis=-1
+            )
+        # If we're doing z-scoring but we need the actual peak value then we must prepend previous RAW values too.
+        if return_peak_val and scaler is not None:
+            data_raw = np.concatenate(
+                (_data_raw if _data_raw is not None else data_raw[..., :1], data_raw),
+                axis=-1,
+            )
+        # We will modify _overs later, so for now we take note of how many samples were prepended.
+        n_prepended = 1 if _overs is None else _overs.shape[-1]
+
+        # Find threshold crossing where sample k is over and sample k-1 is not over.
+        b_cross_over = np.logical_and(~overs[..., :-1], overs[..., 1:])
+
+        feat_idx, samp_idx = np.where(b_cross_over)
+        # Because we looked at samples [1:]...
+        samp_idx += 1
+        # Note: There is an assumption that the 0th sample only serves as a reference and is not part of the output;
+        #  this will be trimmed at the very end. For now the offset is useful for bookkeeping (peak finding, etc.).
+
+        # Optionally drop crossings during refractory period
+        if refrac_width > 2 and len(samp_idx) > 0:
+            # TODO: Consider putting this into its own unit. The downside to moving it out is that some of the remaining
+            #  computation in this function is unnecessary for to-be-dropped events.
+            uq_feats, feat_splits = np.unique(feat_idx, return_index=True)
+            ieis = np.diff(np.hstack(([samp_idx[0] + 1], samp_idx)))
+            # Reset elapsed time at feature boundaries.
+            ieis[feat_splits] = samp_idx[feat_splits] + _elapsed[uq_feats]
+            b_drop = ieis <= refrac_width
+            drop_idx = np.where(b_drop)[0]
+            final_drop = []
+            while len(drop_idx) > 0:
+                d_idx = drop_idx[0]
+                # Update next iei so its interval refers to the event before the to-be-dropped event.
+                #  but only if the next iei belongs to the same feature.
+                if ((d_idx + 1) < len(ieis)) and (d_idx + 1) not in feat_splits:
+                    ieis[d_idx + 1] += ieis[d_idx]
+                # We will later remove this event from samp_idx and feat_idx
+                final_drop.append(d_idx)
+                # Remove the dropped event from drop_idx.
+                drop_idx = drop_idx[1:]
+
+                # If the next event is now outside the refractory period then it will not be dropped.
+                # TODO: Maybe we don't want this feature? e.g., if the 3rd in a triplet is not within the refractory
+                #  period of the first, but it is within the refractory period of the 2nd which is now dropped,
+                #  maybe the 3rd event should still be dropped? If so then we can refactor above to drop all at once.
+                if len(drop_idx) > 0 and ieis[drop_idx[0]] > refrac_width:
+                    drop_idx = drop_idx[1:]
+
+            samp_idx = np.delete(samp_idx, final_drop)
+            feat_idx = np.delete(feat_idx, final_drop)
+
+        hold_idx = overs.shape[-1] - 1
+        if len(samp_idx) == 0:
+            # No events.
+            result_val = np.ones(
+                samp_idx.shape, dtype=data.dtype if return_peak_val else bool
+            )
+        elif not (min_width > 1 or align_on_peak or return_peak_val):
+            # No postprocessing required.
+            result_val = np.ones(samp_idx.shape, dtype=bool)
+        else:
+            # Do postprocessing of events.
+
+            # We extract max_width-length vectors of `overs` values for each event.
+            # Pad using last sample until last crossover has at least max_width following samples.
+            n_pad = max(0, max(samp_idx) + max_width - overs.shape[-1])
+            pad_width = ((0, 0),) * (overs.ndim - 1) + ((0, n_pad),)
+            overs_padded = np.pad(overs, pad_width, mode="edge")
+
+            # Multi-index to extract the vectors for each event
+            s_idx = np.arange(max_width)[None, :] + samp_idx[:, None]
+            ep_overs = overs_padded[feat_idx[:, None], s_idx]  # (n_events, max_width)
+
+            # Find the event lengths: i.e., the first non-over-threshold value for each event.
+            # Warning: Values are invalid for events that don't cross back.
+            ev_len = ep_overs[..., 1:].argmin(axis=-1)
+            ev_len += 1
+
+            # Identify peaks that successfully cross back
+            b_ev_crossback = np.any(~ep_overs[..., 1:], axis=-1)
+
+            if min_width > 1:
+                # Drop events that have crossed back but fail min_width
+                b_short = np.logical_and(b_ev_crossback, ev_len < min_width)
+                b_long = ~b_short
+                samp_idx = samp_idx[b_long]
+                feat_idx = feat_idx[b_long]
+                ev_len = ev_len[b_long]
+                b_ev_crossback = b_ev_crossback[b_long]
+
+            # We are returning a sparse array and unfinished peaks must be buffered for the next iteration.
+            # Find the earliest unfinished event. If none, we still buffer the final sample.
+            b_unf = ~b_ev_crossback
+            hold_idx = samp_idx[b_unf].min() if np.any(b_unf) else hold_idx
+
+            # Trim events that are past the hold_idx. They will be processed next iteration.
+            b_pass_ev = samp_idx < hold_idx
+            samp_idx = samp_idx[b_pass_ev]
+            feat_idx = feat_idx[b_pass_ev]
+            ev_len = ev_len[b_pass_ev]
+
+            if np.any(b_unf):
+                # Must hold back at least 1 sample before start of unfinished events so we can re-detect.
+                hold_idx = max(hold_idx - 1, 0)
+
+            if not return_peak_val:
+                result_val = np.ones(samp_idx.shape, dtype=bool)
+
+            # For remaining _finished_ peaks, get the peak location -- for alignment or if returning its value.
+            if align_on_peak or return_peak_val:
+                # We process peaks in batches based on their length, otherwise short peaks could give
+                #  incorrect argmax results.
+                # TODO: Check performance of using a masked array instead. Might take longer to create the mask.
+                pk_offset = np.zeros_like(ev_len)
+                uq_lens, len_grps = np.unique(ev_len, return_inverse=True)
+                for len_idx, ep_len in enumerate(uq_lens):
+                    b_grp = len_grps == len_idx
+                    ep_resamp = np.arange(ep_len)[None, :] + samp_idx[b_grp, None]
+                    eps = data[feat_idx[b_grp, None], ep_resamp]
+                    if threshold >= 0:
+                        pk_offset[b_grp] = np.argmax(eps, axis=1)
+                    else:
+                        pk_offset[b_grp] = np.argmin(eps, axis=1)
+
+                # Now that we have the offset to the peak for each event, we can find the peak value and/or align.
+                if return_peak_val:
+                    if scaler is None:
+                        result_val = data[feat_idx, samp_idx + pk_offset]
+                    else:
+                        result_val = data_raw[feat_idx, samp_idx + pk_offset]
+                if align_on_peak:
+                    samp_idx += pk_offset
+
+        # Save data for next iteration
+        _overs = overs[..., hold_idx:]
+        if align_on_peak or return_peak_val:
+            _data = data[..., hold_idx:]
+            if return_peak_val and scaler is not None:
+                _data_raw = data_raw[..., hold_idx:]
+        _elapsed += hold_idx
+        _elapsed[feat_idx] = hold_idx - samp_idx
+        # ^ multiple-write to same index is fine because last value is desired.
+
+        # Prepare sparse matrix output
+        # Note: The first of the "held" samples is part of this iteration's return.
+        #  Likewise, the first prepended sample was part of the previous iteration's return.
+        n_out_samps = hold_idx
+        t0 = msg_in.axes["time"].offset - (n_prepended - 1) * msg_in.axes["time"].gain
+        samp_idx -= 1  # Discard first prepended sample.
+        result = sparse.COO(
+            (feat_idx, samp_idx),
+            data=result_val,
+            shape=data.shape[:-1] + (n_out_samps,),
+        )
+        msg_out = replace(
+            template,
+            data=result,
+            axes={**template.axes, "time": replace(template.axes["time"], offset=t0)},
+        )
+
+
+class ThresholdSettings(ez.Settings):
+    threshold: float = -3.5
+    max_peak_dur: float = 0.002
+    refrac_dur: float = 0.001
+    align_on_peak: bool = False
+    return_peak_val: bool = False
+    auto_scale_tau: float = 0.0
+
+
+class ThresholdCrossing(GenAxisArray):
+    SETTINGS = ThresholdSettings
+
+    def construct_generator(self):
+        self.STATE.gen = threshold_crossing(
+            threshold=self.SETTINGS.threshold,
+            max_peak_dur=self.SETTINGS.max_peak_dur,
+            refrac_dur=self.SETTINGS.refrac_dur,
+            align_on_peak=self.SETTINGS.align_on_peak,
+            return_peak_val=self.SETTINGS.return_peak_val,
+            auto_scale_tau=self.SETTINGS.auto_scale_tau,
+        )

ezmsg/event/rate.py

@@ -0,0 +1,82 @@
+"""
+Count number of events in a given time window. Optionally, divide by window duration to get rate.
+"""
+
+from dataclasses import replace
+import typing
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
+from ezmsg.sigproc.base import GenAxisArray
+
+from .window import windowing
+
+
+@consumer
+def event_rate(
+    bin_duration: float = 0.05,
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+
+
+    Args:
+        bin_duration:
+
+    Returns:
+        A primed generator object that yields an :obj:`AxisArray` object of event rates for every
+        :obj:`AxisArray` of sparse events it receives via `send`.
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    win_proc = windowing(
+        axis="time",
+        newaxis="win",
+        window_dur=bin_duration,
+        window_shift=bin_duration,
+        zero_pad_until="none",
+    )
+    out_dims: typing.Optional[typing.List[str]] = None
+    out_axes: typing.Optional[typing.Dict[str, AxisArray.Axis]] = None
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+
+        win_msg = win_proc.send(msg_in)
+
+        b_reset = out_dims is None
+        if b_reset:
+            # Fixup `dims`
+            out_dims = list(win_msg.dims)
+            out_dims.remove("time")
+            out_dims[out_dims.index("win")] = "time"
+            # Fixup axes
+            out_axes = {k: v for k, v in win_msg.axes.items() if k != "time"}
+
+        # Sum over time
+        time_ax = win_msg.get_axis_idx("time")
+        counts_per_bin = np.sum(win_msg.data, axis=time_ax)
+        # Scale by 1 / bin_duration to get rates
+        rates_per_bin = counts_per_bin / bin_duration
+        # Densify
+        rates_per_bin = rates_per_bin.todense()
+        msg_out = replace(
+            win_msg,
+            data=rates_per_bin,
+            dims=out_dims,
+            axes={**out_axes, "time": win_msg.axes["win"]},
+        )
+
+
+class EventRateSettings(ez.Settings):
+    bin_duration: float = 0.05
+
+
+class EventRate(GenAxisArray):
+    SETTINGS = EventRateSettings
+
+    def construct_generator(self):
+        self.STATE.gen = event_rate(
+            bin_duration=self.SETTINGS.bin_duration,
+        )

ezmsg/event/sparse.py

@@ -0,0 +1,32 @@
+from dataclasses import replace
+import typing
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.sigproc.base import GenAxisArray
+from ezmsg.util.generator import consumer
+from ezmsg.util.messages.axisarray import AxisArray
+
+
+@consumer
+def to_dense() -> typing.Generator[AxisArray, AxisArray, None]:
+    msg_out = AxisArray(np.array([]), dims=[""])
+    while True:
+        msg_in: AxisArray = yield msg_out
+        if hasattr(msg_in.data, "todense"):
+            msg_out = replace(msg_in, data=msg_in.data.todense())
+        else:
+            msg_out = msg_in
+
+
+class DensifySettings(ez.Settings):
+    pass
+
+
+class Densify(GenAxisArray):
+    """:obj:`Unit` for :obj:`bandpower`."""
+
+    SETTINGS = DensifySettings
+
+    def construct_generator(self):
+        self.STATE.gen = to_dense()

ezmsg/event/util/array.py

@@ -0,0 +1,17 @@
+import sparse
+
+
+def sliding_win_oneaxis(
+    s: sparse.SparseArray, nwin: int, axis: int, step: int = 1
+) -> sparse.SparseArray:
+    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 working with coords and values only.
+    return result

ezmsg/event/util/simulate.py

@@ -0,0 +1,117 @@
+import numpy as np
+import numpy.typing as npt
+
+
+def generate_events(
+    fs: float,
+    dur: float,
+    n_chans: int,
+    rate_range: tuple[float, float],
+    chunk_dur: float,
+) -> list[npt.NDArray]:
+    n_times = int(fs * dur)
+    frates = np.random.uniform(rate_range[0], rate_range[1], n_chans)
+    frates[:3] = np.random.uniform(150, 200, 3)  # Boost rate of first 3 chans.
+    chunk_len = int(fs * chunk_dur)
+
+    # Create a list of spike times for each channel
+    rng = np.random.default_rng()
+    spike_offsets = []
+    for ch_ix, fr in enumerate(frates):
+        lam, size = fs / fr, int(fr * dur)
+        isi = rng.poisson(lam=lam, size=size)
+        spike_samp_inds = np.cumsum(isi)
+        spike_samp_inds = spike_samp_inds[spike_samp_inds < n_times]
+
+        # Add some special cases
+        if ch_ix == 0:
+            # -- Refractory within chunk --
+            # In channel 0, we replace the first event with a triplet; events 2-3 will be eliminated by refractory check
+            spike_samp_inds = spike_samp_inds[spike_samp_inds > 30]
+            spike_samp_inds = np.hstack(([1, 4, 6], spike_samp_inds))
+        elif ch_ix in [1, 2]:
+            # -- Unfinished events at chunk boundaries --
+            # Drop spike samples within 34 samples of the end of the 0th chunk
+            b_drop = np.logical_and(
+                spike_samp_inds >= chunk_len - 34, spike_samp_inds < chunk_len
+            )
+            spike_samp_inds = spike_samp_inds[~b_drop]
+            if ch_ix == 1:
+                # In channel 1, we add a spike that is in the very last sample of the 0th chunk.
+                # It will be detected while processing the 1th chunk.
+                spike_samp_inds = np.insert(
+                    spike_samp_inds,
+                    np.searchsorted(spike_samp_inds, chunk_len),
+                    chunk_len - 1,
+                )
+            elif ch_ix == 2:
+                # In channel 2, we make a long event at the end of the 0th chunk.
+                # It will be detected while processing the 1th chunk.
+                spike_samp_inds = np.insert(
+                    spike_samp_inds,
+                    np.searchsorted(spike_samp_inds, chunk_len - 10),
+                    np.arange(chunk_len - 10, chunk_len),
+                )
+        elif ch_ix == 3:
+            # -- Refractoriness across chunk boundaries --
+            # In channel 3, we add a spike 2 samples before the end of 1th chunk, and another within its
+            #  refractory period at the beginning of 2th chunk.
+            ins_ev_start = 2 * chunk_len - 2
+            # Clear events that are within target period.
+            b_drop = np.logical_and(
+                spike_samp_inds >= ins_ev_start - 30,
+                spike_samp_inds < ins_ev_start + 30,
+            )
+            spike_samp_inds = spike_samp_inds[~b_drop]
+            # Add the two events; one 2 samples before the chunk boundary and another 10 samples later.
+            spike_samp_inds = np.insert(
+                spike_samp_inds,
+                np.searchsorted(spike_samp_inds, ins_ev_start),
+                [ins_ev_start, ins_ev_start + 10],
+            )
+            # Note: We must also drop events in other channels near the end of chunk 2 to make sure
+            #  they don't cause the event in channel 3 to be held back to the next iteration.
+        elif ch_ix == 4:
+            # -- Spike in first sample of non-first chunk --
+            # In channel 4, we add a spike at the very beginning of chunk 1th chunk after making sure 0th was empty.
+            spike_samp_inds = spike_samp_inds[spike_samp_inds > chunk_len]
+            spike_samp_inds = np.insert(
+                spike_samp_inds, np.searchsorted(spike_samp_inds, chunk_len), chunk_len
+            )
+        spike_offsets.append(spike_samp_inds)
+
+    # Clear all events that occur in 4th - 5th chunks to test flow logic.
+    # Additionally clear events in the last sample so we don't have lingering events.
+    for ch_ix, so_arr in enumerate(spike_offsets):
+        b_drop = np.logical_and(so_arr >= chunk_len * 3, so_arr < chunk_len * 5)
+        b_drop = np.logical_or(b_drop, so_arr == n_times - 1)
+        if ch_ix != 3:
+            # See above for special case in channel 3
+            b_drop = np.logical_or(
+                b_drop,
+                np.logical_and(so_arr >= 2 * chunk_len - 30, so_arr < 2 * chunk_len),
+            )
+        spike_offsets[ch_ix] = so_arr[~b_drop]
+
+    return spike_offsets
+
+
+def generate_white_noise_with_events(
+    fs: float,
+    dur: float,
+    n_chans: int,
+    rate_range: tuple[float, float],
+    chunk_dur: float,
+    threshold: float,
+) -> npt.NDArray:
+    n_times = int(fs * dur)
+    spike_offsets = generate_events(fs, dur, n_chans, rate_range, chunk_dur)
+
+    rng = np.random.default_rng()
+    mixed = rng.normal(size=(n_times, n_chans), loc=0, scale=0.1)
+    mixed = np.clip(mixed, -np.abs(threshold), np.abs(threshold))
+    for ch_ix, ch_spk_offs in enumerate(spike_offsets):
+        mixed[ch_spk_offs, ch_ix] = threshold + np.random.random(
+            size=(len(ch_spk_offs),)
+        )
+    return mixed

ezmsg/event/window.py

@@ -0,0 +1,225 @@
+from dataclasses import replace
+import typing
+
+import numpy as np
+import sparse
+import ezmsg.core as ez
+from ezmsg.sigproc.base import GenAxisArray
+from ezmsg.util.generator import consumer
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    slice_along_axis,
+    sliding_win_oneaxis,
+)
+
+from .util.array import sliding_win_oneaxis as sparse_sliding_win
+
+
+@consumer
+def windowing(
+    axis: typing.Optional[str] = None,
+    newaxis: str = "win",
+    window_dur: typing.Optional[float] = None,
+    window_shift: typing.Optional[float] = None,
+    zero_pad_until: str = "input",
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Construct a generator that yields windows of data from an input :obj:`AxisArray`.
+
+    Args:
+        axis: The axis along which to segment windows.
+            If None, defaults to the first dimension of the first seen AxisArray.
+        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.
+
+    Returns:
+        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.
+    """
+    # Check arguments
+    if newaxis is None:
+        ez.logger.warning("`newaxis` must not be None. Setting to 'win'.")
+        newaxis = "win"
+    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."
+        )
+
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    # State variables
+    buffer: typing.Optional[sparse.SparseArray] = None
+    window_samples: typing.Optional[int] = None
+    window_shift_samples: typing.Optional[int] = None
+    shift_deficit: int = 0
+    b_1to1 = window_shift is None
+    newaxis_warned: bool = b_1to1
+    out_newaxis: typing.Optional[AxisArray.Axis] = None
+    out_dims: typing.Optional[typing.List[str]] = None
+
+    check_inputs = {"samp_shape": None, "fs": None, "key": None}
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+
+        if window_dur is None:
+            msg_out = msg_in
+            continue
+
+        axis = axis or msg_in.dims[0]
+        axis_idx = msg_in.get_axis_idx(axis)
+        axis_info = msg_in.get_axis(axis)
+        fs = 1.0 / axis_info.gain
+
+        if not newaxis_warned and newaxis in msg_in.dims:
+            ez.logger.warning(
+                f"newaxis {newaxis} present in input dims. Using {newaxis}_win instead"
+            )
+            newaxis_warned = True
+            newaxis = f"{newaxis}_win"
+
+        samp_shape = msg_in.data.shape[:axis_idx] + msg_in.data.shape[axis_idx + 1 :]
+
+        # If buffer unset or input stats changed, create a new buffer
+        b_reset = buffer is None
+        b_reset = b_reset or samp_shape != check_inputs["samp_shape"]
+        b_reset = b_reset or fs != check_inputs["fs"]
+        b_reset = b_reset or msg_in.key != check_inputs["key"]
+        if b_reset:
+            # Update check variables
+            check_inputs["samp_shape"] = samp_shape
+            check_inputs["fs"] = fs
+            check_inputs["key"] = msg_in.key
+
+            window_samples = int(window_dur * fs)
+            if not b_1to1:
+                window_shift_samples = int(window_shift * 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 = msg_in.data.shape[axis_idx]
+            n_zero = max(0, window_samples - req_samples)
+            buffer = sparse.zeros(
+                msg_in.data.shape[:axis_idx]
+                + (n_zero,)
+                + msg_in.data.shape[axis_idx + 1 :]
+            )
+
+        # Add new data to buffer.
+        buffer = sparse.concatenate((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]]
+        # 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, slice(n_skip, None), axis_idx)
+                buffer_offset = buffer_offset[n_skip:]
+                shift_deficit -= n_skip
+
+        # Prepare reusable parts of output
+        if out_newaxis is None:
+            out_dims = msg_in.dims[:axis_idx] + [newaxis] + msg_in.dims[axis_idx:]
+            out_newaxis = replace(
+                axis_info,
+                gain=0.0 if b_1to1 else axis_info.gain * window_shift_samples,
+                offset=0.0,  # offset modified per-msg below
+            )
+
+        # Generate outputs.
+        # Preliminary copy of axes without the axes that we are modifying.
+        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)
+
+        # How we update .data and .axes[newaxis] depends on the windowing mode.
+        if b_1to1:
+            # one-to-one mode -- Each send yields exactly one window containing only the most recent samples.
+            buffer = slice_along_axis(buffer, slice(-window_samples, None), axis_idx)
+            out_dat = buffer.reshape(
+                buffer.shape[:axis_idx] + (1,) + buffer.shape[axis_idx:]
+            )
+            out_newaxis = replace(out_newaxis, offset=buffer_offset[-window_samples])
+        elif buffer.shape[axis_idx] >= window_samples:
+            # Deterministic window shifts.
+            out_dat = sparse_sliding_win(
+                buffer, window_samples, axis_idx, step=window_shift_samples
+            )
+            offset_view = sliding_win_oneaxis(buffer_offset, window_samples, 0)[
+                ::window_shift_samples
+            ]
+            out_newaxis = replace(out_newaxis, offset=offset_view[0, 0])
+
+            # Drop expired beginning of buffer and update shift_deficit
+            multi_shift = window_shift_samples * out_dat.shape[axis_idx]
+            shift_deficit = max(0, multi_shift - buffer.shape[axis_idx])
+            buffer = slice_along_axis(buffer, slice(multi_shift, None), axis_idx)
+        else:
+            # Not enough data to make a new window. Return empty data.
+            empty_data_shape = (
+                msg_in.data.shape[:axis_idx]
+                + (0, window_samples)
+                + msg_in.data.shape[axis_idx + 1 :]
+            )
+            out_dat = sparse.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)
+
+        msg_out = replace(
+            msg_in, data=out_dat, dims=out_dims, axes={**out_axes, newaxis: out_newaxis}
+        )
+
+
+class WindowSettings(ez.Settings):
+    axis: typing.Optional[str] = None
+    newaxis: typing.Optional[str] = None  # new axis for output. No new axes if None
+    window_dur: typing.Optional[float] = None  # Sec. passthrough if None
+    window_shift: typing.Optional[float] = None  # Sec. Use "1:1 mode" if None
+    zero_pad_until: str = "full"  # "full", "shift", "input", "none"
+
+
+class Window(GenAxisArray):
+    """:obj:`Unit` for :obj:`bandpower`."""
+
+    SETTINGS = WindowSettings
+
+    def construct_generator(self):
+        self.STATE.gen = windowing(
+            axis=self.SETTINGS.axis,
+            newaxis=self.SETTINGS.newaxis,
+            window_dur=self.SETTINGS.window_dur,
+            window_shift=self.SETTINGS.window_shift,
+            zero_pad_until=self.SETTINGS.zero_pad_until,
+        )

ezmsg_event-0.2.0.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.3
+Name: ezmsg-event
+Version: 0.2.0
+Summary: Namespace package for ezmsg to work with signal events like neural spikes and heartbeats
+Author-email: Chadwick Boulay <chadwick.boulay@gmail.com>
+Requires-Python: >=3.10
+Requires-Dist: ezmsg-sigproc>=1.5.0
+Requires-Dist: numpy>=1.26
+Requires-Dist: sparse>=0.15.4
+Provides-Extra: test
+Requires-Dist: pytest>=8.3.3; extra == 'test'
+Description-Content-Type: text/markdown
+
+# ezmsg-event
+
+ezmsg namespace package for working with signal events like neural spikes and heartbeats

ezmsg_event-0.2.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+ezmsg/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ezmsg/event/__init__.py,sha256=8K4IcOA3-pfzadoM6s2Sfg5460KlJUocGgyTJTJl96U,52
+ezmsg/event/__version__.py,sha256=H-qsvrxCpdhaQzyddR-yajEqI71hPxLa4KxzpP3uS1g,411
+ezmsg/event/message.py,sha256=IXjmuWECZaU0NY0vb52hWbJvTP-d2NaKurBg4nOg7ZM,606
+ezmsg/event/peak.py,sha256=EHENdpa9huKJIBiUHuzFLbq2HptY1RKgE96eNnkOe2w,15065
+ezmsg/event/rate.py,sha256=uyj7M6R7_pJSGo4_GY74snrM_8eXYnvVtUjwrUScRgU,2224
+ezmsg/event/sparse.py,sha256=FaRM4YLpFLqkLU0tNjrLvjh4Om9EzPs8rOpW8ajVc48,790
+ezmsg/event/window.py,sha256=A8B412jczoxqEQa95d12_fU2pzeal-a-mW2NakP2qmA,9827
+ezmsg/event/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ezmsg/event/util/array.py,sha256=ZBqaf4jmsXSb56D9Hea7cWD9qJvszfv3x9X3NIUB1f8,653
+ezmsg/event/util/simulate.py,sha256=-fBfFeEktjAn-mtNR_6E0QjijHMQTg08nn1q2B-ywqE,5127
+ezmsg_event-0.2.0.dist-info/METADATA,sha256=ir4Vx3O6CV6w6s_Lf84hyn-nrNVngYcNsaDEJuFEyqA,540
+ezmsg_event-0.2.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+ezmsg_event-0.2.0.dist-info/licenses/LICENSE,sha256=KImy0vSiKLkW28zWqI8rIaLwvjGKxOjE5kZmVxZDdy4,1066
+ezmsg_event-0.2.0.dist-info/RECORD,,

ezmsg_event-0.2.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.26.3
+Root-Is-Purelib: true
+Tag: py3-none-any

ezmsg_event-0.2.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 ezmsg-org
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.