Skip to content

visualise_ssc_encoding.py

SSC DyNED encoding visualisation - 280 lines. View on GitHub (speech-neuro/visualise_ssc_encoding.py).

"""Visualise SSC neuromorphic data through the DyNED encoding pipeline.

10 samples on a single sheet, 4 columns:
  1. Original  -  raw cochlea spike raster
  2. DyNED Spikes  -  binary step signal (Greys)
  3. DyNEDc Compressed  -  compressed bitstream (zero-padded) with CR annotation
  4. Reconstructed  -  decoded from quantised values, with SNR annotation

Usage:
    uv run python speech-neuro/visualise_ssc_encoding.py
"""

import os
import sys
import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

try:
    _SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
    _SCRIPT_DIR = os.getcwd()
sys.path.insert(0, os.path.join(_SCRIPT_DIR, ".."))

import h5py
from dyned import (
    sigma_delta_quantisation,
    generate_step_signal,
    DyNEDcCompressorV4,
)

OUTPUT_DIR = os.path.join(_SCRIPT_DIR, "ssc_encoding_vis")
os.makedirs(OUTPUT_DIR, exist_ok=True)

SSC_DIR = os.path.join(_SCRIPT_DIR, "..", "assets", "SSC")
SSC_N_CHANNELS = 700
SSC_N_BINS = 50
N_SAMPLES = 10
SEED = 42

SSC_LABELS = [
    "backward",
    "bed",
    "bird",
    "cat",
    "dog",
    "down",
    "eight",
    "five",
    "follow",
    "forward",
    "four",
    "go",
    "happy",
    "house",
    "learn",
    "left",
    "marvin",
    "nine",
    "no",
    "off",
    "on",
    "one",
    "right",
    "seven",
    "sheila",
    "six",
    "stop",
    "three",
    "tree",
    "two",
    "up",
    "visual",
    "wow",
    "yes",
    "zero",
]


def ssc_to_dense(times, units, n_channels=SSC_N_CHANNELS, n_bins=SSC_N_BINS):
    """Convert raw SSC spike events to a dense [n_channels, n_bins] histogram."""
    times = np.asarray(times, dtype=np.float32)
    units = np.asarray(units, dtype=np.int64)
    if len(times) == 0:
        return np.zeros((n_channels, n_bins), dtype=np.float32)
    t_min, t_max = times.min(), times.max()
    if t_max - t_min <= 0:
        return np.zeros((n_channels, n_bins), dtype=np.float32)
    bin_edges = np.linspace(t_min, t_max, n_bins + 1)
    bin_indices = np.clip(np.digitize(times, bin_edges) - 1, 0, n_bins - 1)
    dense = np.zeros((n_channels, n_bins), dtype=np.float32)
    for t_idx, ch in zip(bin_indices, units):
        if 0 <= ch < n_channels:
            dense[ch, t_idx] += 1.0
    return dense


def dyned_encode(dense, levels=256):
    """Per-time-bin DyNED sigma-delta encoding.

    Returns: quantised_dense, step_signal, log_dense (all [n_channels, n_bins])
    """
    log_dense = np.log1p(dense)
    n_channels, n_bins = dense.shape
    quantised_dense = np.zeros_like(log_dense)
    step_signal = np.zeros_like(log_dense)
    for t in range(n_bins):
        frame = log_dense[:, t]
        quantised, _ = sigma_delta_quantisation(frame, levels=levels)
        step = generate_step_signal(quantised)
        quantised_dense[:, t] = quantised
        step_signal[:, t] = step.astype(np.float32)
    return quantised_dense, step_signal, log_dense


def compute_snr(original, reconstructed):
    """Compute SNR in dB."""
    noise = original - reconstructed
    sig_power = np.sum(original**2)
    noise_power = np.sum(noise**2)
    return 10 * np.log10(sig_power / (noise_power + 1e-12))


