neuralpixel-convert-to-NWB/plot_all_units_clip_group_mean.py at main · liujh22/neuralpixel-convert-to-NWB · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from pynwb import NWBHDF5IO


ROOT = Path(__file__).resolve().parent
NWB_PATH = ROOT / "NWB" / "NP-SC12-jh.nwb"
OUT_DIR = ROOT / "NWB" / "analysis_NP-SC12-jh" / "trimmed_transition_control"


def rate_in_windows(spikes: np.ndarray, windows: list[tuple[float, float]]) -> float:
    n = 0
    d = 0.0
    for s, e in windows:
        n += int(np.sum((spikes >= s) & (spikes < e)))
        d += e - s
    return n / d if d > 0 else np.nan


def main():
    with NWBHDF5IO(str(NWB_PATH), "r", load_namespaces=True) as io:
        nwb = io.read()
        units = nwb.units.to_dataframe()

    spikes_all = [np.asarray(x, dtype=float) for x in units["spike_times"].values]

    # movie block and trimmed clip rules
    movie_start = 120.0
    routine_starts = [movie_start + i * 43.0 for i in range(3)]
    clip_offsets_trim = [(0.0, 9.5), (11.0, 20.5), (22.0, 31.5), (33.0, 42.5)]

    # R1/R2/R3: each repeat combines its 4 trimmed clips
    routine_rates = []
    for rs in routine_starts:
        ws = [(rs + a, rs + b) for a, b in clip_offsets_trim]
        rr = [rate_in_windows(spk, ws) for spk in spikes_all]
        routine_rates.append(np.asarray(rr))

    # C1..C4: each clip averaged over 3 repeats
    clip_rates = []
    for a, b in clip_offsets_trim:
        ws = [(rs + a, rs + b) for rs in routine_starts]
        cr = [rate_in_windows(spk, ws) for spk in spikes_all]
        clip_rates.append(np.asarray(cr))

    r_mean = np.array([x.mean() for x in routine_rates])
    r_sem = np.array([x.std(ddof=1) / np.sqrt(len(x)) for x in routine_rates])
    c_mean = np.array([x.mean() for x in clip_rates])
    c_sem = np.array([x.std(ddof=1) / np.sqrt(len(x)) for x in clip_rates])

    # clip-only figure (requested)
    fig, ax = plt.subplots(figsize=(6.2, 4.2))
    x = np.arange(4)
    ax.bar(x, c_mean, yerr=c_sem, capsize=4, color="tab:orange")
    ax.set_xticks(x)
    ax.set_xticklabels(["C1", "C2", "C3", "C4"])
    ax.set_ylabel("Rate (Hz)")
    ax.set_title("All units mean by video clip (trimmed)")
    fig.tight_layout()
    fig.savefig(OUT_DIR / "all_units_group_by_video_clip_mean.png", dpi=180)
    plt.close(fig)

    # routine + clip figure (same style as your right-lower panel intent)
    fig, ax = plt.subplots(figsize=(7.2, 4.4))
    x1 = np.arange(3)
    x2 = np.arange(4) + 4.2
    ax.bar(x1, r_mean, yerr=r_sem, capsize=4, color="tab:blue", label="routine 1/2/3")
    ax.bar(x2, c_mean, yerr=c_sem, capsize=4, color="tab:orange", label="clip 1/2/3/4")
    ax.set_xticks(list(x1) + list(x2))
    ax.set_xticklabels(["R1", "R2", "R3", "C1", "C2", "C3", "C4"])
    ax.set_ylabel("Rate (Hz)")
    ax.set_title("All units mean: routine/clip groups (trimmed)")
    ax.legend(frameon=False, fontsize=8)
    fig.tight_layout()
    fig.savefig(OUT_DIR / "all_units_routine_clip_group_mean.png", dpi=180)
    plt.close(fig)

    print(f"Saved: {OUT_DIR / 'all_units_group_by_video_clip_mean.png'}")
    print(f"Saved: {OUT_DIR / 'all_units_routine_clip_group_mean.png'}")


if __name__ == "__main__":
    main()