diff --git a/nir/data_ir/graph.py b/nir/data_ir/graph.py index a527554..0c99876 100644 --- a/nir/data_ir/graph.py +++ b/nir/data_ir/graph.py @@ -1,12 +1,31 @@ from __future__ import annotations from dataclasses import dataclass from typing import Dict, Union +import warnings import numpy as np from nir.ir import NIRGraph, NIRNode @dataclass -class TimeGriddedData: +class ObservableData: + """ + Base class for observable data in SNNs. + """ + + def get_event(self, n_events: int | None) -> EventData: + pass + + def get_time_gridded( + self, + dt: float, + dimension_order: tuple = ("time", "batch", "neuron"), + dynamic_before_transition: bool = True, + ) -> TimeGriddedData: + pass + + +@dataclass +class TimeGriddedData(ObservableData): """ Either boolean entries indicate whether a binary event is present at a particular time step, or a real-valued signal provides the measurement @@ -18,23 +37,51 @@ class TimeGriddedData: Input data. For binary data the dtype should be bool. dt: float Time step size. + dimension_order: tuple, optional + The order of dimensions in the spike tensors. Defaults to ('time', + 'batch', 'neuron') for (time, batch, neurons). + dynamic_before_transition: bool, optional + If True, it is assumed that the framework evolves the state (e.g., + membrane potential) of the neurons before checking if the threshold has + been crossed and generating an event (transition). If False, the state + is updated after the event generation. """ data: np.ndarray dt: float # pylint: disable=invalid-name + dimension_order: tuple = ("time", "batch", "neuron") + dynamic_before_transition: bool = True def __post_init__(self): - if not isinstance(self.data, np.ndarray) or self.data.ndim != 3: - raise ValueError( - "Data must be of shape (n_samples, n_time_steps, n_neurons)" - "and of type np.ndarray" - ) - - def __getitem__(self, idx): - return self.data[idx] + if not isinstance(self.data, np.ndarray): + raise TypeError("Data must be a numpy array") + if self.data.ndim != 3: + raise ValueError("Data must be a 3D array") + if ( + "time" not in self.dimension_order + or "batch" not in self.dimension_order + or "neuron" not in self.dimension_order + ): + raise ValueError("dimension_order must contain 'time', 'batch', and 'neuron'") + + def _view_as(self, order): + if order == self.dimension_order: + return self.data + perm = tuple(self.dimension_order.index(dim) for dim in order) + return self.data.transpose(perm) + + def __getitem__(self, idx, out_order=("time", "batch", "neuron")): + """ + Get a slice of the data, with the option to specify the output + dimension order. The `idx` refers to the reordered data. + """ + return self._view_as(out_order)[idx] - def __setitem__(self, idx, val): - self.data[idx] = val + def __setitem__(self, idx, value, in_order=("time", "batch", "neuron")): + """ + Set a slice of the data, given as `value` in `in_order` dimension order. + """ + self._view_as(in_order)[idx] = value @property def shape(self): @@ -42,40 +89,64 @@ def shape(self): @property def n_samples(self): - return self.data.shape[0] + return self.shape[self.dimension_order.index("batch")] @property def n_time_steps(self): - return self.data.shape[1] + return self.shape[self.dimension_order.index("time")] @property def n_neurons(self): - return self.data.shape[2] + return self.shape[self.dimension_order.index("neuron")] @property def t_max(self): return self.n_time_steps * self.dt - def to_event(self, n_events: int, time_shift: float = 0.0) -> EventData: + def toggle_dynamic_before_transition(self): + """ + Toggle the dynamic_before_transition flag and update the data + accordingly. + """ + time_axis = self.dimension_order.index("time") + # if dynamic_before_transition is True, shift events back by one time step + if self.dynamic_before_transition: + self.data = np.roll(self.data, shift=-1, axis=time_axis) + # build an indexer equivalent to [:, -1, :] but for the correct axis + idx = [slice(None)] * self.data.ndim + idx[time_axis] = -1 + self.data[tuple(idx)] = False + else: + self.data = np.roll(self.data, shift=1, axis=time_axis) + # build an indexer equivalent to [:, 0, :] but for the correct axis + idx = [slice(None)] * self.data.ndim + idx[time_axis] = 0 + self.data[tuple(idx)] = False + + self.dynamic_before_transition = not self.dynamic_before_transition + + def get_event(self, n_events: int | None = None) -> EventData: """ + Convert the time-gridded data to event-based data, where each neuron + can have at most `n_events` events. If a neuron has more than + `n_events`, the earliest events are kept and the rest are dropped. + Arguments --------- n_spikes : int Maximum number of events stored for each neuron. - time_shift : float, optional - Shift the event times by this value from the beginning of each time - step. Must be in interval [0, dt). Default is 0.0. """ - + if n_events == None: + n_events = self.n_time_steps if not self.data.dtype == bool: raise ValueError("Data must be boolean to convert to EventData.") idx = np.full((self.n_samples, n_events), -1) time = np.full((self.n_samples, n_events), np.inf) - if time_shift < 0 or time_shift >= self.dt: - raise ValueError("time_shift must be in interval [0, dt)") for sample in range(self.n_samples): - time_step, neuron = np.where(self.data[sample]) + sample_idx = [slice(None)] * self.data.ndim + sample_idx[self.dimension_order.index("batch")] = sample + time_step, neuron = np.where(self.data[tuple(sample_idx)]) order = np.argsort(time_step) # sort events by time time_step = time_step[order] @@ -83,13 +154,48 @@ def to_event(self, n_events: int, time_shift: float = 0.0) -> EventData: num_events = min(len(time_step), n_events) idx[sample, :num_events] = neuron[:num_events] - time[sample, :num_events] = time_step[:num_events] * self.dt + time_shift + time[sample, :num_events] = ( + time_step[:num_events] + self.dynamic_before_transition + ) * self.dt return EventData(idx, time, self.n_neurons, self.t_max) + def get_time_gridded( + self, + dt: float, + dimension_order: tuple = ("time", "batch", "neuron"), + dynamic_before_transition: bool = True, + ) -> TimeGriddedData: + """ + Return a new TimeGriddedData object with the specified dt and + dynamic_before_transition flag. If the current object already has the + desired dt and dynamic_before_transition, return self. + """ + + if self.dt != dt: + event_data = self.get_event() + return event_data.get_time_gridded( + dt=dt, + dimension_order=dimension_order, + dynamic_before_transition=dynamic_before_transition + ) + + if self.dimension_order == dimension_order: + return self + else: + new_data = TimeGriddedData( + data=self._view_as(dimension_order), + dt=dt, + dimension_order=self.dimension_order, + dynamic_before_transition=self.dynamic_before_transition, + ) + if self.dynamic_before_transition != dynamic_before_transition: + new_data.toggle_dynamic_before_transition() + return new_data + @dataclass -class EventData: +class EventData(ObservableData): """ Event-based data represented as a list of event indices and their corresponding timestamps. Each event is discrete and carries no magnitude; @@ -124,27 +230,76 @@ def shape(self): def n_samples(self): return self.idx.shape[0] - def to_time_gridded( - self, dt: float # pylint: disable=invalid-name + def get_event(self, n_events: int | None) -> EventData: + """ + Return a new EventData object with at most `n_events` events per sample. + If a sample has more than `n_events`, the earliest events are kept and + the rest are dropped. + + Arguments + --------- + n_events : int + Maximum number of events stored for each sample. If None, return all events. + """ + if n_events is None or n_events >= self.idx.shape[1]: + return self + + new_idx = np.full((self.n_samples, n_events), -1) + new_time = np.full((self.n_samples, n_events), np.inf) + + for sample in range(self.n_samples): + valid_events = self.idx[sample] != -1 + valid_times = self.time[sample][valid_events] + valid_indices = self.idx[sample][valid_events] + + num_events = min(len(valid_times), n_events) + if num_events > 0: + order = np.argsort(valid_times) # sort events by time + new_idx[sample, :num_events] = valid_indices[order][:num_events] + new_time[sample, :num_events] = valid_times[order][:num_events] + + return EventData(new_idx, new_time, self.n_neurons, self.t_max) + + def get_time_gridded( + self, + dt: float, + dimension_order: tuple = ("time", "batch", "neuron"), + dynamic_before_transition: bool = True, # pylint: disable=invalid-name ) -> TimeGriddedData: """ Arguments --------- dt : float Time step size. + dynamic_before_transition : bool, optional + If True, the membrane potential is updated before checking if the + threshold has been crossed and generating an event (transition). If + False, the state is updated after the event generation. Default is + True. """ - n_time_steps = int(self.t_max / dt) - discrete_data = np.zeros( - (self.n_samples, n_time_steps, self.n_neurons), dtype=bool - ) - + n_time_steps = round(self.t_max / dt) + discrete_data = np.zeros((n_time_steps, self.n_samples, self.n_neurons), dtype=bool) for sample in range(self.n_samples): valid_spikes = self.idx[sample] != -1 valid_times = self.time[sample][valid_spikes] - steps = np.floor((valid_times / dt)).astype(int) + eps = dt * 1e-10 # small epsilon to avoid floating point issues + steps = np.ceil(valid_times / dt - eps).astype(int) - dynamic_before_transition neurons = self.idx[sample][valid_spikes] - discrete_data[sample, steps, neurons] = True - return TimeGriddedData(discrete_data, dt) + mask = steps < n_time_steps + if np.any(mask): + steps, neurons = steps[mask], neurons[mask] + warnings.warn( + "Some events got dropped because they occur after the " + "maximum time of the recording." + ) + + discrete_data[steps, sample, neurons] = True + if dimension_order != ("time", "batch", "neuron"): + perm = tuple(("time", "batch", "neuron").index(dim) for dim in dimension_order) + discrete_data = discrete_data.transpose(perm) + return TimeGriddedData( + data=discrete_data, dt=dt, dynamic_before_transition=dynamic_before_transition + ) @dataclass @@ -173,9 +328,44 @@ def __post_init__(self): if self.idx.shape != self.time.shape or self.idx.shape != self.value.shape: raise ValueError("idx, time and value must have the same shape") - def to_time_gridded( + def get_event(self, n_events: int | None) -> ValuedEventData: + """ + Return a new ValuedEventData object with at most `n_events` events per + sample. If a sample has more than `n_events`, the earliest events are + kept and the rest are dropped. + + Arguments + --------- + n_events : int + Maximum number of events stored for each sample. If None, return all events. + """ + if n_events is None or n_events >= self.idx.shape[1]: + return self + + new_idx = np.full((self.n_samples, n_events), -1) + new_time = np.full((self.n_samples, n_events), np.inf) + new_value = np.zeros((self.n_samples, n_events)) + + for sample in range(self.n_samples): + valid_events = self.idx[sample] != -1 + valid_times = self.time[sample][valid_events] + valid_indices = self.idx[sample][valid_events] + valid_values = self.value[sample][valid_events] + + num_events = min(len(valid_times), n_events) + if num_events > 0: + order = np.argsort(valid_times) # sort events by time + new_idx[sample, :num_events] = valid_indices[order][:num_events] + new_time[sample, :num_events] = valid_times[order][:num_events] + new_value[sample, :num_events] = valid_values[order][:num_events] + + return ValuedEventData(new_idx, new_time, self.n_neurons, self.t_max, new_value) + + def get_time_gridded( self, dt: float, # pylint: disable=invalid-name + dimension_order: tuple = ("time", "batch", "neuron"), + dynamic_before_transition: bool = True, ) -> TimeGriddedData: """ Currently, the values are assigned directly to the corresponding time @@ -185,6 +375,11 @@ def to_time_gridded( ---------- dt : float Time step size. + dynamic_before_transition : bool, optional + If True, the membrane potential is updated before checking if the + threshold has been crossed and generating an event (transition). If + False, the state is updated after the event generation. Default is + True. """ n_samples = self.n_samples n_time_steps = int(self.t_max / dt) @@ -193,12 +388,20 @@ def to_time_gridded( for sample in range(n_samples): valid_spikes = self.idx[sample] != -1 valid_times = self.time[sample][valid_spikes] - steps = np.floor((valid_times / dt)).astype(int) + if dynamic_before_transition: + steps = np.floor(valid_times / dt).astype(int) + else: + steps = np.ceil(valid_times / dt).astype(int) neurons = self.idx[sample][valid_spikes] value = self.value[sample][valid_spikes] discrete_data[sample, steps, neurons] = value - return TimeGriddedData(discrete_data, dt) + return TimeGriddedData( + data=discrete_data, + dt=dt, + dimension_order=dimension_order, + dynamic_before_transition=dynamic_before_transition, + ) @dataclass @@ -217,9 +420,7 @@ class NIRNodeData: def __post_init__(self): if not isinstance(self.observables, dict): - raise TypeError( - "observables must be a dictionary of EventData or TimeGriddedData" - ) + raise TypeError("observables must be a dictionary of EventData or TimeGriddedData") def __getitem__(self, idx): return self.observables[idx] @@ -278,6 +479,4 @@ def check_nodes(self, graph: NIRGraph): if not isinstance(graph_node, NIRNode): raise TypeError(f"Node {key} is not a NIRNode in the NIRGraph") if not node.check_observables(graph_node): - raise ValueError( - f"Observables for node {key} do not match the NIRNode" - ) + raise ValueError(f"Observables for node {key} do not match the NIRNode") diff --git a/tests/test_data_ir.py b/tests/test_data_ir.py index 8b012e6..6af5b6d 100644 --- a/tests/test_data_ir.py +++ b/tests/test_data_ir.py @@ -7,7 +7,7 @@ def test_generate_time_gridded_data(): dt = 0.1 gridded = nir.TimeGriddedData(spikes, dt) node = nir.NIRNodeData({"spikes": gridded}) - graph = nir.NIRGraphData({"node": node}) # noqa: F841 + graph = nir.NIRGraphData({"node": node}) assert np.allclose(graph.nodes["node"].observables["spikes"].data, spikes) assert graph.nodes["node"].observables["spikes"].dt == dt @@ -19,7 +19,7 @@ def test_generate_event_data(): t_max = 1.0 event = nir.EventData(idx, time, n_neurons, t_max) node = nir.NIRNodeData({"spikes": event}) - graph = nir.NIRGraphData({"node": node}) # noqa: F841 + graph = nir.NIRGraphData({"node": node}) assert np.allclose(graph.nodes["node"].observables["spikes"].idx, idx) assert np.allclose(graph.nodes["node"].observables["spikes"].time, time) assert graph.nodes["node"].observables["spikes"].n_neurons == n_neurons @@ -34,7 +34,7 @@ def test_generate_valued_event_data(): t_max = 1.0 valued_event = nir.ValuedEventData(idx, time, n_neurons, t_max, value) node = nir.NIRNodeData({"current": valued_event}) - graph = nir.NIRGraphData({"node": node}) # noqa: F841 + graph = nir.NIRGraphData({"node": node}) assert np.allclose(graph.nodes["node"].observables["current"].idx, idx) assert np.allclose(graph.nodes["node"].observables["current"].time, time) assert np.allclose(graph.nodes["node"].observables["current"].value, value) @@ -43,23 +43,18 @@ def test_generate_valued_event_data(): def test_binary_conversion(): - # time_shift = 0.0 * dt - spikes = np.random.randint(0, 2, size=(10, 10, 10)).astype(bool) - dt = 0.1 - gridded_1 = nir.TimeGriddedData(spikes, dt) - event = gridded_1.to_event(n_events=100) - gridded_2 = event.to_time_gridded(dt=dt) - assert np.array_equal(gridded_1.data, gridded_2.data) - assert gridded_1.dt == gridded_2.dt - - # time_shift = 0.5 * dt - spikes = np.random.randint(0, 2, size=(10, 10, 10)).astype(bool) - dt = 0.1 - gridded_1 = nir.TimeGriddedData(spikes, dt) - event = gridded_1.to_event(n_events=100, time_shift=0.5 * dt) - gridded_2 = event.to_time_gridded(dt=dt) - assert np.array_equal(gridded_1.data, gridded_2.data) - assert gridded_1.dt == gridded_2.dt + for dynamic_before_transition in [True, False]: + spikes = np.random.randint(0, 2, size=(5, 10, 10)).astype(bool) + dt = 0.1 + gridded_1 = nir.TimeGriddedData( + spikes, dt, dynamic_before_transition=dynamic_before_transition + ) + event = gridded_1.get_event(n_events=100) + gridded_2 = event.get_time_gridded(dt, dynamic_before_transition=dynamic_before_transition) + assert np.array_equal(gridded_1.data, gridded_2.data) + assert gridded_1.dt == gridded_2.dt + assert gridded_1.dynamic_before_transition == gridded_2.dynamic_before_transition + assert gridded_1.dimension_order == gridded_2.dimension_order def test_valued_conversion(): @@ -70,13 +65,16 @@ def test_valued_conversion(): dt = 0.01 t_max = 0.3 valued_event = nir.ValuedEventData(idx, time, n_neurons, t_max, value) - gridded = valued_event.to_time_gridded(dt=dt) # noqa: F841 + gridded = valued_event.get_time_gridded(dt=dt) expected = np.zeros((1, 30, 2)) expected[0, 5, 0] = 1 expected[0, 10, 1] = 2 expected[0, 15, 0] = 4 expected[0, 20, 1] = 3 assert np.array_equal(gridded.data, expected) + assert gridded.dt == dt + assert gridded.dimension_order == ("time", "batch", "neuron") # default value + assert gridded.dynamic_before_transition == True # default value def test_check_nodes():