""" Module defining `EventMonitor` and `SpikeMonitor`. """ import numpy as np from brian2.core.names import Nameable from brian2.core.spikesource import SpikeSource from brian2.core.variables import Variables from brian2.groups.group import CodeRunner, Group from brian2.units.fundamentalunits import Quantity __all__ = ["EventMonitor", "SpikeMonitor"] class EventMonitor(Group, CodeRunner): """ Record events from a `NeuronGroup` or another event source. The recorded events can be accessed in various ways: the attributes `~EventMonitor.i` and `~EventMonitor.t` store all the indices and event times, respectively. Alternatively, you can get a dictionary mapping neuron indices to event trains, by calling the `event_trains` method. Parameters ---------- source : `NeuronGroup`, `SpikeSource` The source of events to record. event : str The name of the event to record variables : str or sequence of str, optional Which variables to record at the time of the event (in addition to the index of the neuron). Can be the name of a variable or a list of names. record : bool, optional Whether or not to record each event in `i` and `t` (the `count` will always be recorded). Defaults to ``True``. when : str, optional When to record the events, by default records events in the same slot where the event is emitted. See :ref:`scheduling` for possible values. order : int, optional The priority of of this group for operations occurring at the same time step and in the same scheduling slot. Defaults to the order where the event is emitted + 1, i.e. it will be recorded directly afterwards. name : str, optional A unique name for the object, otherwise will use ``source.name+'_eventmonitor_0'``, etc. codeobj_class : class, optional The `CodeObject` class to run code with. See Also -------- SpikeMonitor """ invalidates_magic_network = False add_to_magic_network = True def __init__( self, source, event, variables=None, record=True, when=None, order=None, name="eventmonitor*", codeobj_class=None, ): if not isinstance(source, SpikeSource): raise TypeError( f"{self.__class__.__name__} can only monitor groups " "producing spikes (such as NeuronGroup), but the given " f"argument is of type {type(source)}." ) #: The source we are recording from self.source = source #: Whether to record times and indices of events self.record = record #: The array of event counts (length = size of target group) self.count = None del self.count # this is handled by the Variable mechanism if event not in source.events: if event == "spike": threshold_text = " Did you forget to set a 'threshold'?" else: threshold_text = "" raise ValueError( f"Recorded group '{source.name}' does not define an event " f"'{event}'.{threshold_text}" ) if when is None: if order is not None: raise ValueError("Cannot specify order if when is not specified.") # TODO: Would be nicer if there was a common way of accessing the # relevant object for NeuronGroup and SpikeGeneratorGroup if hasattr(source, "thresholder"): parent_obj = source.thresholder[event] else: parent_obj = source when = parent_obj.when order = parent_obj.order + 1 elif order is None: order = 0 #: The event that we are listening to self.event = event if variables is None: variables = {} elif isinstance(variables, str): variables = {variables} #: The additional variables that will be recorded self.record_variables = set(variables) for variable in variables: if variable not in source.variables: raise ValueError( f"'{variable}' is not a variable of the recorded group" ) if self.record: self.record_variables |= {"i", "t"} # Some dummy code so that code generation takes care of the indexing # and subexpressions code = [f"_to_record_{v} = _source_{v}" for v in sorted(self.record_variables)] code = "\n".join(code) self.codeobj_class = codeobj_class # Since this now works for general events not only spikes, we have to # pass the information about which variable to use to the template, # it can not longer simply refer to "_spikespace" eventspace_name = f"_{event}space" # Handle subgroups correctly start = getattr(source, "start", 0) stop = getattr(source, "stop", len(source)) source_N = getattr(source, "_source_N", len(source)) Nameable.__init__(self, name=name) self.variables = Variables(self) self.variables.add_reference(eventspace_name, source) for variable in self.record_variables: source_var = source.variables[variable] self.variables.add_reference(f"_source_{variable}", source, variable) self.variables.add_auxiliary_variable( f"_to_record_{variable}", dimensions=source_var.dim, dtype=source_var.dtype, ) self.variables.add_dynamic_array( variable, size=0, dimensions=source_var.dim, dtype=source_var.dtype, read_only=True, ) self.variables.add_arange("_source_idx", size=len(source)) self.variables.add_array( "count", size=len(source), dtype=np.int32, read_only=True, index="_source_idx", ) self.variables.add_constant("_source_start", start) self.variables.add_constant("_source_stop", stop) self.variables.add_constant("_source_N", source_N) self.variables.add_array( "N", size=1, dtype=np.int32, read_only=True, scalar=True ) record_variables = { varname: self.variables[varname] for varname in self.record_variables } template_kwds = { "eventspace_variable": source.variables[eventspace_name], "record_variables": record_variables, "record": self.record, } needed_variables = {eventspace_name} | self.record_variables CodeRunner.__init__( self, group=self, code=code, template="spikemonitor", name=None, # The name has already been initialized clock=source.clock, when=when, order=order, needed_variables=needed_variables, template_kwds=template_kwds, ) self.variables.create_clock_variables(self._clock, prefix="_clock_") self.add_dependency(source) self.written_readonly_vars = { self.variables[varname] for varname in self.record_variables } self._enable_group_attributes() def resize(self, new_size): # Note that this does not set N, this has to be done in the template # since we use a restricted pointer to access it (which promises that # we only change the value through this pointer) for variable in self.record_variables: self.variables[variable].resize(new_size) def reinit(self): """ Clears all recorded spikes """ raise NotImplementedError() @property def it(self): """ Returns the pair (`i`, `t`). """ if not self.record: raise AttributeError( "Indices and times have not been recorded." "Set the record argument to True to record " "them." ) return self.i, self.t @property def it_(self): """ Returns the pair (`i`, `t_`). """ if not self.record: raise AttributeError( "Indices and times have not been recorded." "Set the record argument to True to record " "them." ) return self.i, self.t_ def _values_dict(self, first_pos, sort_indices, used_indices, var): sorted_values = self.state(var, use_units=False)[sort_indices] dim = self.variables[var].dim event_values = {} current_pos = 0 # position in the all_indices array for idx in range(len(self.source)): if current_pos < len(used_indices) and used_indices[current_pos] == idx: if current_pos < len(used_indices) - 1: event_values[idx] = Quantity( sorted_values[ first_pos[current_pos] : first_pos[current_pos + 1] ], dim=dim, ) else: event_values[idx] = Quantity( sorted_values[first_pos[current_pos] :], dim=dim ) current_pos += 1 else: event_values[idx] = Quantity([], dim=dim) return event_values def values(self, var): """ Return a dictionary mapping neuron indices to arrays of variable values at the time of the events (sorted by time). Parameters ---------- var : str The name of the variable. Returns ------- values : dict Dictionary mapping each neuron index to an array of variable values at the time of the events Examples -------- >>> from brian2 import * >>> G = NeuronGroup(2, '''counter1 : integer ... counter2 : integer ... max_value : integer''', ... threshold='counter1 >= max_value', ... reset='counter1 = 0') >>> G.run_regularly('counter1 += 1; counter2 += 1') # doctest: +ELLIPSIS CodeRunner(...) >>> G.max_value = [50, 100] >>> mon = EventMonitor(G, event='spike', variables='counter2') >>> run(10*ms) >>> counter2_values = mon.values('counter2') >>> print(counter2_values[0]) [ 50 100] >>> print(counter2_values[1]) [100] """ if not self.record: raise AttributeError( "Indices and times have not been recorded." "Set the record argument to True to record " "them." ) indices = self.i[:] # We have to make sure that the sort is stable, otherwise our spike # times do not necessarily remain sorted. sort_indices = np.argsort(indices, kind="mergesort") used_indices, first_pos = np.unique(self.i[:][sort_indices], return_index=True) return self._values_dict(first_pos, sort_indices, used_indices, var) def all_values(self): """ Return a dictionary mapping recorded variable names (including ``t``) to a dictionary mapping neuron indices to arrays of variable values at the time of the events (sorted by time). This is equivalent to (but more efficient than) calling `values` for each variable and storing the result in a dictionary. Returns ------- all_values : dict Dictionary mapping variable names to dictionaries which themselves are mapping neuron indicies to arrays of variable values at the time of the events. Examples -------- >>> from brian2 import * >>> G = NeuronGroup(2, '''counter1 : integer ... counter2 : integer ... max_value : integer''', ... threshold='counter1 >= max_value', ... reset='counter1 = 0') >>> G.run_regularly('counter1 += 1; counter2 += 1') # doctest: +ELLIPSIS CodeRunner(...) >>> G.max_value = [50, 100] >>> mon = EventMonitor(G, event='spike', variables='counter2') >>> run(10*ms) >>> all_values = mon.all_values() >>> print(all_values['counter2'][0]) [ 50 100] >>> print(all_values['t'][1]) [ 9.9] ms """ if not self.record: raise AttributeError( "Indices and times have not been recorded." "Set the record argument to True to record " "them." ) indices = self.i[:] sort_indices = np.argsort(indices, kind="mergesort") used_indices, first_pos = np.unique(self.i[:][sort_indices], return_index=True) all_values_dict = {} for varname in self.record_variables - {"i"}: all_values_dict[varname] = self._values_dict( first_pos, sort_indices, used_indices, varname ) return all_values_dict def event_trains(self): """ Return a dictionary mapping neuron indices to arrays of event times. Equivalent to calling ``values('t')``. Returns ------- event_trains : dict Dictionary that stores an array with the event times for each neuron index. See Also -------- SpikeMonitor.spike_trains """ return self.values("t") @property def num_events(self): """ Returns the total number of recorded events. """ return self.N[:] def __repr__(self): classname = self.__class__.__name__ return f"<{classname}, recording event '{self.event}' from '{self.group.name}'>" class SpikeMonitor(EventMonitor): """ Record spikes from a `NeuronGroup` or other spike source. The recorded spikes can be accessed in various ways (see Examples below): the attributes `~SpikeMonitor.i` and `~SpikeMonitor.t` store all the indices and spike times, respectively. Alternatively, you can get a dictionary mapping neuron indices to spike trains, by calling the `spike_trains` method. If you record additional variables with the ``variables`` argument, these variables can be accessed by their name (see Examples). Parameters ---------- source : (`NeuronGroup`, `SpikeSource`) The source of spikes to record. variables : str or sequence of str, optional Which variables to record at the time of the spike (in addition to the index of the neuron). Can be the name of a variable or a list of names. record : bool, optional Whether or not to record each spike in `i` and `t` (the `count` will always be recorded). Defaults to ``True``. when : str, optional When to record the events, by default records events in the same slot where the event is emitted. See :ref:`scheduling` for possible values. order : int, optional The priority of of this group for operations occurring at the same time step and in the same scheduling slot. Defaults to the order where the event is emitted + 1, i.e. it will be recorded directly afterwards. name : str, optional A unique name for the object, otherwise will use ``source.name+'_spikemonitor_0'``, etc. codeobj_class : class, optional The `CodeObject` class to run code with. Examples -------- >>> from brian2 import * >>> spikes = SpikeGeneratorGroup(3, [0, 1, 2], [0, 1, 2]*ms) >>> spike_mon = SpikeMonitor(spikes) >>> net = Network(spikes, spike_mon) >>> net.run(3*ms) >>> print(spike_mon.i[:]) [0 1 2] >>> print(spike_mon.t[:]) [ 0. 1. 2.] ms >>> print(spike_mon.t_[:]) [ 0. 0.001 0.002] >>> from brian2 import * >>> G = NeuronGroup(2, '''counter1 : integer ... counter2 : integer ... max_value : integer''', ... threshold='counter1 >= max_value', ... reset='counter1 = 0') >>> G.run_regularly('counter1 += 1; counter2 += 1') # doctest: +ELLIPSIS CodeRunner(...) >>> G.max_value = [50, 100] >>> mon = SpikeMonitor(G, variables='counter2') >>> net = Network(G, mon) >>> net.run(10*ms) >>> print(mon.i[:]) [0 0 1] >>> print(mon.counter2[:]) [ 50 100 100] """ def __init__( self, source, variables=None, record=True, when=None, order=None, name="spikemonitor*", codeobj_class=None, ): #: The array of spike counts (length = size of target group) self.count = None del self.count # this is handled by the Variable mechanism super().__init__( source, event="spike", variables=variables, record=record, when=when, order=order, name=name, codeobj_class=codeobj_class, ) @property def num_spikes(self): """ Returns the total number of recorded spikes. """ return self.num_events # We "re-implement" the following functions only to get more specific # doc strings (and to make sure that the methods are included in the # reference documentation for SpikeMonitor). def spike_trains(self): """ Return a dictionary mapping neuron indices to arrays of spike times. Returns ------- spike_trains : dict Dictionary that stores an array with the spike times for each neuron index. Examples -------- >>> from brian2 import * >>> spikes = SpikeGeneratorGroup(3, [0, 1, 2], [0, 1, 2]*ms) >>> spike_mon = SpikeMonitor(spikes) >>> run(3*ms) >>> spike_trains = spike_mon.spike_trains() >>> spike_trains[1] array([ 1.]) * msecond """ return self.event_trains() def values(self, var): """ Return a dictionary mapping neuron indices to arrays of variable values at the time of the spikes (sorted by time). Parameters ---------- var : str The name of the variable. Returns ------- values : dict Dictionary mapping each neuron index to an array of variable values at the time of the spikes. Examples -------- >>> from brian2 import * >>> G = NeuronGroup(2, '''counter1 : integer ... counter2 : integer ... max_value : integer''', ... threshold='counter1 >= max_value', ... reset='counter1 = 0') >>> G.run_regularly('counter1 += 1; counter2 += 1') # doctest: +ELLIPSIS CodeRunner(...) >>> G.max_value = [50, 100] >>> mon = SpikeMonitor(G, variables='counter2') >>> run(10*ms) >>> counter2_values = mon.values('counter2') >>> print(counter2_values[0]) [ 50 100] >>> print(counter2_values[1]) [100] """ return super().values(var) def all_values(self): """ Return a dictionary mapping recorded variable names (including ``t``) to a dictionary mapping neuron indices to arrays of variable values at the time of the spikes (sorted by time). This is equivalent to (but more efficient than) calling `values` for each variable and storing the result in a dictionary. Returns ------- all_values : dict Dictionary mapping variable names to dictionaries which themselves are mapping neuron indicies to arrays of variable values at the time of the spikes. Examples -------- >>> from brian2 import * >>> G = NeuronGroup(2, '''counter1 : integer ... counter2 : integer ... max_value : integer''', ... threshold='counter1 >= max_value', ... reset='counter1 = 0') >>> G.run_regularly('counter1 += 1; counter2 += 1') # doctest: +ELLIPSIS CodeRunner(...) >>> G.max_value = [50, 100] >>> mon = SpikeMonitor(G, variables='counter2') >>> run(10*ms) >>> all_values = mon.all_values() >>> print(all_values['counter2'][0]) [ 50 100] >>> print(all_values['t'][1]) [ 9.9] ms """ return super().all_values() def __repr__(self): classname = self.__class__.__name__ return f"<{classname}, recording from '{self.group.name}'>"