braintools.analysis.spike.cross_correlogram

braintools.analysis.spike.cross_correlogram#

braintools.analysis.spike.cross_correlogram(a, b, bin_size, max_lag)#

Compute the spike-time cross-correlogram of paired neurons.

For every neuron \(n\), counts the pairs of spikes \((t^a_k, t^b_{k'})\) of neuron \(n\) of a and neuron \(n\) of b by their lag \(\tau = t^b_{k'} - t^a_k\), in n_bins = 2 * ceil(max_lag / bin_size) + 1 bins centred on the lags \(j \cdot \text{bin\_size}\), \(j = -M, \dots, M\). A positive lag means the spike of b follows the spike of a.

Parameters:

  • a (SpikeTrains) – Spike trains with the same number of neurons N; neuron n of a is paired with neuron n of b (row-wise pairs). The times of b are converted to the unit of a.

  • b (SpikeTrains) – Spike trains with the same number of neurons N; neuron n of a is paired with neuron n of b (row-wise pairs). The times of b are converted to the unit of a.

  • bin_size (Any) – Width of the lag bins, static and positive. A plain number is in milliseconds.

  • max_lag (Any) – Largest lag of interest, static and non-negative. A plain number is in milliseconds.

Return type:

Tuple[Array, Any]

Returns:

  • counts (jax.Array) – int32 pair counts of shape (N, n_bins).

  • lags (jax.Array or Quantity) – Bin centres (-M, ..., M) * bin_size, shape (n_bins,); a Quantity in the unit of a when it carries units, otherwise plain milliseconds.

Raises:

  • TypeError – If a or b is not a SpikeTrains, or bin_size/max_lag are traced.

  • ValueError – If the neuron counts differ, bin_size is not positive or max_lag is negative.

See also

auto_correlogram

The correlogram of a train with itself.

psth

Event-aligned histograms with the same counting.

braintools.spike.SpikeTrains.select

Pick neurons to pair up.

Notes

Bin \(j\) collects the lags whose nearest bin centre is \(j b\). A lag exactly half-way between two centres, \(\pm (j + \tfrac12) b\), goes to the bin farther from lag 0 on both sides, so the correlogram of a train with itself is exactly mirror-symmetric. Bin \(j > 0\) is \([(j - \tfrac12) b, (j + \tfrac12) b)\), bin \(-j\) is \((-(j + \tfrac12) b, -(j - \tfrac12) b]\), and the centre bin is the open interval \((-\tfrac12 b, \tfrac12 b)\). Together the bins cover a little more than [-max_lag, max_lag] [1]. Lags within a few machine epsilons of a half-way point count as lying on it, so grid-aligned spike times with float round-off are binned consistently.

elephant.spike_train_correlation.cross_correlation_histogram works on binned trains and counts the difference of bin indices. Both agree when the spike times lie on the bin_size grid; for off-grid spikes this function bins the exact time differences instead, so it does not depend on where the bin grid starts.

For every spike of a, the spikes of b before each bin edge are counted with searchsorted, in O(N K_a n_bins log K_b) time. The neurons are processed with jax.lax.map in batches of a static size chosen so that a batch’s temporaries hold at most about \(2^{21}\) elements (at least one neuron, K_a (n_bins + 1) elements), so peak memory is O(N (K_a + K_b) + max(2^{21}, K_a n_bins)) instead of O(N K_a n_bins). Under an outer jax.vmap the per-batch memory is multiplied by the size of that vmap. The function works under jax.jit and jax.vmap with static bin_size and max_lag; the counts are not differentiable. For all pairs of neurons, vmap over index pairs with select() (see the example).

References

Examples

>>> import brainunit as u
>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.spike import SpikeTrains
>>> from braintools.analysis.spike import cross_correlogram
>>> raster = jnp.zeros((20, 3)).at[jnp.array([2, 10]), 0].set(1)
>>> raster = raster.at[jnp.array([3, 11]), 1].set(1).at[5, 2].set(1)
>>> st = SpikeTrains.from_dense(raster, dt=1. * u.ms)
>>> counts, lags = cross_correlogram(st.select([0]), st.select([1]),
...                                  bin_size=1. * u.ms, max_lag=2. * u.ms)
>>> counts
Array([[0, 0, 0, 2, 0]], dtype=int32)
>>> lags
Quantity([-2. -1.  0.  1.  2.], "ms")
>>>
>>> # all pairs i < j, by vmapping over index pairs
>>> i, j = jnp.triu_indices(3, k=1)
>>> pairs = jax.vmap(lambda i, j: cross_correlogram(
...     st.select(i), st.select(j), 1. * u.ms, 2. * u.ms)[0][0])(i, j)
>>> pairs.shape
(3, 5)