From cd99d6b874d9d09b3bb87b8485cc787885af71f1 Mon Sep 17 00:00:00 2001
From: YurenHao0426 <blackhao0426@gmail.com>
Date: Tue, 13 Jan 2026 23:49:05 -0600
Subject: init commit

---
 files/data_io/utils/visualize.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)
 create mode 100644 files/data_io/utils/visualize.py

(limited to 'files/data_io/utils/visualize.py')

diff --git a/files/data_io/utils/visualize.py b/files/data_io/utils/visualize.py
new file mode 100644
index 0000000..6f0de95
--- /dev/null
+++ b/files/data_io/utils/visualize.py
@@ -0,0 +1,19 @@
+import matplotlib.pyplot as plt
+import torch
+
+def plot_raster(spikes: torch.Tensor, title=None):
+    """
+    Plot raster diagram of spike activity (T,B,N) or (T,N).
+    """
+    s = spikes.detach().cpu()
+    if s.ndim == 3:
+        s = s[:, 0, :]  # take first batch
+    t, n = s.shape
+    for i in range(n):
+        times = torch.nonzero(s[:, i]).squeeze().numpy()
+        plt.scatter(times, i * np.ones_like(times), s=2, c='black')
+    plt.xlabel("Time step")
+    plt.ylabel("Neuron index")
+    if title:
+        plt.title(title)
+    plt.show()
-- 
cgit v1.2.3