diff --git a/nirtorch/__init__.py b/nirtorch/__init__.py index 498c87d..7da59e9 100644 --- a/nirtorch/__init__.py +++ b/nirtorch/__init__.py @@ -12,6 +12,8 @@ from .to_nir import extract_nir_graph from .nir_interpreter import nir_to_torch from .torch_tracer import torch_to_nir +from .from_nir_data import from_nir_data +from .to_nir_data import to_nir_data __all__ = [ "extract_nir_graph", @@ -19,4 +21,6 @@ "load", "nir_to_torch", "torch_to_nir", + "from_nir_data", + "to_nir_data", ] diff --git a/nirtorch/from_nir_data.py b/nirtorch/from_nir_data.py new file mode 100644 index 0000000..3ef6415 --- /dev/null +++ b/nirtorch/from_nir_data.py @@ -0,0 +1,57 @@ +""" +Convert NIRGraphData to a torch-compatible spike representation. +""" + +import torch +from nir.data_ir import NIRGraphData, NIRNodeData + + +def from_nir_data( + nir_graph_data: NIRGraphData, + dt: float, + dynamic_before_transition: bool = True, + dimension_order: tuple = ("time", "batch", "neuron"), + time_unit: float = 1.0, +) -> dict: + """ + Args: + nir_graph_data: A NIRGraphData object containing the spike data. + dt: The time step size of the spike data. + 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. + dimension_order: The desired ordering of dimensions in the output spike + tensors. Defaults to ('time', 'batch', 'neuron'). + If the output spike tensors should have a different dimension + ordering, this can be used to specify the correct ordering. + time_unit: The unit of time for the spike data and dt. Defaults to 1.0, + which corresponds to seconds. For milliseconds, set this to 1e-3. + Returns: + A dictionary mapping node names to spike tensors. + """ + torch_dict = {} + + for node_key, nir_node_data in nir_graph_data.nodes.items(): + if isinstance(nir_node_data, NIRNodeData): + for observable, data in nir_node_data.observables.items(): + if observable == "spikes": + data = data.get_time_gridded( + dt=dt * time_unit, + dimension_order=dimension_order, + dynamic_before_transition=dynamic_before_transition, + ) + torch_spikes = torch.tensor(data.data).float() + else: + raise NotImplementedError( + "Only spikes are supported as observables yet." + ) + else: + raise NotImplementedError( + "The translation of nested NIRGraphData is not supported." + ) + + torch_dict[node_key] = torch_spikes + + return torch_dict diff --git a/nirtorch/to_nir_data.py b/nirtorch/to_nir_data.py new file mode 100644 index 0000000..ad07261 --- /dev/null +++ b/nirtorch/to_nir_data.py @@ -0,0 +1,50 @@ +""" +Convert spikes from a torch tensor to a NIRGraphData object. +""" + +from nir.data_ir import NIRGraphData, NIRNodeData, TimeGriddedData + + +def to_nir_data( + torch_dict: dict, + dt: float, + time_unit: float = 1.0, + dimension_order: tuple = ("time", "batch", "neuron"), + dynamic_before_transition: bool = True, +) -> NIRGraphData: + """ + Args: + torch_dict: A dictionary mapping node names to spike tensors. If the + shape is not (n_time_steps, batch_size, n_neurons), the `shape` + argument can be used to specify the correct ordering of dimensions. + dt: The time step size of the spike data. + time_unit: The unit of time for the spike data and dt. Defaults to 1.0, + which corresponds to seconds. For milliseconds, set this to 1e-3. + dimension_order: The order of dimensions in the input spike tensors. + Defaults to ("time", "batch", "neuron"). + dynamic_before_transition: 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. + Returns: + A NIRGraphData object containing the spike data. + """ + nir_nodes = {} + + for key, spikes in torch_dict.items(): + spikes = spikes.detach().cpu().numpy().astype(bool) + + nir_node_data = NIRNodeData( + { + "spikes": TimeGriddedData( + spikes, + dt * time_unit, + dimension_order=dimension_order, + dynamic_before_transition=dynamic_before_transition, + ) + } + ) + nir_nodes[key] = nir_node_data + + nir_data = NIRGraphData(nir_nodes) + return nir_data diff --git a/tests/test_nir_data.py b/tests/test_nir_data.py new file mode 100644 index 0000000..36b728f --- /dev/null +++ b/tests/test_nir_data.py @@ -0,0 +1,50 @@ +import nir +from nirtorch import from_nir_data, to_nir_data +import numpy as np +import torch + + +def test_event_data_from_nir(): + nir_data = nir.NIRGraphData( + nodes={ + "lif": nir.NIRNodeData( + observables={ + "spikes": nir.EventData( + idx=np.random.randint(0, 10, (3, 5)), + time=np.random.rand(3, 5) * 0.1, + n_neurons=10, + t_max=0.1, + ) + } + ) + } + ) + + torch_dict = from_nir_data(nir_data, dt=0.001) + + assert "lif" in torch_dict, "Node 'lif' not found in the converted dictionary" + assert torch_dict["lif"].shape == (100, 3, 10), ( + f"Expected shape (100, 3, 10) for node 'lif', but got {torch_dict['lif'].shape}" + ) + + +def test_stable_conversion(): + original_spikes = {"lif": torch.randint(0, 2, (4, 10, 10), dtype=torch.float32)} + + nir_data = to_nir_data(original_spikes, dt=0.001) + converted_spikes = from_nir_data(nir_data, dt=0.001) + + assert torch.equal(original_spikes["lif"], converted_spikes["lif"]), ( + f"Mismatch in spikes for node 'lif'. Got {converted_spikes['lif'].shape}, expected {original_spikes['lif'].shape}" + ) + + +def test_stable_conversion_with_time_unit(): + original_spikes = {"lif": torch.randint(0, 2, (4, 10, 10), dtype=torch.float32)} + + nir_data = to_nir_data(original_spikes, dt=1e-3, time_unit=1e-3) + converted_spikes = from_nir_data(nir_data, dt=1e-3, time_unit=1e-3) + + assert torch.equal(original_spikes["lif"], converted_spikes["lif"]), ( + f"Mismatch in spikes for node 'lif' with time unit conversion. Got {converted_spikes['lif'].shape}, expected {original_spikes['lif'].shape}" + )