Integrate raster plot

src/lava/utils/plots.py

@@ -0,0 +1,45 @@
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: BSD-3-Clause
+# See: https://spdx.org/licenses/
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def raster_plot(spks, stride=6, fig=None, color='b', alpha=1):
+    """Generate raster plot of spiking activity.
+
+    Parameters
+    ----------
+    spks : np.ndarray shape (num_neurons, num_timesteps)
+        Spiking activity of neurons, a spike is indicated by a one
+    stride : int
+        Stride for plotting neurons
+    """
+    num_neurons = spks.shape[0]
+    num_time_steps = spks.shape[1]
+
+    if stride >= num_neurons:
+        raise ValueError("Stride must be less than the number of neurons")
+
+    time_steps = np.arange(0, num_time_steps, 1)
+    if fig is None:
+        fig = plt.figure(figsize=(10, 5))
+
+    plt.xlim(-1, num_time_steps)
+    plt.yticks([])
+
+    plt.xlabel('Time steps')
+    plt.ylabel('Neurons')
+
+    for i in range(0, num_neurons, stride):
+        spike_times = time_steps[spks[i] == 1]
+        plt.plot(spike_times,
+                 i * np.ones(spike_times.shape),
+                 linestyle=' ',
+                 marker='o',
+                 markersize=1.5,
+                 color=color,
+                 alpha=alpha)
+
+    return fig

tutorials/end_to_end/tutorial02_excitatory_inhibitory_network.ipynb

@@ -5,7 +5,7 @@
    "id": "3fd51524",
    "metadata": {},
    "source": [
-    "*Copyright (C) 2022 Intel Corporation*<br>\n",
+    "*Copyright (C) 2022-23 Intel Corporation*<br>\n",
     "*SPDX-License-Identifier: BSD-3-Clause*<br>\n",
     "*See: https://spdx.org/licenses/*\n",
     "\n",
@@ -889,51 +889,6 @@
     "To this end, we display neurons on the vertical axis and mark the time step when a neuron spiked."
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "id": "abac04ec",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def raster_plot(spks, stride=6, fig=None, color='b', alpha=1):\n",
-    "    \"\"\"Generate raster plot of spiking activity.\n",
-    "    \n",
-    "    Parameters\n",
-    "    ----------\n",
-    "    \n",
-    "    spks : np.ndarray shape (num_neurons, timesteps)\n",
-    "        Spiking activity of neurons, a spike is indicated by a one    \n",
-    "    stride : int\n",
-    "        Stride for plotting neurons\n",
-    "    \"\"\"\n",
-    "    num_time_steps = spks.shape[1]\n",
-    "    assert stride < num_time_steps, \"Stride must be smaller than number of time steps\"\n",
-    "    \n",
-    "    time_steps = np.arange(0, num_time_steps, 1)\n",
-    "    if fig is None:\n",
-    "        fig = plt.figure(figsize=(10,5))\n",
-    "    timesteps = spks.shape[1]\n",
-    "    \n",
-    "    plt.xlim(-1, num_time_steps)\n",
-    "    plt.yticks([])\n",
-    "    \n",
-    "    plt.xlabel('Time steps')\n",
-    "    plt.ylabel('Neurons')\n",
-    "    \n",
-    "    for i in range(0, dim, stride):\n",
-    "        spike_times = time_steps[spks[i] == 1]\n",
-    "        plt.plot(spike_times,\n",
-    "                 i * np.ones(spike_times.shape),\n",
-    "                 linestyle=' ',\n",
-    "                 marker='o',\n",
-    "                 markersize=1.5,\n",
-    "                 color=color,\n",
-    "                 alpha=alpha)\n",
-    "        \n",
-    "    return fig       "
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 18,
@@ -954,6 +909,7 @@
     }
    ],
    "source": [
+    "from lava.utils.plots import raster_plot\n",
     "fig = raster_plot(spks=spks_balanced)"
    ]
   },