vis_utils.py
Shared visualisation helpers - 486 lines.
View on GitHub (vis_utils.py).
Source
Section titled “Source”"""Shared visualisation helpers that mirror the CIFAR-10 (non-DVS) plot style
in `image/cifar10_conv_dyned_dynedc_snn.py`.
Goal: speech-side plots produced by `speech/regenerate_v4_figures.py` look
visually consistent with the CIFAR-10 ones (same panel layouts, axes,
colormaps, threshold overlays).
Each helper writes:
1. The PNG (matching CIFAR's layout).
2. A companion `<name>_data.npz` (and CSV where appropriate) holding every
array needed to redraw the plot from disk alone.
The CIFAR script itself is not modified - these helpers reproduce its
look. They are used by the speech regen script and could be reused later
to back the CIFAR script's plotting if the user wants.
"""
from __future__ import annotations
import csv
import math
import os
from pathlib import Path
from typing import Dict, Iterable, Optional, Sequence
import matplotlib
try:
matplotlib.use("Agg")
except Exception:
pass
import matplotlib.pyplot as plt
import numpy as np
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _to_numpy(x):
if hasattr(x, "detach"):
x = x.detach()
if hasattr(x, "cpu"):
x = x.cpu()
if hasattr(x, "numpy"):
x = x.numpy()
return np.asarray(x)
def _ensure_dir(path: str | os.PathLike) -> str:
p = str(path)
os.makedirs(p, exist_ok=True)
return p
def dump_plot_data(out_dir: str | os.PathLike, name: str, **arrays) -> str:
"""Save raw plot data alongside a PNG so the figure can be regenerated.
Designed to be called next to a `plt.savefig(...)` line inside training
scripts: pass the same arrays that were plotted and they're written to
`<out_dir>/<name>_data.npz`. Tensors are auto-converted to numpy; `None`
values are skipped.
Use a separate suffix (e.g. `name="foo"` -> `foo_data.npz`) so the data
file never collides with the PNG/CSV next to it.
"""
out_dir = _ensure_dir(out_dir)
payload = {}
for k, v in arrays.items():
if v is None:
continue
payload[k] = _to_numpy(v)
path = os.path.join(out_dir, f"{name}_data.npz")
np.savez(path, **payload)
return path
# ---------------------------------------------------------------------------
# 1. Network activity (2x2) - mirrors CIFAR's visualize_network_activity
# Panels: [output spike raster | output membrane traces |
# mean firing rate hist | side imshow (conv filters / encoded input)]
# ---------------------------------------------------------------------------
def plot_network_activity(
out_dir: str | os.PathLike,
*,
output_spike_raster: np.ndarray, # [T, num_outputs] binary
output_membrane: np.ndarray, # [T, num_outputs] continuous
mean_firing_rates: np.ndarray, # 1D, per-neuron mean rate
side_panel_image: np.ndarray, # 2D image for bottom-right (filters / features)
side_panel_title: str,
side_panel_xlabel: str = "",
side_panel_cmap: str = "gray",
filename: str = "network_activity",
) -> None:
out_dir = _ensure_dir(out_dir)
spk = np.asarray(output_spike_raster)
mem = np.asarray(output_membrane)
rates = np.asarray(mean_firing_rates).flatten()
side = np.asarray(side_panel_image)
fig, axes = plt.subplots(2, 2, figsize=(15, 15))
axes[0, 0].imshow(spk, aspect="auto", cmap="binary")
axes[0, 0].set_title("Spike Raster (First Sample)")
axes[0, 0].set_xlabel("Neuron Index")
axes[0, 0].set_ylabel("Time Step")
axes[0, 1].plot(mem)
axes[0, 1].set_title("Membrane Potential (First Sample)")
axes[0, 1].set_xlabel("Time Step")
axes[0, 1].set_ylabel("Membrane Potential")
axes[1, 0].hist(rates, bins=50)
axes[1, 0].set_title("Average Firing Rate Distribution")
axes[1, 0].set_xlabel("Firing Rate")
axes[1, 0].set_ylabel("Count")
axes[1, 1].imshow(side, cmap=side_panel_cmap, aspect="auto")
axes[1, 1].set_title(side_panel_title)
if side_panel_xlabel:
axes[1, 1].set_xlabel(side_panel_xlabel)
plt.tight_layout()
plt.savefig(os.path.join(out_dir, f"{filename}.png"), dpi=150)
plt.close()
np.savez(
os.path.join(out_dir, f"{filename}_data.npz"),
output_spike_raster=spk,
output_membrane=mem,
mean_firing_rates=rates,
side_panel_image=side,
side_panel_title=np.array(side_panel_title),
side_panel_xlabel=np.array(side_panel_xlabel),
side_panel_cmap=np.array(side_panel_cmap),
)
def replot_network_activity(out_dir: str | os.PathLike, filename: str = "network_activity") -> None:
d = np.load(os.path.join(out_dir, f"{filename}_data.npz"), allow_pickle=False)
plot_network_activity(
out_dir,
output_spike_raster=d["output_spike_raster"],
output_membrane=d["output_membrane"],
mean_firing_rates=d["mean_firing_rates"],
side_panel_image=d["side_panel_image"],
side_panel_title=str(d["side_panel_title"]),
side_panel_xlabel=str(d["side_panel_xlabel"]),
side_panel_cmap=str(d["side_panel_cmap"]),
filename=filename,
)
# ---------------------------------------------------------------------------
# 2. Per-layer spike rasters (2x2) - mirrors visualize_layer_spike_rasters
# ---------------------------------------------------------------------------
def plot_layer_spike_rasters(
out_dir: str | os.PathLike,
*,
layer_arrays: Dict[str, np.ndarray], # keys: spk1, spk2, spk3, output
layer_titles: Dict[str, str],
suptitle: str = "Per-Layer Spike Rasters (Sample 0)",
filename: str = "layer_spike_rasters",
) -> None:
"""All four panels use cmap='binary' with a per-panel rate annotation.
`layer_arrays['output']` should already be binarised for speech (mem_out > threshold)
or a real spike record for CIFAR (spk_out).
"""
out_dir = _ensure_dir(out_dir)
keys = ["spk1", "spk2", "spk3", "output"]
arrs = {k: np.asarray(layer_arrays[k]) for k in keys}
fig, axes = plt.subplots(2, 2, figsize=(16, 12))
for ax, key in zip(axes.flat, keys):
d = arrs[key]
ax.imshow(d, aspect="auto", cmap="binary", interpolation="nearest")
ax.set_title(layer_titles.get(key, key))
ax.set_xlabel("Neuron Index")
ax.set_ylabel("Time Step")
rate = float(d.mean())
ax.text(
0.02,
0.98,
f"Rate: {rate:.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
color="red",
bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
)
plt.suptitle(suptitle, fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(out_dir, f"{filename}.png"), dpi=150)
plt.close()
np.savez(
os.path.join(out_dir, f"{filename}_data.npz"),
keys=np.array(keys),
titles=np.array([layer_titles.get(k, k) for k in keys]),
suptitle=np.array(suptitle),
**{k: arrs[k] for k in keys},
)
def replot_layer_spike_rasters(out_dir: str | os.PathLike, filename: str = "layer_spike_rasters") -> None:
d = np.load(os.path.join(out_dir, f"{filename}_data.npz"), allow_pickle=False)
keys = [str(k) for k in d["keys"]]
plot_layer_spike_rasters(
out_dir,
layer_arrays={k: d[k] for k in keys},
layer_titles=dict(zip(keys, [str(t) for t in d["titles"]])),
suptitle=str(d["suptitle"]),
filename=filename,
)
# ---------------------------------------------------------------------------
# 3. Membrane distributions (2x2) - mirrors visualize_membrane_distributions
# ---------------------------------------------------------------------------
def plot_membrane_distributions(
out_dir: str | os.PathLike,
*,
mem_arrays: Dict[str, np.ndarray], # keys: mem1, mem2, mem3, mem_out
layer_titles: Dict[str, str],
threshold: Optional[float],
filename: str = "membrane_distributions",
) -> None:
out_dir = _ensure_dir(out_dir)
keys = ["mem1", "mem2", "mem3", "mem_out"]
flats = {k: np.asarray(mem_arrays[k]).flatten() for k in keys}
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
for ax, key in zip(axes.flat, keys):
vals = flats[key]
ax.hist(vals, bins=100, density=True, alpha=0.7, color="steelblue")
ax.set_title(f"{layer_titles.get(key, key)} Membrane Potential Distribution")
ax.set_xlabel("Membrane Potential")
ax.set_ylabel("Density")
if threshold is not None:
ax.axvline(x=threshold, color="red", linestyle="--", label="Threshold")
ax.legend()
ax.text(
0.02,
0.98,
f"Mean: {vals.mean():.3f}\nStd: {vals.std():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
)
plt.suptitle("Membrane Potential Distributions", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(out_dir, f"{filename}.png"), dpi=150)
plt.close()
np.savez(
os.path.join(out_dir, f"{filename}_data.npz"),
keys=np.array(keys),
titles=np.array([layer_titles.get(k, k) for k in keys]),
threshold=np.array(threshold if threshold is not None else np.nan),
**{k: flats[k] for k in keys},
)
def replot_membrane_distributions(out_dir: str | os.PathLike, filename: str = "membrane_distributions") -> None:
d = np.load(os.path.join(out_dir, f"{filename}_data.npz"), allow_pickle=False)
keys = [str(k) for k in d["keys"]]
threshold = float(d["threshold"])
plot_membrane_distributions(
out_dir,
mem_arrays={k: d[k] for k in keys},
layer_titles=dict(zip(keys, [str(t) for t in d["titles"]])),
threshold=None if math.isnan(threshold) else threshold,
filename=filename,
)
# ---------------------------------------------------------------------------
# 4. Per-class spike patterns - mirrors visualize_per_class_spikes
# ---------------------------------------------------------------------------
def plot_per_class_spikes(
out_dir: str | os.PathLike,
*,
per_class_arr: np.ndarray, # [num_classes, T, num_outputs]
class_labels: Sequence[str],
grid: tuple[int, int] | None = None, # (rows, cols); auto if None
filename: str = "per_class_spikes",
) -> None:
out_dir = _ensure_dir(out_dir)
arr = np.asarray(per_class_arr)
n = arr.shape[0]
if grid is None:
if n <= 10:
rows, cols = 2, math.ceil(n / 2)
else:
cols = 7
rows = math.ceil(n / cols)
else:
rows, cols = grid
fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 3 * rows))
axes_flat = np.atleast_1d(axes).flatten()
summary_means = []
for cls_idx in range(len(axes_flat)):
ax = axes_flat[cls_idx]
if cls_idx >= n:
ax.axis("off")
summary_means.append(np.nan)
continue
m = arr[cls_idx]
if not np.isfinite(m).any():
ax.set_title(f"{class_labels[cls_idx]} (no samples)", fontsize=8)
summary_means.append(np.nan)
continue
ax.imshow(m, aspect="auto", cmap="binary", interpolation="nearest")
ax.set_title(class_labels[cls_idx], fontsize=8)
ax.set_xlabel("Neuron")
ax.set_ylabel("Time Step")
ax.tick_params(labelsize=6)
summary_means.append(float(m.mean()))
plt.suptitle("Per-Class Average Spike Patterns (Output Layer)", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(out_dir, f"{filename}.png"), dpi=150)
plt.close()
np.savez(
os.path.join(out_dir, f"{filename}_data.npz"),
per_class_arr=arr,
class_labels=np.array(list(class_labels)),
grid=np.array([rows, cols]),
)
csv_path = os.path.join(out_dir, f"{filename}_summary.csv")
with open(csv_path, "w", newline="") as f:
w = csv.writer(f)
w.writerow(["class_index", "class_label", "mean_output_activity"])
for i, label in enumerate(class_labels):
w.writerow([i, label, "" if i >= len(summary_means) else summary_means[i]])
def replot_per_class_spikes(out_dir: str | os.PathLike, filename: str = "per_class_spikes") -> None:
d = np.load(os.path.join(out_dir, f"{filename}_data.npz"), allow_pickle=False)
grid = d["grid"]
plot_per_class_spikes(
out_dir,
per_class_arr=d["per_class_arr"],
class_labels=[str(c) for c in d["class_labels"]],
grid=(int(grid[0]), int(grid[1])),
filename=filename,
)
# ---------------------------------------------------------------------------
# 5. Weight distributions - mirrors visualize_weight_distributions
# ---------------------------------------------------------------------------
def plot_weight_distributions(
out_dir: str | os.PathLike,
*,
weights: Dict[str, np.ndarray],
initial_weights: Optional[Dict[str, np.ndarray]] = None,
suptitle: str = "Weight Distributions: Before vs After Training",
filename: str = "weight_distributions",
) -> None:
out_dir = _ensure_dir(out_dir)
names = list(weights.keys())
flats = {n: np.asarray(weights[n]).flatten() for n in names}
init_flats = (
{n: np.asarray(initial_weights[n]).flatten() for n in names if n in initial_weights} if initial_weights else {}
)
n = len(names)
fig, axes = plt.subplots(1, n, figsize=(6 * n, 5))
if n == 1:
axes = [axes]
for ax, name in zip(axes, names):
vals = flats[name]
if name in init_flats:
ax.hist(init_flats[name], bins=100, alpha=0.5, label="Before", density=True, color="blue")
ax.hist(vals, bins=100, alpha=0.5, label="After", density=True, color="red")
ax.legend()
else:
ax.hist(vals, bins=100, density=True, color="steelblue", alpha=0.85)
ax.set_title(f"{name} Weight Distribution")
ax.set_xlabel("Weight Value")
ax.set_ylabel("Density")
plt.suptitle(suptitle, fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(out_dir, f"{filename}.png"), dpi=150)
plt.close()
payload = {
"names": np.array(names),
"suptitle": np.array(suptitle),
}
for name in names:
payload[f"weight__{name}"] = np.asarray(weights[name])
if name in init_flats:
payload[f"initial__{name}"] = np.asarray(initial_weights[name])
np.savez(os.path.join(out_dir, f"{filename}_data.npz"), **payload)
def replot_weight_distributions(out_dir: str | os.PathLike, filename: str = "weight_distributions") -> None:
d = np.load(os.path.join(out_dir, f"{filename}_data.npz"), allow_pickle=False)
names = [str(n) for n in d["names"]]
weights = {n: d[f"weight__{n}"] for n in names}
init = {}
for n in names:
key = f"initial__{n}"
if key in d.files:
init[n] = d[key]
plot_weight_distributions(
out_dir,
weights=weights,
initial_weights=init or None,
suptitle=str(d["suptitle"]),
filename=filename,
)
# ---------------------------------------------------------------------------
# 6. Firing-rates history - mirrors visualize_firing_rates_history
# ---------------------------------------------------------------------------
def plot_firing_rates_history(
out_dir: str | os.PathLike,
*,
metrics_csv_path: str,
rate_columns: Optional[Iterable[str]] = None,
title: str = "Per-Layer Firing Rates Over Training",
filename: str = "firing_rates_history",
) -> None:
if not os.path.exists(metrics_csv_path):
return
out_dir = _ensure_dir(out_dir)
with open(metrics_csv_path, "r") as f:
reader = csv.reader(f)
header = next(reader)
rows = list(reader)
if not rows:
return
cols = {h: i for i, h in enumerate(header)}
epoch_col = "epoch" if "epoch" in cols else header[0]
epochs = [int(float(r[cols[epoch_col]])) for r in rows]
if rate_columns is None:
rate_columns = [
h for h in header if h.startswith("firing_rate_") or h.startswith("spk") or h.startswith("layer")
]
rate_columns = [c for c in rate_columns if c in cols]
if not rate_columns:
return
plt.figure(figsize=(12, 6))
for col in rate_columns:
vals = [float(r[cols[col]]) for r in rows]
plt.plot(epochs, vals, label=col, linewidth=1.5)
plt.xlabel("Epoch")
plt.ylabel("Mean Firing Rate")
plt.title(title)
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(out_dir, f"{filename}.png"), dpi=150)
plt.close()