SpikeTrains

class braintools.spike.SpikeTrains(times, counts, t_start, t_stop, overflow=False)

Spike times of a population in a padded, fixed-shape layout.

Neuron \(n\) has counts[n] spikes, stored in ascending order in times[n, :counts[n]]; the remaining capacity - counts[n] entries of the row are padding with value +inf. Because every array has a static shape, a SpikeTrains can be created, transformed and consumed inside jax.jit and jax.vmap. It is a JAX pytree whose leaves are all five fields.

Parameters:

• times (Any) – Spike times with shape (n_neurons, capacity), sorted ascending per neuron and padded with +inf. A brainunit Quantity when the time unit is known, otherwise a plain array.

• counts (Any) – Number of valid entries of every row, shape (n_neurons,); stored as int32.

• t_start (Any) – Start of the recording, in the unit of times.

• t_stop (Any) – End of the recording, in the unit of times.

• overflow (Any) – Whether some neuron had more than capacity spikes, so that spikes were dropped (only possible for traced input, see below).

Variables:

• times (jax.Array or Quantity) – Padded spike times, shape (n_neurons, capacity).

• counts (jax.Array) – Valid entries per neuron, shape (n_neurons,), int32.

• t_stop (t_start,) – Scalars in the unit of times. Spike times lie in [t_start, t_stop].

• overflow (jax.Array) – Scalar bool.

See also

braintools.visualize.spike_raster

Plot a SpikeTrains as a raster.

Notes

Capacity and overflow. The capacity (the padded row length K) is part of the array shapes, so it must be known when a function is traced.

• With concrete (non-traced) input and capacity=None, the capacity is computed as max(1, largest spike count).

• With traced input (inside jit, vmap or grad), capacity=None raises ValueError; pass an upper bound instead.

• With concrete input, a neuron with more than capacity spikes raises ValueError.

• With traced input, the first capacity spikes of every neuron are kept, counts is clipped to capacity and overflow is set to True. Call check() on the result outside the transform to turn the flag into an error.

Units. When dt (for from_dense()) or the input times carry brainunit units, times, t_start and t_stop are Quantities in that unit, and time arguments of the methods (bin_size, t0, …) must be Quantities with compatible units. Without units, plain numbers are used throughout.

Transforms. All methods work under jit and vmap except the host-only ones: from_list(), to_events() and to_list() (their results are ragged) and, for traced data, the default t_start/t_stop of from_events(). Methods whose output shape depends on t_start/t_stop (to_dense(), bin()) need those values to be concrete, or an explicit n_steps/n_bins. Methods assume an unbatched object; inside vmap every trial is unbatched.

Examples

>>> import brainunit as u
>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> spikes = jnp.array([[1, 0],
...                     [0, 0],
...                     [1, 1],
...                     [1, 0]])
>>> st = SpikeTrains.from_dense(spikes, dt=1. * u.ms)
>>> st.n_neurons, st.capacity
(2, 3)
>>> st.counts
Array([3, 1], dtype=int32)
>>> st.times
Quantity([[ 0.  2.  3.]
          [ 2. inf inf]], "ms")
>>> st.bin(2. * u.ms)
Array([[1, 0],
       [2, 1]], dtype=int32)
>>>
>>> bool((st.to_dense(dt=1. * u.ms) == spikes.astype(bool)).all())
True

bin(bin_size, n_bins=None)

Count spikes in consecutive bins starting at t_start.

Bin b covers [t_start + b * bin_size, t_start + (b + 1) * bin_size). When bin_size does not divide the duration, the last bin is partial. A spike exactly at t_stop is counted in the last bin when that bin reaches t_stop; spikes past the last bin are ignored.

Parameters:

• bin_size (Any) – Bin width, compatible with the spike times.

• n_bins (int | None) – Number of bins. Defaults to ceil(duration / bin_size) (at least 1 when there are spikes, so a zero-length recording keeps them), which needs concrete t_start/t_stop/bin_size; pass it explicitly under jit/vmap. Spikes beyond the last bin are ignored.

Returns:

Spike counts of shape (n_bins, N), int32.

Return type:

Array

Raises:

ValueError – If bin_size is not positive, or n_bins is missing for traced bounds.

Examples

>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_dense(jnp.ones((5, 1)), dt=1.)
>>> st.bin(2.)[:, 0]
Array([2, 2, 1], dtype=int32)

property capacity: int

Padded number of spike slots per neuron K (static).

