From 26e8d8172645a4ca07533d133b1957b5552bf176 Mon Sep 17 00:00:00 2001 From: Ben Kroehs Date: Mon, 11 May 2026 17:11:34 +0200 Subject: [PATCH 1/2] feat: Add NIRData conversion utilities --- nirtorch/__init__.py | 4 +++ nirtorch/from_nir_data.py | 67 +++++++++++++++++++++++++++++++++++++++ nirtorch/to_nir_data.py | 44 +++++++++++++++++++++++++ tests/test_nir_data.py | 51 +++++++++++++++++++++++++++++ 4 files changed, 166 insertions(+) create mode 100644 nirtorch/from_nir_data.py create mode 100644 nirtorch/to_nir_data.py create mode 100644 tests/test_nir_data.py 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..3bd389b --- /dev/null +++ b/nirtorch/from_nir_data.py @@ -0,0 +1,67 @@ +""" +Convert NIRGraphData to a torch-compatible spike representation. +""" + +import torch +from nir.data_ir import EventData, NIRGraphData, NIRNodeData + + +def from_nir_data( + nir_graph_data: NIRGraphData, + dt: float, + dynamic_before_transition: bool = True, + shape: tuple = ("T", "B", "N"), + 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. + shape: The desired shape of the output spike tensors. Defaults to + ('T', 'B', 'N') for (time, batch, neurons). + If the output spike tensors should have a different shape, this can + be used to specify the correct ordering of dimensions. + 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. + """ + nirdata_shape = ("B", "T", "N") + 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": + if isinstance(data, EventData): + data = data.to_time_gridded( + dt=dt * time_unit, + dynamic_before_transition=dynamic_before_transition + ) + else: + if data.dynamic_before_transition != dynamic_before_transition: + data = data.toggle_dynamic_before_transition() + torch_spikes = torch.tensor(data.data) + # reorder dimensions to shape specified by `shape` argument + torch_spikes = torch_spikes.permute( + nirdata_shape.index(shape[0]), + nirdata_shape.index(shape[1]), + nirdata_shape.index(shape[2]), + ) + 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.float() + + return torch_dict diff --git a/nirtorch/to_nir_data.py b/nirtorch/to_nir_data.py new file mode 100644 index 0000000..cfdee82 --- /dev/null +++ b/nirtorch/to_nir_data.py @@ -0,0 +1,44 @@ +""" +Convert spikes from a torch tensor to a NIRGraphData object. +""" + +import torch +from nir.data_ir import NIRGraphData, NIRNodeData, TimeGriddedData + + +def to_nir_data( + torch_dict: dict, + dt: float, + time_unit: float = 1.0, + shape: tuple = ("T", "B", "N"), +) -> 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. + shape: The shape of the input spike tensors. Defaults to + ('T', 'B', 'N') for (time, batch, neurons). If the spike tensors + have a different shape, this can be used to specify the correct + ordering of dimensions. + Returns: + A NIRGraphData object containing the spike data. + """ + nir_nodes = {} + + for key, spikes in torch_dict.items(): + # reorder dimensions to (batch, time, neurons) + spikes = torch.permute( + spikes, (shape.index("B"), shape.index("T"), shape.index("N")) + ) + + spikes = spikes.detach().cpu().numpy().astype(bool) + + nir_node_data = NIRNodeData({"spikes": TimeGriddedData(spikes, dt * time_unit)}) + 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..42e4d37 --- /dev/null +++ b/tests/test_nir_data.py @@ -0,0 +1,51 @@ +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"]), ( + "Mismatch in spikes for node 'lif'" + ) + + +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"]), ( + "Mismatch in spikes for node 'lif' with time unit conversion" + ) From b4afc5b3bb28dac52f80c9396c59e05f99defa68 Mon Sep 17 00:00:00 2001 From: Ben Kroehs Date: Thu, 13 Aug 2026 13:10:24 +0200 Subject: [PATCH 2/2] chore: Adapt to new NIR API --- nirtorch/from_nir_data.py | 34 ++++++++++++---------------------- nirtorch/to_nir_data.py | 30 ++++++++++++++++++------------ tests/test_nir_data.py | 5 ++--- 3 files changed, 32 insertions(+), 37 deletions(-) diff --git a/nirtorch/from_nir_data.py b/nirtorch/from_nir_data.py index 3bd389b..3ef6415 100644 --- a/nirtorch/from_nir_data.py +++ b/nirtorch/from_nir_data.py @@ -3,14 +3,14 @@ """ import torch -from nir.data_ir import EventData, NIRGraphData, NIRNodeData +from nir.data_ir import NIRGraphData, NIRNodeData def from_nir_data( nir_graph_data: NIRGraphData, dt: float, dynamic_before_transition: bool = True, - shape: tuple = ("T", "B", "N"), + dimension_order: tuple = ("time", "batch", "neuron"), time_unit: float = 1.0, ) -> dict: """ @@ -22,37 +22,27 @@ def from_nir_data( 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. - shape: The desired shape of the output spike tensors. Defaults to - ('T', 'B', 'N') for (time, batch, neurons). - If the output spike tensors should have a different shape, this can - be used to specify the correct ordering of dimensions. + 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. """ - nirdata_shape = ("B", "T", "N") 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": - if isinstance(data, EventData): - data = data.to_time_gridded( - dt=dt * time_unit, - dynamic_before_transition=dynamic_before_transition - ) - else: - if data.dynamic_before_transition != dynamic_before_transition: - data = data.toggle_dynamic_before_transition() - torch_spikes = torch.tensor(data.data) - # reorder dimensions to shape specified by `shape` argument - torch_spikes = torch_spikes.permute( - nirdata_shape.index(shape[0]), - nirdata_shape.index(shape[1]), - nirdata_shape.index(shape[2]), + 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." @@ -62,6 +52,6 @@ def from_nir_data( "The translation of nested NIRGraphData is not supported." ) - torch_dict[node_key] = torch_spikes.float() + torch_dict[node_key] = torch_spikes return torch_dict diff --git a/nirtorch/to_nir_data.py b/nirtorch/to_nir_data.py index cfdee82..ad07261 100644 --- a/nirtorch/to_nir_data.py +++ b/nirtorch/to_nir_data.py @@ -2,7 +2,6 @@ Convert spikes from a torch tensor to a NIRGraphData object. """ -import torch from nir.data_ir import NIRGraphData, NIRNodeData, TimeGriddedData @@ -10,7 +9,8 @@ def to_nir_data( torch_dict: dict, dt: float, time_unit: float = 1.0, - shape: tuple = ("T", "B", "N"), + dimension_order: tuple = ("time", "batch", "neuron"), + dynamic_before_transition: bool = True, ) -> NIRGraphData: """ Args: @@ -20,24 +20,30 @@ def to_nir_data( 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. - shape: The shape of the input spike tensors. Defaults to - ('T', 'B', 'N') for (time, batch, neurons). If the spike tensors - have a different shape, this can be used to specify the correct - ordering of dimensions. + 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(): - # reorder dimensions to (batch, time, neurons) - spikes = torch.permute( - spikes, (shape.index("B"), shape.index("T"), shape.index("N")) - ) - spikes = spikes.detach().cpu().numpy().astype(bool) - nir_node_data = NIRNodeData({"spikes": TimeGriddedData(spikes, dt * time_unit)}) + 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) diff --git a/tests/test_nir_data.py b/tests/test_nir_data.py index 42e4d37..36b728f 100644 --- a/tests/test_nir_data.py +++ b/tests/test_nir_data.py @@ -4,7 +4,6 @@ import torch - def test_event_data_from_nir(): nir_data = nir.NIRGraphData( nodes={ @@ -36,7 +35,7 @@ def test_stable_conversion(): converted_spikes = from_nir_data(nir_data, dt=0.001) assert torch.equal(original_spikes["lif"], converted_spikes["lif"]), ( - "Mismatch in spikes for node 'lif'" + f"Mismatch in spikes for node 'lif'. Got {converted_spikes['lif'].shape}, expected {original_spikes['lif'].shape}" ) @@ -47,5 +46,5 @@ def test_stable_conversion_with_time_unit(): converted_spikes = from_nir_data(nir_data, dt=1e-3, time_unit=1e-3) assert torch.equal(original_spikes["lif"], converted_spikes["lif"]), ( - "Mismatch in spikes for node 'lif' with time unit conversion" + f"Mismatch in spikes for node 'lif' with time unit conversion. Got {converted_spikes['lif'].shape}, expected {original_spikes['lif'].shape}" )