neurospatial.behavior.epochs

epochs

Array-native epoch selection (restrict / in_epochs).

"Give me my running periods" / "give me trial N" -- restrict time-series and spike data to a set of time intervals (epochs). Epochs are plain (start, end) arrays, so this stays array-first and never imports pynapple; a pynapple IntervalSet is nonetheless accepted transparently because it is duck-typed (it exposes .start / .end arrays), not isinstance-d.

Functions:

Name	Description
restrict	The headline one-liner: t, pos = restrict(times, positions, epochs=run_epochs). Slices times and any number of arrays aligned to times (e.g. positions) by the same in-epoch mask, order preserved. With no extra arrays it restricts an event-time array by its own timestamps (restrict(spike_train, epochs=...) -> the in-epoch spikes).
in_epochs	The boolean mask primitive: True where t falls inside any interval. Inclusive endpoints by default (closed="both"), matching :mod:neurospatial.behavior.segmentation and pynapple.
restrict_spike_trains	Ragged per-unit spike times are not aligned to a common times; each train is restricted by its own timestamps. Accepts a plain sequence of trains or a :class:~neurospatial.encoding.SpikeTrains container.

Aligned vs ragged

restrict is for arrays that share one time axis (position samples, a single spike train). restrict_spike_trains is for a ragged collection where each unit has its own timestamps and there is no shared times -- so each train is masked against itself.

Functions

in_epochs

in_epochs(t: NDArray[float64], epochs: Any, *, closed: _Closed = 'both') -> NDArray[np.bool_]

Boolean mask: True where each element of t falls in any epoch.

Vectorized over samples (broadcasting, no Python loop) and taking the union across intervals (overlapping intervals behave as their union).

Parameters:

Name	Type	Description	Default
t	ndarray	Timestamps to test (any shape; the mask has the same shape).	required
epochs	IntervalSet-like, tuple, list, or ndarray	Epochs in any form accepted by :func:_as_intervals ((start, end) scalars/arrays, an (n, 2) array, or a duck-typed IntervalSet). Empty epochs -> all-False.	required
closed	('both', 'left', 'right', 'neither')	Which endpoints are inclusive. Default "both" (start <= t <= end), matching :mod:neurospatial.behavior.segmentation and pynapple.	"both"

Returns:

Type	Description
ndarray of bool	Same shape as t; True where t is inside any interval.

Notes

NaN timestamps evaluate False at every >=/<=/>/< comparison, so they are (deliberately) excluded from every epoch.

Examples:

>>> import numpy as np
>>> from neurospatial.behavior import in_epochs
>>> in_epochs(np.array([0.0, 2.5, 5.0, 6.0]), (0.0, 5.0))
array([ True,  True,  True, False])

Source code in src/neurospatial/behavior/epochs.py

def in_epochs(
    t: NDArray[np.float64],
    epochs: Any,
    *,
    closed: _Closed = "both",
) -> NDArray[np.bool_]:
    """Boolean mask: ``True`` where each element of ``t`` falls in any epoch.

    Vectorized over samples (broadcasting, no Python loop) and taking the
    **union** across intervals (overlapping intervals behave as their union).

    Parameters
    ----------
    t : ndarray
        Timestamps to test (any shape; the mask has the same shape).
    epochs : IntervalSet-like, tuple, list, or ndarray
        Epochs in any form accepted by :func:`_as_intervals`
        (``(start, end)`` scalars/arrays, an ``(n, 2)`` array, or a
        duck-typed ``IntervalSet``). Empty epochs -> all-``False``.
    closed : {"both", "left", "right", "neither"}, optional
        Which endpoints are inclusive. Default ``"both"`` (``start <= t <=
        end``), matching :mod:`neurospatial.behavior.segmentation` and pynapple.

    Returns
    -------
    ndarray of bool
        Same shape as ``t``; ``True`` where ``t`` is inside any interval.

    Notes
    -----
    NaN timestamps evaluate ``False`` at every ``>=``/``<=``/``>``/``<``
    comparison, so they are (deliberately) excluded from every epoch.

    Examples
    --------
    >>> import numpy as np
    >>> from neurospatial.behavior import in_epochs
    >>> in_epochs(np.array([0.0, 2.5, 5.0, 6.0]), (0.0, 5.0))
    array([ True,  True,  True, False])
    """
    # Validate ``closed`` up front -- before ``_as_intervals`` and the
    # empty-epochs early return -- so a bad value always raises (even with no
    # intervals) AND is reported before a malformed ``epochs``.
    if closed not in ("both", "left", "right", "neither"):
        raise ValueError(
            f"closed must be one of 'both', 'left', 'right', 'neither', got {closed!r}."
        )

    intervals = _as_intervals(epochs)
    return _mask_in_intervals(t, intervals, closed=closed)

