visualize icephys data for DANDI paper

gistfile1.txt

# analyze data from https://dandiarchive.org/dandiset/001753/0.260220.1838

"""Visualize intracellular electrophysiology data from an NWB file."""

import h5py
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import cm

NWB_PATH = "/Users/bdichter/Downloads/sub-1461085915_ses-1461686589_icephys.nwb"


def load_trace(f, group, key):
    """Load a trace with proper unit conversion, returning (time_s, data_scaled, unit)."""
    grp = f[f"{group}/{key}"]
    raw = grp["data"][:]
    conv = grp["data"].attrs.get("conversion", 1.0)
    unit = grp["data"].attrs.get("unit", "?")
    rate = grp["starting_time"].attrs.get("rate")
    t = np.arange(len(raw)) / rate
    return t, raw * conv, unit


def classify_traces(f):
    """Separate acquisition keys into voltage-clamp and current-clamp lists."""
    vc, cc = [], []
    for key in sorted(f["acquisition"].keys()):
        ntype = f["acquisition"][key].attrs.get("neurodata_type", "")
        if "VoltageClamp" in ntype:
            vc.append(key)
        elif "CurrentClamp" in ntype:
            cc.append(key)
    return vc, cc


def stim_key_for(acq_key):
    """Convert acquisition key (e.g. data_00006_AD0) to stimulus key (data_00006_DA0)."""
    return acq_key.replace("_AD0", "_DA0")


def get_stim_amplitude_pA(f, stim_key):
    """Get the max absolute stimulus amplitude in pA."""
    grp = f[f"stimulus/presentation/{stim_key}"]
    raw = grp["data"][:]
    conv = grp["data"].attrs.get("conversion", 1.0)
    scaled = raw * conv
    if np.all(np.isnan(scaled)):
        return 0.0
    return np.nanmax(np.abs(scaled)) * 1e12  # convert A to pA


def classify_stimulus_type(f, acq_key):
    """Classify a current-clamp sweep by its stimulus waveform."""
    sk = stim_key_for(acq_key)
    grp = f[f"stimulus/presentation/{sk}"]
    raw = grp["data"][:]
    conv = grp["data"].attrs.get("conversion", 1.0)
    rate = grp["starting_time"].attrs.get("rate")
    stim = raw * conv * 1e12  # pA
    dur_ms = len(stim) / rate * 1e3

    if np.all(np.isnan(stim)):
        return "Noise / Unknown"

    valid = stim[~np.isnan(stim)]
    unique_vals = len(np.unique(valid))
    mn, mx = np.nanmin(stim), np.nanmax(stim)

    if dur_ms < 200:
        return "Short Square Pulse"
    if dur_ms < 1000:
        return "Square Pulse Step"
    if dur_ms > 10000 and unique_vals > 50:
        return "Ramp"
    if dur_ms > 10000:
        return "Noise / Unknown"
    if mx <= 0:
        return "Hyperpolarizing Step"
    return "Long Step + Hyperpol."


def group_by_stimulus_type(f, cc_keys):
    """Group current-clamp keys by stimulus type."""
    groups = {}
    for key in cc_keys:
        stype = classify_stimulus_type(f, key)
        groups.setdefault(stype, []).append(key)
    return groups


def plot_square_pulse_vm(f, cc_keys):
    """Plot membrane potential for square pulse step sweeps only."""
    groups = group_by_stimulus_type(f, cc_keys)
    keys = groups["Square Pulse Step"]
    n_sweeps = len(keys)

    fig, ax = plt.subplots(figsize=(10, 5))

    amplitudes = [get_stim_amplitude_pA(f, stim_key_for(k)) for k in keys]
    amp_arr = np.array(amplitudes)
    norm = plt.Normalize(amp_arr.min(), amp_arr.max())
    cmap = cm.viridis

    for key, amp in zip(keys, amplitudes):
        color = cmap(norm(amp))
        t_resp, resp, _ = load_trace(f, "acquisition", key)
        ax.plot(t_resp * 1e3, resp * 1e3, color=color, alpha=0.8, linewidth=0.8)

    ax.set_ylabel("Membrane Potential (mV)")
    ax.set_xlabel("Time (ms)")
    ax.set_title(f"Square Pulse Step — Membrane Potential ({n_sweeps} sweeps)")
    ax.set_box_aspect(1)
    ax.spines[["top", "right"]].set_visible(False)

    sm = cm.ScalarMappable(cmap=cmap, norm=norm)
    fig.colorbar(sm, ax=ax, label="Stimulus Amplitude (pA)")
    fig.tight_layout()
    return fig


def main():
    with h5py.File(NWB_PATH, "r") as f:
        vc_keys, cc_keys = classify_traces(f)
        print(f"Found {len(vc_keys)} voltage-clamp and {len(cc_keys)} current-clamp sweeps")

        fig = plot_square_pulse_vm(f, cc_keys)
        fig.savefig("cc_overview.png", dpi=150, bbox_inches="tight")
        print("Saved cc_overview.png")


if __name__ == "__main__":
    main()