Source code for brian2.synapses.spikequeue

The spike queue class stores future synaptic events
produced by a given presynaptic neuron group (or postsynaptic for backward
propagation in STDP).

import bisect

import numpy as np

from brian2.utils.arrays import calc_repeats
from brian2.utils.logger import get_logger

__all__ = ["SpikeQueue"]

logger = get_logger(__name__)


[docs]class SpikeQueue: """ Data structure saving the spikes and taking care of delays. Parameters ---------- source_start : int The start of the source indices (for subgroups) source_end : int The end of the source indices (for subgroups) Notes ----- **Data structure** A spike queue is implemented as a 2D array `X` that is circular in the time direction (rows) and dynamic in the events direction (columns). The row index corresponding to the current timestep is `currentime`. Each element contains the target synapse index. **Offsets** Offsets are used to solve the problem of inserting multiple synaptic events with the same delay. This is difficult to vectorise. If there are n synaptic events with the same delay, these events are given an offset between 0 and n-1, corresponding to their relative position in the data structure. """ def __init__(self, source_start, source_end): #: The start of the source indices (for subgroups) self._source_start = source_start #: The end of the source indices (for subgroups) self._source_end = source_end self.dtype = np.int32 # TODO: Ths is fixed for now self.X = np.zeros((1, 1), dtype=self.dtype) # target synapses self.X_flat = self.X.reshape( 1, ) #: The current time (in time steps) self.currenttime = 0 #: number of events in each time step self.n = np.zeros(1, dtype=int) #: The dt used for storing the spikes (will be set in `prepare`) self._dt = None self._state_tuple = (self._source_start, self._source_end, self.dtype)
[docs] def prepare(self, delays, dt, synapse_sources): """ Prepare the data structures This is called every time the network is run. The size of the of the data structure (number of rows) is adjusted to fit the maximum delay in `delays`, if necessary. A flag is set if delays are homogeneous, in which case insertion will use a faster method implemented in `insert_homogeneous`. """ if self._dt is not None: # store the current spikes spikes = self._extract_spikes() # adapt the spikes to the new dt if it changed if self._dt != dt: spiketimes = spikes[:, 0] * self._dt spikes[:, 0] = np.round(spiketimes / dt).astype(int) else: spikes = None if len(delays): delays = np.array(np.round(delays / dt)).astype(int) max_delays = max(delays) min_delays = min(delays) else: max_delays = min_delays = 0 self._delays = delays # Prepare the data structure used in propagation synapse_sources = synapse_sources[:] ss = np.ravel(synapse_sources) # mergesort to retain relative order, keeps the output lists in sorted order ss_sort_indices = np.argsort(ss, kind="mergesort") ss_sorted = ss[ss_sort_indices] splitinds = np.searchsorted( ss_sorted, np.arange(self._source_start, self._source_end + 1) ) self._neurons_to_synapses = [ ss_sort_indices[splitinds[j] : splitinds[j + 1]] for j in range(len(splitinds) - 1) ] max_events = max(map(len, self._neurons_to_synapses)) n_steps = max_delays + 1 # Adjust the maximum delay and number of events per timestep if necessary # Check if delays are homogeneous self._homogeneous = max_delays == min_delays # Resize if (n_steps > self.X.shape[0]) or (max_events > self.X.shape[1]): # Resize # Choose max_delay if is is larger than the maximum delay n_steps = max(n_steps, self.X.shape[0]) max_events = max(max_events, self.X.shape[1]) self.X = np.zeros( (n_steps, max_events), dtype=self.dtype ) # target synapses self.X_flat = self.X.reshape( n_steps * max_events, ) self.n = np.zeros(n_steps, dtype=int) # number of events in each time step # Re-insert the spikes into the data structure if spikes is not None: self._store_spikes(spikes) self._dt = dt
def _extract_spikes(self): """ Get all the stored spikes Returns ------- spikes : ndarray A 2d array with two columns, where each row describes a spike. The first column gives the time (as integer time steps) and the second column gives the index of the target synapse. """ spikes = np.zeros((np.sum(self.n), 2), dtype=int) counter = 0 for idx, n in enumerate(self.n): t = (idx - self.currenttime) % len(self.n) for target in self.X[idx, :n]: spikes[counter, :] = np.array([t, target]) counter += 1 return spikes def _store_spikes(self, spikes): """ Store a list of spikes at the given positions after clearing all spikes in the queue. Parameters ---------- spikes : ndarray A 2d array with two columns, where each row describes a spike. The first column gives the time (as integer time steps) and the second column gives the index of the target synapse. """ # Clear all spikes self.n[:] = 0 for t, target in spikes: row_idx = (t + self.currenttime) % len(self.n) self.X[row_idx, self.n[row_idx]] = target self.n[row_idx] += 1 def _full_state(self): return (self._dt, self._extract_spikes(), self.X.shape) def _restore_from_full_state(self, state): if state is None: # It is possible that _full_state was called in `SynapticPathway`, # before the `SpikeQueue` was created. In that case, delete all spikes in # the queue self._store_spikes(np.empty((0, 2), dtype=int)) self._dt = None else: self._dt, spikes, X_shape = state # Restore the previous shape n_steps, max_events = X_shape self.X = np.zeros((n_steps, max_events), dtype=self.dtype) self.X_flat = self.X.reshape( n_steps * max_events, ) self.n = np.zeros(n_steps, dtype=int) self._store_spikes(spikes) ################################ SPIKE QUEUE DATASTRUCTURE ################
[docs] def advance(self): """ Advances by one timestep """ self.n[self.currenttime] = 0 # erase self.currenttime = (self.currenttime + 1) % len(self.n)
[docs] def peek(self): """ Returns the all the synaptic events corresponding to the current time, as an array of synapse indexes. """ return self.X[self.currenttime, : self.n[self.currenttime]]
[docs] def push(self, sources): """ Push spikes to the queue. Parameters ---------- sources : ndarray of int The indices of the neurons that spiked. """ if len(sources) and len(self._delays): start = self._source_start stop = self._source_end if start > 0: start_idx = bisect.bisect_left(sources, start) else: start_idx = 0 if stop <= sources[-1]: stop_idx = bisect.bisect_left(sources, stop, lo=start_idx) else: stop_idx = len(sources) + 1 sources = sources[start_idx:stop_idx] if len(sources) == 0: return synapse_indices = self._neurons_to_synapses indices = np.concatenate( [synapse_indices[source - start] for source in sources] ).astype(np.int32) if self._homogeneous: # homogeneous delays self._insert_homogeneous(self._delays[0], indices) else: # vectorise over synaptic events self._insert(self._delays[indices], indices)
def _insert(self, delay, target): """ Vectorised insertion of spike events. Parameters ---------- delay : ndarray Delays in timesteps. target : ndarray Target synaptic indices. """ delay = np.array(delay, dtype=int) offset = calc_repeats(delay) # Calculate row indices in the data structure timesteps = (self.currenttime + delay) % len(self.n) # (Over)estimate the number of events to be stored, to resize the array # It's an overestimation for the current time, but I believe a good one # for future events m = max(self.n) + len(target) if m >= self.X.shape[1]: # overflow self._resize(m + 1) self.X_flat[timesteps * self.X.shape[1] + offset + self.n[timesteps]] = target self.n[timesteps] += offset + 1 # that's a trick (to update stack size) def _insert_homogeneous(self, delay, target): """ Inserts events at a fixed delay. Parameters ---------- delay : int Delay in timesteps. target : ndarray Target synaptic indices. """ timestep = (self.currenttime + delay) % len(self.n) nevents = len(target) m = ( self.n[timestep] + nevents + 1 ) # If overflow, then at least one self.n is bigger than the size if m >= self.X.shape[1]: self._resize(m + 1) # was m previously (not enough) k = timestep * self.X.shape[1] + self.n[timestep] self.X_flat[k : k + nevents] = target self.n[timestep] += nevents def _resize(self, maxevents): """ Resizes the underlying data structure (number of columns = spikes per dt). Parameters ---------- maxevents : int The new number of columns. It will be rounded to the closest power of 2. """ # old and new sizes old_maxevents = self.X.shape[1] new_maxevents = int(2 ** np.ceil(np.log2(maxevents))) # maybe 2 is too large # new array newX = np.zeros((self.X.shape[0], new_maxevents), dtype=self.X.dtype) newX[:, :old_maxevents] = self.X[:, :old_maxevents] # copy old data self.X = newX self.X_flat = self.X.reshape( self.X.shape[0] * new_maxevents, )