viewer/raster_plot.py at main · Senzai-Lab/viewer · 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
from viewer import RasterView, Spikes, EventBars,show
import numpy as np
import pandas as pd
from pathlib import Path
from scipy.io import loadmat

if __name__ == "__main__":
    save_path = Path('.tmp/')
    metadata = pd.read_csv(save_path / 'metadata.csv', index_col=0).to_dict()
    spikes = Spikes(
        name="units",
        ts=np.load(save_path / 'spike_times.npy'),
        spike_units=np.load(save_path / 'spike_clusters.npy'),
    )

    sleep_states = loadmat(r"R:\Basic_Sciences\Phys\SenzaiLab\Elissa_Belluccini\Kilosort\Mouse11_20260603_225to1665Stag_510to1935Stag\ProbeA\eeg\eeg.SleepState.states.mat")['SleepState']
    sleep = sleep_states[0][0][0][0][0]
    wake = sleep[0]
    nrem = sleep[1]
    rem = sleep[2]

    starts = np.concatenate([wake[:, 0], nrem[:, 0], rem[:, 0]])
    ends = np.concatenate([wake[:, 1], nrem[:, 1], rem[:, 1]])
    labels = np.concatenate([np.full(len(wake), 'wake'), np.full(len(nrem), 'nrem'), np.full(len(rem), 'rem')])

    order = np.argsort(starts, kind='stable')

    starts = starts[order]
    ends = ends[order]
    labels = labels[order]
    sleep_bars = EventBars(
        starts=starts, ends=ends, labels=labels,
        label_order=['wake', 'nrem', 'rem'],
        colors={
            'wake': '#6b8cff', # blue
            'nrem': '#47d18c', # green
            'rem': '#f26d7d' # red
            })

    show(
        [(spikes, RasterView(metadata=metadata))],
        overlays=(sleep_bars,),
        title="Raster")