Initial commit for neuropixels-analysis

scripts/run_sorting.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+"""
+Run spike sorting on preprocessed recording.
+
+Usage:
+    python run_sorting.py preprocessed/ --sorter kilosort4 --output sorting/
+"""
+
+import argparse
+from pathlib import Path
+
+import spikeinterface.full as si
+
+
+# Default parameters for each sorter
+SORTER_DEFAULTS = {
+    'kilosort4': {
+        'batch_size': 30000,
+        'nblocks': 1,
+        'Th_learned': 8,
+        'Th_universal': 9,
+    },
+    'kilosort3': {
+        'do_CAR': False,  # Already done in preprocessing
+    },
+    'spykingcircus2': {
+        'apply_preprocessing': False,
+    },
+    'mountainsort5': {
+        'filter': False,
+        'whiten': False,
+    },
+}
+
+
+def run_sorting(
+    input_path: str,
+    output_dir: str,
+    sorter: str = 'kilosort4',
+    sorter_params: dict = None,
+    n_jobs: int = -1,
+):
+    """Run spike sorting."""
+
+    print(f"Loading preprocessed recording from: {input_path}")
+    recording = si.load_extractor(Path(input_path) / 'preprocessed')
+
+    print(f"Recording: {recording.get_num_channels()} channels, {recording.get_total_duration():.1f}s")
+
+    # Get sorter parameters
+    params = SORTER_DEFAULTS.get(sorter, {}).copy()
+    if sorter_params:
+        params.update(sorter_params)
+
+    print(f"Running {sorter} with params: {params}")
+
+    output_path = Path(output_dir)
+
+    # Run sorter (note: parameter is 'folder' not 'output_folder' in newer SpikeInterface)
+    sorting = si.run_sorter(
+        sorter,
+        recording,
+        folder=output_path / f'{sorter}_output',
+        verbose=True,
+        **params,
+    )
+
+    print(f"\nSorting complete!")
+    print(f"  Units found: {len(sorting.unit_ids)}")
+    print(f"  Total spikes: {sum(len(sorting.get_unit_spike_train(uid)) for uid in sorting.unit_ids)}")
+
+    # Save sorting
+    sorting.save(folder=output_path / 'sorting')
+    print(f"  Saved to: {output_path / 'sorting'}")
+
+    return sorting
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Run spike sorting')
+    parser.add_argument('input', help='Path to preprocessed recording')
+    parser.add_argument('--output', '-o', default='sorting/', help='Output directory')
+    parser.add_argument('--sorter', '-s', default='kilosort4',
+                       choices=['kilosort4', 'kilosort3', 'spykingcircus2', 'mountainsort5'])
+    parser.add_argument('--n-jobs', type=int, default=-1, help='Number of parallel jobs')
+
+    args = parser.parse_args()
+
+    run_sorting(
+        args.input,
+        args.output,
+        sorter=args.sorter,
+        n_jobs=args.n_jobs,
+    )
+
+
+if __name__ == '__main__':
+    main()