Skills › Scientific Agent Skills › neuropixels-analysis › run_sorting.py
run_sorting.py
102 lines · skills/neuropixels-analysis/scripts/run_sorting.py at 5f2bc03
1#!/usr/bin/env python
2"""
3Run spike sorting on preprocessed recording.
5Usage:
6 python run_sorting.py preprocessed/ --sorter kilosort4 --output sorting/
7"""
9import argparse
10from pathlib import Path
12import spikeinterface.full as si
13from _common import load_saved, validate_recording
16# Default parameters for each sorter
17SORTER_DEFAULTS = {
18 'kilosort4': {
19 'do_CAR': False,
20 'batch_size': 60000,
21 'nblocks': 1,
22 'Th_learned': 8,
23 'Th_universal': 9,
24 },
25 'kilosort3': {
26 'car': False, # Kilosort3 parameter; filtering/whitening stay sorter-specific
27 },
28 'spykingcircus2': {
29 'apply_preprocessing': False,
30 },
31 'mountainsort5': {
32 'filter': False,
33 'whiten': True,
34 },
35}
38def run_sorting(
39 input_path: str,
40 output_dir: str,
41 sorter: str = 'kilosort4',
42 sorter_params: dict = None,
43 n_jobs: int = -1,
44):
45 """Run spike sorting."""
47 print(f"Loading preprocessed recording from: {input_path}")
48 recording = load_saved(input_path, 'preprocessed')
49 validate_recording(recording)
50 si.set_global_job_kwargs(n_jobs=n_jobs, chunk_duration='1s')
52 print(f"Recording: {recording.get_num_channels()} channels, {recording.get_total_duration():.1f}s")
54 # Get sorter parameters
55 params = SORTER_DEFAULTS.get(sorter, {}).copy()
56 if sorter_params:
57 params.update(sorter_params)
59 print(f"Running {sorter} with params: {params}")
61 output_path = Path(output_dir)
63 # Run sorter (note: parameter is 'folder' not 'output_folder' in newer SpikeInterface)
64 sorting = si.run_sorter(
65 sorter,
66 recording,
67 folder=output_path / f'{sorter}_output',
68 verbose=True,
69 **params,
70 )
72 print(f"\nSorting complete!")
73 print(f" Units found: {len(sorting.unit_ids)}")
74 print(f" Total spikes: {sorting.count_total_num_spikes()}")
76 # Save sorting
77 sorting.save(folder=output_path / 'sorting')
78 print(f" Saved to: {output_path / 'sorting'}")
80 return sorting
83def main():
84 parser = argparse.ArgumentParser(description='Run spike sorting')
85 parser.add_argument('input', help='Path to preprocessed recording')
86 parser.add_argument('--output', '-o', default='sorting/', help='Output directory')
87 parser.add_argument('--sorter', '-s', default='kilosort4',
88 choices=['kilosort4', 'kilosort3', 'spykingcircus2', 'mountainsort5'])
89 parser.add_argument('--n-jobs', type=int, default=-1, help='Number of parallel jobs')
91 args = parser.parse_args()
93 run_sorting(
94 args.input,
95 args.output,
96 sorter=args.sorter,
97 n_jobs=args.n_jobs,
98 )
101if __name__ == '__main__':
102 main()