Type:

int

check()

Validate the container and raise if spikes were dropped.

In eager mode this raises on overflow and checks the layout invariants (counts in [0, K], valid times sorted and within [t_start, t_stop], padding +inf). On a traced object it cannot raise and returns self unchanged; call it on the output of the transform instead.

Returns:

self, to allow chaining.

Return type:

SpikeTrains

Raises:

ValueError – If overflow is set or an invariant is violated.

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> make = jax.jit(lambda s: SpikeTrains.from_dense(s, dt=1., capacity=2))
>>> st = make(jnp.ones((3, 1)))
>>> bool(st.overflow)
True
>>> try:
...     st.check()
... except ValueError:
...     print('overflow')
overflow

property duration: Any

t_stop - t_start.

Type:

float or Quantity

classmethod from_dense(spikes, dt=None, t_start=0.0, capacity=None)

Build from a dense spike raster with time on axis 0.

Every nonzero entry spikes[i, n] is one spike of neuron n at time t_start + i * dt; t_stop = t_start + T * dt.

Parameters:

• spikes (Any) – Raster with shape (T, N) (bool or numeric).

• dt (Any) – Time step. Defaults to brainstate.environ.get_dt(). A Quantity makes the result carry its unit.

• t_start (Any) – Time of row 0. Must be a Quantity when dt is one, unless it is zero.

• capacity (int | None) – Padded number of spike slots per neuron. Required when spikes is traced; see the class notes for the overflow policy.

Returns:

The spike trains. Works under jit and vmap when capacity is given.

Return type:

SpikeTrains

Raises:

• ValueError – If spikes is not 2-D, if capacity is None for traced input, or if a neuron has more than capacity spikes (concrete input).

• TypeError – If t_start and dt disagree about units.

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_dense(jnp.array([[0, 1], [1, 1]]), dt=0.5)
>>> st.times
Array([[0.5, inf],
       [0. , 0.5]], dtype=float32)
>>>
>>> trials = jnp.zeros((8, 100, 3)).at[:, ::10, 0].set(1.)
>>> batched = jax.vmap(lambda s: SpikeTrains.from_dense(s, dt=0.1, capacity=10))(trials)
>>> batched.times.shape
(8, 3, 10)

classmethod from_events(times, neuron_ids, n_neurons, t_start=None, t_stop=None, capacity=None)

Build from flat event arrays (times, neuron_ids).

Events may come in any order; they are sorted per neuron.

Parameters:

• times (Any) – Spike times, shape (M,).

• neuron_ids (Any) – Neuron index of every spike, shape (M,), in [0, n_neurons).

• n_neurons (int) – Number of neurons N (static).

• t_start (Any) – Recording bounds, in units compatible with times. Default to the minimum and maximum of times (0 for no events), which requires concrete times.

• t_stop (Any) – Recording bounds, in units compatible with times. Default to the minimum and maximum of times (0 for no events), which requires concrete times.

• capacity (int | None) – Padded number of spike slots per neuron. Required when any input is traced; see the class notes for the overflow policy.

Returns:

The spike trains. Works under jit and vmap when capacity, t_start and t_stop are given.

Return type:

SpikeTrains

Raises:

ValueError – If the shapes differ; if capacity, t_start or t_stop is missing for traced input; or, for concrete input, if a neuron id is out of range, a time lies outside [t_start, t_stop], or a neuron has more than capacity spikes.

Notes

For traced input, events with an out-of-range neuron id or a time outside [t_start, t_stop] are dropped instead of raising.

Examples

>>> import brainunit as u
>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_events(jnp.array([3., 1., 2.]) * u.ms,
...                              jnp.array([0, 0, 1]), n_neurons=2)
>>> st.times
Quantity([[ 1.  3.]
          [ 2. inf]], "ms")
>>> st.t_start, st.t_stop
(Quantity(1., "ms"), Quantity(3., "ms"))

classmethod from_list(trains, t_start=None, t_stop=None, capacity=None)

Build from a ragged list with one spike-time array per neuron (host-only).

Parameters:

• trains (Sequence[Any]) – trains[n] holds the spike times of neuron n; empty trains are allowed. Either all non-empty trains are Quantities with compatible units, or none is.

• t_start (Any) – Recording bounds; default to the earliest and latest spike.

• t_stop (Any) – Recording bounds; default to the earliest and latest spike.

• capacity (int | None) – Padded number of spike slots per neuron; defaults to the largest train length (at least 1).