def main():
    print("SSC DyNED Encoding Pipeline Visualisation")

    h5_path = os.path.join(SSC_DIR, "ssc_valid.h5")
    if not os.path.exists(h5_path):
        print(f"Error: {h5_path} not found")
        sys.exit(1)

    rng = np.random.default_rng(SEED)

    with h5py.File(h5_path, "r") as f:
        all_labels = f["labels"][:]

        # Pick one sample per unique label (first N_SAMPLES)
        unique_labels = np.unique(all_labels)
        chosen = []
        for lab in rng.permutation(unique_labels):
            if len(chosen) >= N_SAMPLES:
                break
            candidates = np.where(all_labels == lab)[0]
            chosen.append(rng.choice(candidates))

        # Pre-compute all data
        rows = []
        for sample_idx in chosen:
            times = f["spikes/times"][sample_idx]
            units = f["spikes/units"][sample_idx]
            label_idx = int(f["labels"][sample_idx])
            label_name = SSC_LABELS[label_idx] if label_idx < len(SSC_LABELS) else str(label_idx)

            dense = ssc_to_dense(times, units)
            quantised, step_signal, log_dense = dyned_encode(dense)
            snr = compute_snr(log_dense, quantised)

            # DyNEDc compress
            compressor = DyNEDcCompressorV4(chunk_size=4)
            step_flat = step_signal.flatten().astype(np.uint8)
            compressed, info = compressor.compress(step_flat)
            orig_bits = len(step_flat)
            comp_bits = len(compressed)
            cr = comp_bits / orig_bits if orig_bits > 0 else 1.0
            space_saved = (1.0 - cr) * 100

            # Build compressed visualisation: fill same-shaped image with compressed bits, rest=0
            comp_vis = np.zeros(orig_bits, dtype=np.float32)
            comp_arr = np.array([int(b) for b in compressed[:orig_bits]], dtype=np.float32)
            comp_vis[: len(comp_arr)] = comp_arr
            comp_vis_img = comp_vis.reshape(step_signal.shape)

            rows.append(
                {
                    "label": label_name,
                    "times": np.asarray(times, dtype=np.float32),
                    "units": np.asarray(units, dtype=np.int64),
                    "dense": dense,
                    "log_dense": log_dense,
                    "quantised": quantised,
                    "step_signal": step_signal,
                    "comp_vis": comp_vis_img,
                    "cr": cr,
                    "space_saved": space_saved,
                    "snr": snr,
                }
            )
            print(f"  {label_name}: SNR={snr:.1f} dB | CR={cr:.3f} ({space_saved:.1f}% saved)")

    # -- Single sheet: 10 rows x 4 columns --
    fig, axes = plt.subplots(N_SAMPLES, 4, figsize=(18, 28))

    extent = [0, SSC_N_BINS, 0, SSC_N_CHANNELS]

    for row, data in enumerate(rows):
        t = data["times"]
        u = data["units"]

        # Subsample for plotting
        if len(t) > 20000:
            idx_sub = rng.choice(len(t), 20000, replace=False)
            t_plot, u_plot = t[idx_sub], u[idx_sub]
        else:
            t_plot, u_plot = t, u

        # Col 0: Original cochlea spike raster
        axes[row, 0].scatter(t_plot, u_plot, s=0.1, c="#2C3E50", alpha=0.3, rasterized=True)
        axes[row, 0].set_ylim(0, SSC_N_CHANNELS)
        axes[row, 0].set_ylabel(
            f'"{data["label"]}"', fontsize=10, fontweight="bold", rotation=0, labelpad=50, ha="right"
        )
        if row == 0:
            axes[row, 0].set_title("Original (Cochlea)", fontsize=10, fontweight="bold")
        if row == N_SAMPLES - 1:
            axes[row, 0].set_xlabel("Time (s)")

        # Col 1: DyNED spikes
        axes[row, 1].imshow(data["step_signal"], aspect="auto", origin="lower", cmap="Greys", interpolation="none")
        if row == 0:
            axes[row, 1].set_title("DyNED Spikes", fontsize=10, fontweight="bold")
        if row == N_SAMPLES - 1:
            axes[row, 1].set_xlabel("Time Bin")

        # Col 2: DyNEDc compressed bitstream
        axes[row, 2].imshow(data["comp_vis"], aspect="auto", origin="lower", cmap="Greys", interpolation="none")
        axes[row, 2].text(
            0.98,
            0.95,
            f"CR: {data['cr']:.3f}\n{data['space_saved']:.1f}% saved",
            transform=axes[row, 2].transAxes,
            fontsize=7,
            ha="right",
            va="top",
            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8),
        )
        if row == 0:
            axes[row, 2].set_title("DyNEDc Compressed", fontsize=10, fontweight="bold")
        if row == N_SAMPLES - 1:
            axes[row, 2].set_xlabel("Time Bin")

        # Col 3: Reconstructed (quantised values shown as image with SNR)
        axes[row, 3].imshow(data["quantised"], aspect="auto", origin="lower", cmap="inferno", extent=extent)
        axes[row, 3].text(
            0.98,
            0.95,
            f"SNR: {data['snr']:.1f} dB",
            transform=axes[row, 3].transAxes,
            fontsize=8,
            ha="right",
            va="top",
            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8),
        )
        if row == 0:
            axes[row, 3].set_title("Reconstructed", fontsize=10, fontweight="bold")
        if row == N_SAMPLES - 1:
            axes[row, 3].set_xlabel("Time Bin")

        # Tick-label cleanup: only the bottom row shows x-tick numbers and only
        # the leftmost column shows y-tick numbers. The reference scales sit on
        # the figure edges; interior subplots gain pixel area for content.
        for col in range(4):
            cur = axes[row, col]
            if row != N_SAMPLES - 1:
                cur.tick_params(axis="x", labelbottom=False)
            if col != 0:
                cur.tick_params(axis="y", labelleft=False)

    fig.suptitle("DyNED Encoding Pipeline  -  10 SSC Samples", fontsize=14, fontweight="bold", y=0.995)
    plt.tight_layout()
    out_path = os.path.join(OUTPUT_DIR, "ssc_encoding_grid.png")
    fig.savefig(out_path, dpi=200, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"\nSaved: {out_path}")


if __name__ == "__main__":
    main()