PlotWaveforms.py

# -*- coding: utf-8 -*-
"""
Plot spike waveforms from FRtable.mat files.
Mean +/- SEM shading, color-coded by SU (orange) / MU (blue).
Saves each waveform as PDF and PNG.
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat
from pathlib import Path

# %% ========== SETTINGS ==========
plt.rcParams['svg.fonttype'] = 'none'
plt.rcParams['pdf.fonttype'] = 'truetype'


data_root = '../Optimum1_ProcessedData'
output_dir = 'Figures/waveforms'
figsize = (1, 0.9)

units_to_plot = [
    {'label': 'NSX98-1', 'patient': 'P06', 'session': 1},
    {'label': 'NSX87-3', 'patient': 'P08', 'session': 2},
    {'label': 'NSX65-1', 'patient': 'P13', 'session': 1},
    {'label': 'NSX68-1',  'patient': 'P13', 'session': 1},
    {'label': 'NSX72-1',  'patient': 'P02', 'session': 1},
    {'label': 'NSX72-1',  'patient': 'P02', 'session': 2}
]

# %% ========== FUNCTIONS ==========

def load_waveform_data(mat_file):
    """Load ALLspikes waveforms, labels, SUA classification, and spike timestamps."""
    mat = loadmat(mat_file, squeeze_me=True)

    waveforms = mat['ALLspikes']['waveform'].item()
    labels = mat['ALLspikesInfo']['label'].item()
    is_sua = mat['ALLspikesInfo']['isSUA'].item()

    # Try to load spike timestamps for ISI computation
    # Field is typically 'times' in ALLspikes; adjust if named differently in your mat files
    try:
        times = mat['ALLspikes']['timestamp'].item()
    except (KeyError, ValueError):
        times = None

    if isinstance(labels, str):
        labels = [labels]
        waveforms = [waveforms]
        is_sua = np.array([is_sua])
        if times is not None:
            times = [times]

    return waveforms, labels, is_sua, times


def find_unit_index(mat_labels, unique_label):
    """Match UniqueLabel (P06_NSX124-1S1) to mat label (NSX124-1S1)."""
    parts = unique_label.split('_', 1)
    search_label = parts[1] if len(parts) > 1 else unique_label
    
    for i, lab in enumerate(mat_labels):
        if lab == search_label:
            return i
    
    raise ValueError(f"'{search_label}' not found. Available: {list(mat_labels)}")


def plot_waveform(waveform_data, label, is_su, output_dir, figsize=(2.0, 1.8)):
    """Plot single waveform: mean +/- SEM, SU=orange, MU=blue."""
    wd = np.array(waveform_data)
    if wd.ndim == 3:
        wd = wd.squeeze()
    if wd.ndim == 1:
        wd = wd[:, np.newaxis]
    if wd.shape[0] != 64 and wd.shape[1] == 64:
        wd = wd.T
    
    time_ms = np.arange(wd.shape[0])
    mean_wf = np.mean(wd, axis=1)
    sem_wf = np.std(wd, axis=1) 
    
    color = '#ff7f0e' if is_su else '#1f77b4'
    unit_type = 'SU' if is_su else 'MU'
    
    fig, ax = plt.subplots(figsize=figsize)
    ax.plot(time_ms, mean_wf, color=color, linewidth=1.5)
    ax.fill_between(time_ms, mean_wf - sem_wf, mean_wf + sem_wf,
                    color=color, alpha=0.3, linewidth=0)
    
    for spine in ax.spines.values():
        spine.set_visible(False)
    ax.set_xticks([])
    ax.set_yticks([])
    
    # L-shaped scale bar (bottom-right)
    x_bar = 20  # ms
    y_bar_raw = np.ptp(mean_wf)  # peak-to-peak range
    # Round y_bar to a nice number
    nice_vals = [5, 10, 15, 20, 25, 30, 40, 50, 75, 100]
    y_bar = min(nice_vals, key=lambda v: abs(v - y_bar_raw * 0.4))
    
    # Position: bottom-right of data range
    x0 = time_ms[-1] - x_bar - 2
    y0 = np.min(mean_wf - sem_wf) - y_bar * 0.15
    
    # Horizontal bar (time)
    ax.plot([x0, x0 + x_bar], [y0, y0], color='k', linewidth=1.2, clip_on=False)
    # Vertical bar (voltage)
    ax.plot([x0, x0], [y0, y0 + y_bar], color='k', linewidth=1.2, clip_on=False)
    
    # Labels
    ax.text(x0 + x_bar*0.8, y0 - y_bar * 0.15, f'{x_bar} ms',
            ha='center', va='top', fontsize=6, color='k')
    ax.text(x0 - 1, y0 + y_bar / 2, f'{y_bar} µV',
            ha='right', va='center', fontsize=6, color='k', rotation=90)
    
    plt.tight_layout()
    
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    safe_label = label.replace('/', '_').replace('\\', '_')
    fig.savefig(out / f'waveform_{safe_label}.png', dpi=300, bbox_inches='tight')
    fig.savefig(out / f'waveform_{safe_label}.pdf', bbox_inches='tight')
    plt.close(fig)
    print(f"Saved: {safe_label} ({unit_type}, {wd.shape[1]} spikes)")


def plot_isi_histogram(spike_times, label, output_dir, max_isi_ms=100, refrac_ms=3, figsize=(1, 0.9)):
    """Plot ISI histogram for a single unit.

    spike_times : array of timestamps (seconds or samples at 30 kHz; auto-detected)
    """
    if spike_times is None or len(np.asarray(spike_times).ravel()) < 2:
        print(f"  Skipping ISI: no timestamps for {label}")
        return

    ts = np.sort(np.asarray(spike_times, dtype=float).ravel())
    isis = np.diff(ts)

    # Convert to ms: if median ISI < 0.5 assume seconds; if >> 1000 assume 30 kHz samples
    # samples at 30 kHz → seconds → ms
    isis = isis / 30.0
 
    pct_refrac = 100.0 * np.sum(isis < refrac_ms) / len(isis)

    bins = np.arange(0, max_isi_ms + 1, 1)
    fig, ax = plt.subplots(figsize=figsize)
    ax.hist(isis[isis <= max_isi_ms], bins=bins, color='#1f77b4', edgecolor='#1f77b4', linewidth=0.2)

    ax.set_xlabel('ISI (ms)', fontsize=5, fontfamily='Arial')
    ax.set_ylabel('Count', fontsize=5, fontfamily='Arial')
    ax.tick_params(labelsize=4, pad=1)
    ax.set_title(f'{pct_refrac:.1f}% < {refrac_ms} ms', fontsize=5, fontfamily='Arial', pad=2)
    ax.set_xlim(0, max_isi_ms)
    ax.set_xticks([0, 50, 100])
    for spine in ['top', 'right']:
        ax.spines[spine].set_visible(False)
    ax.spines['left'].set_linewidth(0.6)
    ax.spines['bottom'].set_linewidth(0.6)
    ax.tick_params(width=0.6)

    plt.tight_layout(pad=0.3)

    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    safe_label = label.replace('/', '_').replace('\\', '_')
    fig.savefig(out / f'isi_{safe_label}.png', dpi=300, bbox_inches='tight')
    fig.savefig(out / f'isi_{safe_label}.pdf', bbox_inches='tight')
    plt.close(fig)
    print(f"  ISI saved: {safe_label}  ({pct_refrac:.1f}% violations < {refrac_ms} ms)")


# %% ========== RUN ==========

# Cache loaded mat files to avoid reloading
_cache = {}

for unit in units_to_plot:
    pat = unit['patient']
    ses = unit['session']
    label = unit['label']
    unique_label = f"{pat}_{label}"
    
    mat_path = Path(data_root) / f'Session{ses}' / 'FRTabs' / f'{pat}_FRtable.mat'
    
    # Load (with caching)
    cache_key = str(mat_path)
    if cache_key not in _cache:
        print(f"\nLoading {mat_path}")
        _cache[cache_key] = load_waveform_data(mat_path)
    waveforms, mat_labels, is_sua, times = _cache[cache_key]

    idx = find_unit_index(mat_labels, unique_label)
    plot_waveform(waveforms[idx], unique_label, bool(is_sua[idx]), output_dir, figsize)

    unit_times = times[idx] if times is not None else None
    plot_isi_histogram(unit_times, unique_label, output_dir, figsize=figsize)