restrict

restrict(times: NDArray[float64], *arrays: NDArray[Any], epochs: Any, closed: _Closed = 'both') -> NDArray[Any] | tuple[NDArray[Any], ...]

Restrict times and time-aligned arrays to a set of epochs.

The headline one-liner::

t, pos = restrict(times, positions, epochs=run_epochs)

times is the reference 1-D time axis. Each of *arrays must be aligned to times (len(arr) == len(times); e.g. positions of shape (n, n_dims)). All are sliced by the single mask in_epochs(times, epochs, closed=closed), so alignment and order are preserved.

With no extra arrays, restrict restricts an event-time array by its own timestamps -- restrict(spike_train, epochs=...) returns the in-epoch spikes (here times is the spike train). For ragged per-unit spikes (each unit its own timestamps, no shared axis), use :func:restrict_spike_trains instead.

Parameters:

Name	Type	Description	Default
times	(ndarray, shape(n))	Reference time array.	required
*arrays	ndarray	Zero or more arrays aligned to times (first axis length n).	()
epochs	IntervalSet-like, tuple, list, or ndarray	Epochs in any form accepted by :func:_as_intervals.	required
closed	('both', 'left', 'right', 'neither')	Endpoint inclusivity, forwarded to :func:in_epochs. Default "both".	"both"

Returns:

Type	Description
ndarray or tuple of ndarray	If no extra arrays were passed, the masked times (a bare array). Otherwise (times_kept, *arrays_kept) -- each sliced by the same mask, order preserved.

Raises:

Type	Description
ValueError	If any array's first-axis length differs from len(times) (the error names the offending positional index).

Examples:

>>> import numpy as np
>>> from neurospatial.behavior import restrict
>>> times = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
>>> positions = np.column_stack([times, times])
>>> t, pos = restrict(times, positions, epochs=(1.0, 3.0))
>>> t
array([1., 2., 3.])

Source code in src/neurospatial/behavior/epochs.py

def restrict(
    times: NDArray[np.float64],
    *arrays: NDArray[Any],
    epochs: Any,
    closed: _Closed = "both",
) -> NDArray[Any] | tuple[NDArray[Any], ...]:
    """Restrict ``times`` and time-aligned arrays to a set of epochs.

    The headline one-liner::

        t, pos = restrict(times, positions, epochs=run_epochs)

    ``times`` is the reference 1-D time axis. Each of ``*arrays`` must be
    **aligned to ``times``** (``len(arr) == len(times)``; e.g. ``positions`` of
    shape ``(n, n_dims)``). All are sliced by the single mask
    ``in_epochs(times, epochs, closed=closed)``, so alignment and order are
    preserved.

    With **no** extra arrays, ``restrict`` restricts an event-time array by its
    own timestamps -- ``restrict(spike_train, epochs=...)`` returns the in-epoch
    spikes (here ``times`` *is* the spike train). For ragged per-unit spikes
    (each unit its own timestamps, no shared axis), use
    :func:`restrict_spike_trains` instead.

    Parameters
    ----------
    times : ndarray, shape (n,)
        Reference time array.
    *arrays : ndarray
        Zero or more arrays aligned to ``times`` (first axis length ``n``).
    epochs : IntervalSet-like, tuple, list, or ndarray
        Epochs in any form accepted by :func:`_as_intervals`.
    closed : {"both", "left", "right", "neither"}, optional
        Endpoint inclusivity, forwarded to :func:`in_epochs`. Default
        ``"both"``.

    Returns
    -------
    ndarray or tuple of ndarray
        If no extra arrays were passed, the masked ``times`` (a bare array).
        Otherwise ``(times_kept, *arrays_kept)`` -- each sliced by the same
        mask, order preserved.

    Raises
    ------
    ValueError
        If any array's first-axis length differs from ``len(times)`` (the error
        names the offending positional index).

    Examples
    --------
    >>> import numpy as np
    >>> from neurospatial.behavior import restrict
    >>> times = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
    >>> positions = np.column_stack([times, times])
    >>> t, pos = restrict(times, positions, epochs=(1.0, 3.0))
    >>> t
    array([1., 2., 3.])
    """
    times_arr = np.asarray(times)
    n = len(times_arr)

    coerced: list[NDArray[Any]] = []
    for i, arr in enumerate(arrays):
        arr_np = np.asarray(arr)
        if len(arr_np) != n:
            raise ValueError(
                f"array at position {i} has len {len(arr_np)} but must be "
                f"aligned to times (len {n}); every *arrays entry must have the "
                "same first-axis length as times."
            )
        coerced.append(arr_np)

    mask = in_epochs(times_arr, epochs, closed=closed)

    times_kept = times_arr[mask]
    if not coerced:
        return times_kept
    return (times_kept, *(arr[mask] for arr in coerced))

