feat: Add NIRData conversion functionalities

nirtorch/__init__.py

@@ -12,11 +12,15 @@
 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",
     "extract_torch_graph",
     "load",
     "nir_to_torch",
     "torch_to_nir",
+    "from_nir_data",
+    "to_nir_data",
 ]

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

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

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}"
+    )