Skills › Scientific Agent Skills › neuropixels-analysis › compute_metrics.py
compute_metrics.py
175 lines · skills/neuropixels-analysis/scripts/compute_metrics.py at 5f2bc03
1#!/usr/bin/env python
2"""
3Compute quality metrics and curate units.
5Usage:
6 python compute_metrics.py sorting/ preprocessed/ --output metrics/
7"""
9import argparse
10from pathlib import Path
11import json
13import numpy as np
15from _common import load_saved, validate_recording
16import spikeinterface.full as si
19# Local screening presets, not reproductions of institutional pipelines.
20# "ibl" is a retained legacy name, not IBL's sliding-RP/noise-cutoff classifier.
21CURATION_CRITERIA = {
22 'allen': {
23 'snr': 3.0,
24 'isi_violations_ratio': 0.5,
25 'presence_ratio': 0.9,
26 'amplitude_cutoff': 0.1,
27 },
28 'ibl': {
29 'snr': 4.0,
30 'firing_rate': 0.1,
31 'isi_violations_ratio': 0.1,
32 'presence_ratio': 0.9,
33 'amplitude_cutoff': 0.1,
34 },
35 'strict': {
36 'snr': 5.0,
37 'isi_violations_ratio': 0.01,
38 'presence_ratio': 0.95,
39 'amplitude_cutoff': 0.01,
40 },
41}
44METRIC_NAMES = ['snr', 'isi_violation', 'presence_ratio', 'amplitude_cutoff', 'firing_rate', 'amplitude_cv']
47def curate_units(metrics, method='allen'):
48 """Screen units conservatively; absent/nonfinite evidence is unresolved."""
49 if method not in CURATION_CRITERIA:
50 raise ValueError(f"unknown curation method {method!r}; choose allen, ibl, or strict")
51 criteria = CURATION_CRITERIA[method]
52 labels = {}
53 for unit_id, row in metrics.iterrows():
54 values = {key: row.get(key, np.nan) for key in criteria}
55 if not all(np.isfinite(value) for value in values.values()):
56 labels[unit_id] = 'unsorted'
57 elif values['snr'] < 1.5:
58 labels[unit_id] = 'noise'
59 elif all(values[key] > limit if key in {'snr', 'presence_ratio', 'firing_rate'}
60 else values[key] < limit for key, limit in criteria.items()):
61 labels[unit_id] = 'good'
62 elif values['isi_violations_ratio'] > (0.05 if method == 'strict' else criteria['isi_violations_ratio']):
63 labels[unit_id] = 'mua'
64 else:
65 labels[unit_id] = 'unsorted'
66 return labels
69def compute_metrics(
70 sorting_path: str,
71 recording_path: str,
72 output_dir: str,
73 curation_method: str = 'allen',
74 n_jobs: int = -1,
75):
76 """Compute quality metrics and apply curation."""
78 print(f"Loading sorting from: {sorting_path}")
79 sorting = load_saved(sorting_path, 'sorting')
81 print(f"Loading recording from: {recording_path}")
82 recording = load_saved(recording_path, 'preprocessed')
83 validate_recording(recording)
84 si.set_global_job_kwargs(n_jobs=n_jobs, chunk_duration='1s')
86 print(f"Units: {len(sorting.unit_ids)}")
88 output_path = Path(output_dir)
89 output_path.mkdir(parents=True, exist_ok=True)
91 # Create analyzer
92 print("Creating SortingAnalyzer...")
93 analyzer = si.create_sorting_analyzer(
94 sorting,
95 recording,
96 format='binary_folder',
97 folder=output_path / 'analyzer',
98 sparse=True,
99 )
101 # Compute extensions
102 print("Computing waveforms...")
103 analyzer.compute('random_spikes', max_spikes_per_unit=500)
104 analyzer.compute('waveforms', ms_before=1.0, ms_after=2.0)
105 analyzer.compute('templates', operators=['average', 'std'])
107 print("Computing additional extensions...")
108 analyzer.compute('noise_levels')
109 analyzer.compute('spike_amplitudes')
110 analyzer.compute('amplitude_scalings')
111 analyzer.compute('correlograms', window_ms=50.0, bin_ms=1.0)
112 analyzer.compute('unit_locations', method='monopolar_triangulation')
113 analyzer.compute('template_similarity')
115 # Compute quality metrics
116 print("Computing quality metrics...")
117 metrics = si.compute_quality_metrics(
118 analyzer,
119 metric_names=METRIC_NAMES,
120 n_jobs=n_jobs,
121 )
123 # Save metrics
124 metrics.to_csv(output_path / 'quality_metrics.csv')
125 print(f"Saved metrics to: {output_path / 'quality_metrics.csv'}")
127 labels = curate_units(metrics, method=curation_method)
129 # Save labels
130 with open(output_path / 'curation_labels.json', 'w') as f:
131 json.dump({str(uid): label for uid, label in labels.items()}, f, indent=2)
133 # Summary
134 label_counts = {}
135 for label in labels.values():
136 label_counts[label] = label_counts.get(label, 0) + 1
138 print(f"\nCuration summary:")
139 print(f" Good: {label_counts.get('good', 0)}")
140 print(f" MUA: {label_counts.get('mua', 0)}")
141 print(f" Noise: {label_counts.get('noise', 0)}")
142 print(f" Unsorted: {label_counts.get('unsorted', 0)}")
143 print(f" Total: {len(labels)}")
145 # Metrics summary
146 print(f"\nMetrics summary:")
147 for col in ['snr', 'isi_violations_ratio', 'presence_ratio', 'firing_rate']:
148 if col in metrics.columns:
149 print(f" {col}: {metrics[col].median():.4f} (median)")
151 return analyzer, metrics, labels
154def main():
155 parser = argparse.ArgumentParser(description='Compute quality metrics')
156 parser.add_argument('sorting', help='Path to sorting directory')
157 parser.add_argument('recording', help='Path to preprocessed recording')
158 parser.add_argument('--output', '-o', default='metrics/', help='Output directory')
159 parser.add_argument('--curation', '-c', default='allen',
160 choices=['allen', 'ibl', 'strict'])
161 parser.add_argument('--n-jobs', type=int, default=-1, help='Number of parallel jobs')
163 args = parser.parse_args()
165 compute_metrics(
166 args.sorting,
167 args.recording,
168 args.output,
169 curation_method=args.curation,
170 n_jobs=args.n_jobs,
171 )
174if __name__ == '__main__':
175 main()