import scipy as sp
from collections import OrderedDict
import quantities as pq
from progress_indicator import ProgressIndicator
from . import SpykeException
[docs]def correlogram(trains, bin_size, max_lag=500 * pq.ms, border_correction=True,
per_second=True, unit=pq.ms, progress=None):
""" Return (cross-)correlograms from a dictionary of spike train
lists for different units.
:param dict trains: Dictionary of :class:`neo.core.SpikeTrain` lists.
:param bin_size: Bin size (time).
:type bin_size: Quantity scalar
:param max_lag: Cut off (end time of calculated correlogram).
:type max_lag: Quantity scalar
:param bool border_correction: Apply correction for less data at higher
timelags. Not perfect for bin_size != 1*``unit``, especially with
large ``max_lag`` compared to length of spike trains.
:param bool per_second: If ``True``, counts returned are per second.
Otherwise, counts per spike train are returned.
:param Quantity unit: Unit of X-Axis.
:param progress: A ProgressIndicator object for the operation.
:type progress: :class:`.progress_indicator.ProgressIndicator`
:returns: Two values:
* An ordered dictionary indexed with the indices of ``trains`` of
ordered dictionaries indexed with the same indices. Entries of
the inner dictionaries are the resulting (cross-)correlograms as
numpy arrays. All crosscorrelograms can be indexed in two
different ways: ``c[index1][index2]`` and ``c[index2][index1]``.
* The bins used for the correlogram calculation.
:rtype: dict, Quantity 1D
"""
if not progress:
progress = ProgressIndicator()
bin_size.rescale(unit)
max_lag.rescale(unit)
# Create bins, making sure that 0 is at the center of central bin
half_bins = sp.arange(bin_size / 2, max_lag, bin_size)
all_bins = list(reversed(-half_bins))
all_bins.extend(half_bins)
bins = sp.array(all_bins) * unit
middle_bin = len(bins) / 2 - 1
indices = trains.keys()
num_trains = len(trains[indices[0]])
if not num_trains:
raise SpykeException('Could not create correlogram: No spike trains!')
for u in range(1, len(indices)):
if len(trains[indices[u]]) != num_trains:
raise SpykeException('Could not create correlogram: All units ' +
'need the same number of spike trains!')
progress.set_ticks(sp.sum(range(len(trains) + 1) * num_trains))
corrector = 1
if border_correction:
# Need safe min/max functions
def safe_max(seq):
if len(seq) < 1:
return 0
return max(seq)
def safe_min(seq):
if len(seq) < 1:
return 2 ** 22 # Some arbitrary large value
return min(seq)
max_w = max([max([safe_max(t) for t in l])
for l in trains.itervalues()])
min_w = min([min([safe_min(t) for t in l])
for l in trains.itervalues()])
train_length = (max_w - min_w)
l = int(round(middle_bin)) + 1
cE = max(train_length - (l * bin_size) + 1 * unit, 1 * unit)
corrector = (train_length / sp.concatenate(
(sp.linspace(cE, train_length, l - 1, False),
sp.linspace(train_length, cE, l)))).magnitude
correlograms = OrderedDict()
for i1 in xrange(len(indices)): # For each index
# For all later indices, including itself
for i2 in xrange(i1, len(indices)):
histogram = sp.zeros(len(bins) - 1)
for t in xrange(num_trains):
train1 = trains[indices[i1]][t].rescale(unit).reshape((1, -1))
train2 = trains[indices[i2]][t].rescale(unit).reshape((-1, 1))
histogram += sp.histogram(
sp.subtract(train1, train2), bins=bins)[0]
if i1 == i2: # Correction for autocorrelogram
histogram[middle_bin] -= len(train2)
progress.step()
if per_second:
l = train1.t_stop - train1.t_start
if train2.t_stop - train2.t_start != l:
raise SpykeException(
'A spike train pair does not have equal length,'
'cannot calculate count per second.')
histogram /= l.rescale(pq.s)
crg = corrector * histogram / num_trains
if indices[i1] not in correlograms:
correlograms[indices[i1]] = OrderedDict()
correlograms[indices[i1]][indices[i2]] = crg
if i1 != i2:
if indices[i2] not in correlograms:
correlograms[indices[i2]] = OrderedDict()
correlograms[indices[i2]][indices[i1]] = crg[::-1]
return correlograms, bins