Returns:

The spike trains.

Return type:

SpikeTrains

Raises:

• TypeError – If Quantity and plain trains are mixed.

• ValueError – As for from_events().

Examples

>>> import numpy as np
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_list([np.array([0.5, 0.1]), np.array([]), [0.3]],
...                            t_start=0., t_stop=1.)
>>> st.counts
Array([2, 0, 1], dtype=int32)
>>> [t.tolist() for t in st.to_list()]
[[0.10000000149011612, 0.5], [], [0.30000001192092896]]

property mask: Array

(N, K) bool, True for valid (non-padding) entries.

Type:

jax.Array

property n_neurons: int

Number of neurons N (static).

Type:

int

select(idx)

Return the subset of neurons idx.

Parameters:

idx (Any) – Neurons to keep, in the given order. An integer keeps a single neuron (the result still has a neuron axis). Boolean masks need concrete values.

Returns:

The selected neurons, same capacity and bounds. overflow is carried over (it is not tracked per neuron).

Return type:

SpikeTrains

Examples

>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_dense(jnp.eye(3), dt=1.)
>>> st.select([2, 0]).times
Array([[2.],
       [0.]], dtype=float32)

to_dense(dt=None, n_steps=None)

Convert to a dense boolean raster with time on axis 0.

Spike t goes to row round((t - t_start) / dt); spikes that share a row merge, and spikes on rows >= n_steps are dropped.

Parameters:

• dt (Any) – Time step, compatible with the spike times. Defaults to brainstate.environ.get_dt().

• n_steps (int | None) – Number of rows T. Defaults to round((t_stop - t_start) / dt), plus one row if a spike lies exactly at t_stop (as with the default t_stop of from_events()), so that no spike is dropped. The default needs concrete values; pass it explicitly under jit/vmap.

Returns:

Boolean raster of shape (T, N).

Return type:

Array

Examples

>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_events(jnp.array([0., 2.]), jnp.array([1, 0]),
...                              n_neurons=2, t_start=0., t_stop=3.)
>>> st.to_dense(dt=1.).astype(int)
Array([[0, 1],
       [0, 0],
       [1, 0]], dtype=int32)

to_events()

Return flat (times, neuron_ids) sorted by time (host-only).

Ties in time are ordered by neuron index.

Return type:

Tuple[Any, ndarray]

Returns:

• times (numpy.ndarray or Quantity) – Spike times, shape (M,); a Quantity when the trains carry units.

• neuron_ids (numpy.ndarray) – Neuron index of every spike, shape (M,).

Raises:

TypeError – If called on a traced SpikeTrains (inside jit/vmap).

Examples

>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_dense(jnp.array([[0, 1], [1, 1]]), dt=1.)
>>> times, ids = st.to_events()
>>> times.tolist(), ids.tolist()
([0.0, 1.0, 1.0], [1, 0, 1])

to_list()

Return one spike-time array per neuron (host-only).

Returns:

N arrays of lengths counts; Quantities when the trains carry units.

Return type:

List[Any]

Raises:

TypeError – If called on a traced SpikeTrains (inside jit/vmap).

Examples

>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_dense(jnp.array([[0, 1], [1, 1]]), dt=1.)
>>> [t.tolist() for t in st.to_list()]
[[1.0], [0.0, 1.0]]

tree_flatten()

Flatten into the five leaves (times, counts, t_start, t_stop, overflow).

classmethod tree_unflatten(aux_data, children)

Rebuild from tree_flatten() output without validating the leaves.

window(t0, t1)

Restrict to the spikes in [t0, t1), keeping the capacity.

The bounds are clipped to the recording, so the result has t_start = max(t0, t_start) and t_stop = min(t1, t_stop) (never below the new t_start). A spike exactly at t0 is kept; a spike exactly at t1 is dropped, except that a spike at t_stop is kept when t1 extends past t_stop. overflow is carried over.

Parameters:

• t0 (Any) – Window bounds, compatible with the spike times.

• t1 (Any) – Window bounds, compatible with the spike times.

Returns:

The windowed trains, with the same (N, K) shape.

Return type:

SpikeTrains

Examples

>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> st = SpikeTrains.from_dense(jnp.ones((6, 1)), dt=1.)
>>> w = st.window(2., 4.)
>>> w.times, w.counts
(Array([[ 2.,  3., inf, inf, inf, inf]], dtype=float32), Array([2], dtype=int32))