restrict_spike_trains

restrict_spike_trains(trains: Sequence[NDArray[float64]], epochs: Any, *, closed: _Closed = 'both') -> list[NDArray[np.float64]]

Restrict each ragged per-unit spike train by its own timestamps.

Ragged per-unit spike times are not aligned to a common times axis, so each train is masked against itself: [t[in_epochs(t, epochs)] for t in trains].

Parameters:

Name	Type	Description	Default
trains	sequence of ndarray, or SpikeTrains	Per-unit 1-D spike-time arrays. A :class:~neurospatial.encoding.SpikeTrains container is accepted -- iterating it yields the per-unit train arrays.	required
epochs	IntervalSet-like, tuple, list, or ndarray	Epochs in any form accepted by :func:_as_intervals. Empty epochs -> every returned train is empty (but present, preserving order/count).	required
closed	('both', 'left', 'right', 'neither')	Endpoint inclusivity, forwarded to :func:in_epochs. Default "both".	"both"

Returns:

Type	Description
list of ndarray	One masked train per input train, in the same order. Always a plain list (type-stable), even when given a :class:~neurospatial.encoding.SpikeTrains -- SpikeTrains identity is preserved instead by Session.restrict.

Examples:

>>> import numpy as np
>>> from neurospatial.behavior import restrict_spike_trains
>>> trains = [np.array([0.1, 1.5, 2.9]), np.array([0.5, 3.0, 6.0])]
>>> restrict_spike_trains(trains, (1.0, 3.5))
[array([1.5, 2.9]), array([3.])]

Source code in src/neurospatial/behavior/epochs.py

def restrict_spike_trains(
    trains: Sequence[NDArray[np.float64]],
    epochs: Any,
    *,
    closed: _Closed = "both",
) -> list[NDArray[np.float64]]:
    """Restrict each ragged per-unit spike train by its **own** timestamps.

    Ragged per-unit spike times are not aligned to a common ``times`` axis, so
    each train is masked against itself: ``[t[in_epochs(t, epochs)] for t in
    trains]``.

    Parameters
    ----------
    trains : sequence of ndarray, or SpikeTrains
        Per-unit 1-D spike-time arrays. A
        :class:`~neurospatial.encoding.SpikeTrains` container is accepted --
        iterating it yields the per-unit train arrays.
    epochs : IntervalSet-like, tuple, list, or ndarray
        Epochs in any form accepted by :func:`_as_intervals`. Empty epochs ->
        every returned train is empty (but present, preserving order/count).
    closed : {"both", "left", "right", "neither"}, optional
        Endpoint inclusivity, forwarded to :func:`in_epochs`. Default
        ``"both"``.

    Returns
    -------
    list of ndarray
        One masked train per input train, in the same order. Always a plain
        ``list`` (type-stable), even when given a
        :class:`~neurospatial.encoding.SpikeTrains` -- ``SpikeTrains`` identity
        is preserved instead by ``Session.restrict``.

    Examples
    --------
    >>> import numpy as np
    >>> from neurospatial.behavior import restrict_spike_trains
    >>> trains = [np.array([0.1, 1.5, 2.9]), np.array([0.5, 3.0, 6.0])]
    >>> restrict_spike_trains(trains, (1.0, 3.5))
    [array([1.5, 2.9]), array([3.])]
    """
    # Normalize the intervals ONCE (not per train): a malformed epochs raises up
    # front, and each train is then masked against the same pre-normalized
    # ``(n, 2)`` array via ``_mask_in_intervals`` -- avoiding the per-unit
    # re-normalization that ``in_epochs`` would repeat on every call.
    intervals = _as_intervals(epochs)
    out: list[NDArray[np.float64]] = []
    for train in trains:
        t = np.asarray(train, dtype=np.float64)
        out.append(t[_mask_in_intervals(t, intervals, closed=closed)])
    return out