clusters/wolf.py at main · MattNolanLab/clusters · 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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
from pathlib import Path
from argparse import ArgumentParser
import json

from spikeinterface_gui import run_mainwindow
import spikeinterface.full as si


parser = ArgumentParser()

parser.add_argument('--mouse')
parser.add_argument('--day')
parser.add_argument('--session')
parser.add_argument('--recording', action='store_true', default=False)
parser.add_argument('--protocol', default='kilosort4A')
parser.add_argument('--curation', default='curationA')
parser.add_argument('--MMNAV_folder', default=None)

args = parser.parse_args()

mouse, day, session, protocol, curation, MMNAV_folder = int(args.mouse), int(args.day), args.session, args.protocol, args.curation, args.MMNAV_folder

if MMNAV_folder is None:
    MMNAV_folder = '/run/user/1000/gvfs/smb-share:server=cmvm.datastore.ed.ac.uk,share=cmvm/sbms/groups/CDBS_SIDB_storage/NolanLab/ActiveProjects/Wolf/MMNAV/'
MMNAV_folder = Path(MMNAV_folder)

deriv_folder = MMNAV_folder / 'derivatives'

recording = None
if args.recording:
    raw_folder = MMNAV_folder / 'raw'
    rec_folder = list(raw_folder.glob(f'*/M{mouse:02d}_D{day:02d}_*_{session}'))[0]
    recording = si.read_openephys(rec_folder)
    recording = si.common_reference(si.bandpass_filter(recording))


mouseday_path = deriv_folder / f'M{mouse:02d}/D{day:02d}/{session}/{protocol}'

analyzer_path = mouseday_path / f'sub-{mouse:02d}_day-{day:02d}_ses-{session}_srt-{protocol}_analyzer.zarr'
curation_path = mouseday_path / f'sub-{mouse:02d}_day-{day:02d}_ses-{session}_srt-{protocol}_{curation}.json'

analyzer = si.load_sorting_analyzer(analyzer_path)

no_spikes_units = analyzer.unit_ids[(analyzer.get_extension('quality_metrics').get_data()['num_spikes'] == 0).values]
analyzer_removed = analyzer.remove_units(no_spikes_units)
kept_unit_ids = analyzer_removed.unit_ids

if curation_path.is_file():
    with open(curation_path) as f:
        curation_dict = json.load(f)

    curation_dict['unit_ids'] = list(kept_unit_ids)

    new_manual_labels = []
    for manual_labels in curation_dict['manual_labels']:
        if manual_labels['unit_id'] in kept_unit_ids:
            new_manual_labels.append(manual_labels)

    curation_dict['manual_labels'] = new_manual_labels

    new_removed = []
    for removed_unit_id in curation_dict['removed']:
        if removed_unit_id in kept_unit_ids:
            new_removed.append(removed_unit_id)

    curation_dict['removed'] = new_removed
else:
    curation_dict=None

wolf_layout = dict(
    zone1=['spikelist', 'curation'],
    zone2=['unitlist', 'merge'],
    zone3=['trace', 'tracemap', 'spikeamplitude'],
    zone4=[],
    zone5=['spikerate'],
    zone6=['probe'],
    zone7=['waveform'],
    zone8=['correlogram', 'metrics', 'mainsettings'],
)

overlap = False
if args.recording:
    overlap = True

user_settings = {
    "waveform": {
        "overlap": overlap,
        "plot_selected_spike": True,
        "plot_waveforms_samples": False,
        "y_scalebar": True,
    },
    "spikeamplitude": {
        "max_spikes_per_unit": 10000,
    },
    "spikerate": {
        "bin_s": 30
    },
}

run_mainwindow(
    analyzer_removed,
    mode="desktop",
    curation=True,
    curation_dict = curation_dict,
    skip_extensions = ['waveforms', 'spike_locations'],
    verbose=True,
    layout=wolf_layout,
    user_settings=user_settings,
    recording=recording,
)