Skills › Scientific Agent Skills › neuropixels-analysis › neuropixels_pipeline.py
neuropixels_pipeline.py
381 lines · skills/neuropixels-analysis/scripts/neuropixels_pipeline.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Neuropixels Data Analysis Pipeline (Best Practices Version)
5Based on SpikeInterface, Allen Institute, and IBL recommendations.
7Usage:
8 python neuropixels_pipeline.py /path/to/spikeglx/data /path/to/output
10References:
11 - https://spikeinterface.readthedocs.io/en/stable/how_to/analyze_neuropixels.html
12 - https://github.com/AllenInstitute/ecephys_spike_sorting
13"""
15import argparse
16from pathlib import Path
17import json
18import spikeinterface.full as si
19import numpy as np
21from _common import validate_recording, apply_phase_correction, reference_by_shank
22from compute_metrics import curate_units, METRIC_NAMES
25def load_recording(data_path: str, stream_name: str = 'imec0.ap') -> si.BaseRecording:
26 """Load a SpikeGLX or Open Ephys recording."""
28 data_path = Path(data_path)
30 # Auto-detect format
31 if any(data_path.rglob('*.ap.bin')) or any(data_path.rglob('*.ap.meta')):
32 # SpikeGLX format
33 streams, _ = si.get_neo_streams('spikeglx', data_path)
34 print(f"Available streams: {streams}")
35 recording = si.read_spikeglx(data_path, stream_name=stream_name)
36 elif any(data_path.rglob('*.oebin')):
37 # Open Ephys format
38 recording = si.read_openephys(data_path, stream_name=stream_name)
39 else:
40 raise ValueError(f"Unknown format in {data_path}")
42 validate_recording(recording)
43 print(f"Loaded recording:")
44 print(f" Channels: {recording.get_num_channels()}")
45 print(f" Duration: {recording.get_total_duration():.2f} s")
46 print(f" Sampling rate: {recording.get_sampling_frequency()} Hz")
48 return recording
51def preprocess(
52 recording: si.BaseRecording,
53 apply_phase_shift: bool = True,
54 freq_min: float = 400.,
55) -> tuple:
56 """
57 Apply standard Neuropixels preprocessing.
59 Following SpikeInterface recommendations:
60 1. High-pass filter at 400 Hz (not 300)
61 2. Detect and remove bad channels
62 3. Phase shift using acquisition ADC timing metadata
63 4. Common median reference
64 """
65 print("Preprocessing...")
67 # Step 1: High-pass filter
68 validate_recording(recording)
69 rec = si.highpass_filter(recording, freq_min=freq_min)
70 print(f" Applied high-pass filter at {freq_min} Hz")
72 # Step 2: Detect bad channels
73 bad_channel_ids, channel_labels = si.detect_bad_channels(rec)
74 if len(bad_channel_ids) > 0:
75 print(f" Detected {len(bad_channel_ids)} bad channels: {bad_channel_ids}")
76 rec = rec.remove_channels(bad_channel_ids)
77 else:
78 print(" No bad channels detected")
80 # Step 3: Phase shift from acquisition timing metadata
81 if apply_phase_shift:
82 rec = apply_phase_correction(rec)
83 print(" Applied phase shift correction")
85 # Step 4: Common median reference
86 rec = reference_by_shank(rec)
87 print(" Applied common median reference")
89 return rec, bad_channel_ids
92def check_drift(recording: si.BaseRecording, output_folder: str) -> dict:
93 """
94 Detect peaks and check for drift before spike sorting.
95 """
96 print("Checking for drift...")
98 from spikeinterface.sortingcomponents.peak_detection import detect_peaks
99 from spikeinterface.sortingcomponents.peak_localization import localize_peaks
101 job_kwargs = dict(n_jobs=8, chunk_duration='1s', progress_bar=True)
103 # Get noise levels
104 noise_levels = si.get_noise_levels(recording, return_in_uV=False)
106 # Detect peaks
107 peaks = detect_peaks(recording, method='locally_exclusive', method_kwargs={'noise_levels': noise_levels, 'detect_threshold': 5, 'radius_um': 50.0}, job_kwargs=job_kwargs)
108 print(f" Detected {len(peaks)} peaks")
110 # Localize peaks
111 peak_locations = localize_peaks(recording, peaks, method='center_of_mass', job_kwargs=job_kwargs)
113 # Save drift plot
114 import matplotlib.pyplot as plt
115 fig, ax = plt.subplots(figsize=(12, 6))
117 # Subsample for plotting
118 n_plot = min(100000, len(peaks))
119 idx = np.random.choice(len(peaks), n_plot, replace=False)
121 ax.scatter(
122 peaks['sample_index'][idx] / recording.get_sampling_frequency(),
123 peak_locations['y'][idx],
124 s=1, alpha=0.1, c='k'
125 )
126 ax.set_xlabel('Time (s)')
127 ax.set_ylabel('Depth (μm)')
128 ax.set_title('Peak Activity (Check for Drift)')
130 plt.savefig(f'{output_folder}/drift_check.png', dpi=150, bbox_inches='tight')
131 plt.close()
132 print(f" Saved drift plot to {output_folder}/drift_check.png")
134 # Peak depth spread reflects the population's spatial extent, not displacement
135 # over time. Use a motion estimator (correct_motion below), not depth percentiles.
136 return {
137 'peaks': peaks,
138 'peak_locations': peak_locations,
139 }
142def correct_motion(
143 recording: si.BaseRecording,
144 output_folder: str,
145 preset: str = 'nonrigid_fast_and_accurate'
146) -> si.BaseRecording:
147 """Apply motion correction if needed."""
148 print(f"Applying motion correction (preset: {preset})...")
150 probe = recording.get_probe()
151 if probe.shank_ids is None or len(np.unique(probe.shank_ids)) != 1:
152 raise ValueError("Motion estimation in this pipeline requires one verified shank; select channels per shank first.")
154 # correct_motion returns just the corrected recording by default. Pass
155 # output_motion_info=True only if you also want the motion info dict (a tuple is
156 # returned in that case).
157 rec_corrected = si.correct_motion(
158 recording,
159 preset=preset,
160 folder=f'{output_folder}/motion',
161 n_jobs=8,
162 chunk_duration='1s',
163 progress_bar=True
164 )
166 print(" Motion correction complete")
167 return rec_corrected
170def run_spike_sorting(
171 recording: si.BaseRecording,
172 output_folder: str,
173 sorter: str = 'kilosort4',
174 motion_corrected: bool = False,
175) -> si.BaseSorting:
176 """Run spike sorting."""
177 print(f"Running spike sorting with {sorter}...")
179 sorter_folder = f'{output_folder}/sorting_{sorter}'
180 sorter_params = {}
181 if motion_corrected:
182 # SpikeInterface 0.105.0 uses different motion flags for these sorters.
183 if sorter in {'kilosort2_5', 'kilosort3', 'kilosort4'}:
184 sorter_params['do_correction'] = False
185 elif sorter == 'spykingcircus2':
186 sorter_params['apply_motion_correction'] = False
188 sorting = si.run_sorter(
189 sorter,
190 recording,
191 folder=sorter_folder,
192 verbose=True,
193 **sorter_params,
194 )
196 print(f" Found {len(sorting.unit_ids)} units")
197 print(f" Total spikes: {sorting.count_total_num_spikes()}")
199 return sorting
202def postprocess(
203 sorting: si.BaseSorting,
204 recording: si.BaseRecording,
205 output_folder: str
206) -> tuple:
207 """Run post-processing and compute quality metrics."""
208 print("Post-processing...")
210 job_kwargs = dict(n_jobs=8, chunk_duration='1s', progress_bar=True)
212 # Create analyzer
213 analyzer = si.create_sorting_analyzer(
214 sorting, recording,
215 sparse=True,
216 format='binary_folder',
217 folder=f'{output_folder}/analyzer'
218 )
220 # Compute extensions (order matters)
221 print(" Computing waveforms...")
222 analyzer.compute('random_spikes', method='uniform', max_spikes_per_unit=500)
223 analyzer.compute('waveforms', ms_before=1.5, ms_after=2.0, **job_kwargs)
224 analyzer.compute('templates', operators=['average', 'std'])
225 analyzer.compute('noise_levels')
227 print(" Computing spike features...")
228 analyzer.compute('spike_amplitudes', **job_kwargs)
229 analyzer.compute('amplitude_scalings', **job_kwargs)
230 analyzer.compute('correlograms', window_ms=100, bin_ms=1)
231 analyzer.compute('unit_locations', method='monopolar_triangulation')
232 analyzer.compute('template_similarity')
234 print(" Computing quality metrics...")
235 analyzer.compute('quality_metrics', metric_names=METRIC_NAMES)
237 qm = analyzer.get_extension('quality_metrics').get_data()
239 return analyzer, qm
242def export_results(
243 analyzer,
244 sorting,
245 recording,
246 labels: dict,
247 output_folder: str
248):
249 """Export results to various formats."""
250 print("Exporting results...")
252 # Get good units
253 good_ids = [u for u, l in labels.items() if l == 'good']
254 sorting_good = sorting.select_units(good_ids)
255 sorting_good.save(folder=Path(output_folder) / 'sorting_curated')
256 analyzer.sorting.set_property('curation_label', [labels[uid] for uid in analyzer.sorting.unit_ids])
258 # Export to Phy
259 phy_folder = f'{output_folder}/phy_export'
260 si.export_to_phy(analyzer, phy_folder,
261 compute_pc_features=True,
262 compute_amplitudes=True, additional_properties=['curation_label'])
263 print(f" Phy export: {phy_folder}")
265 # Generate report
266 report_folder = f'{output_folder}/report'
267 si.export_report(analyzer, report_folder, format='png')
268 print(f" Report: {report_folder}")
270 # Save quality metrics
271 qm = analyzer.get_extension('quality_metrics').get_data()
272 qm.to_csv(f'{output_folder}/quality_metrics.csv')
274 # Save labels
275 with open(f'{output_folder}/unit_labels.json', 'w') as f:
276 json.dump({str(k): v for k, v in labels.items()}, f, indent=2)
278 # Save summary
279 summary = {
280 'total_units': len(sorting.unit_ids),
281 'good_units': len(good_ids),
282 'total_spikes': int(sorting.count_total_num_spikes()),
283 'duration_s': float(recording.get_total_duration()),
284 'n_channels': int(recording.get_num_channels()),
285 }
286 with open(f'{output_folder}/summary.json', 'w') as f:
287 json.dump(summary, f, indent=2)
289 print(f" Summary: {summary}")
292def run_pipeline(
293 data_path: str,
294 output_path: str,
295 sorter: str = 'kilosort4',
296 stream_name: str = 'imec0.ap',
297 apply_motion_correction: bool = True,
298 curation_method: str = 'allen',
299 apply_phase_shift: bool = True
300):
301 """Run complete Neuropixels analysis pipeline."""
303 output_path = Path(output_path)
304 output_path.mkdir(parents=True, exist_ok=True)
306 # 1. Load data
307 recording = load_recording(data_path, stream_name)
309 # 2. Preprocess
310 rec_preprocessed, bad_channels = preprocess(recording, apply_phase_shift=apply_phase_shift)
312 # Save preprocessed
313 preproc_folder = output_path / 'preprocessed'
314 job_kwargs = dict(n_jobs=8, chunk_duration='1s', progress_bar=True)
315 rec_preprocessed = rec_preprocessed.save(
316 folder=str(preproc_folder),
317 format='binary',
318 **job_kwargs
319 )
321 # 3. Check drift
322 check_drift(rec_preprocessed, str(output_path))
324 # 4. Estimate and correct motion when explicitly enabled by pipeline settings.
325 if apply_motion_correction:
326 print("Estimating and applying motion correction...")
327 rec_final = correct_motion(rec_preprocessed, str(output_path))
328 else:
329 print("Skipping external motion correction (disabled)")
330 rec_final = rec_preprocessed
332 # 5. Spike sorting
333 sorting = run_spike_sorting(
334 rec_final, str(output_path), sorter, motion_corrected=apply_motion_correction
335 )
337 # 6. Post-processing
338 analyzer, qm = postprocess(sorting, rec_final, str(output_path))
340 # 7. Curation
341 labels = curate_units(qm, method=curation_method)
343 # 8. Export
344 export_results(analyzer, sorting, rec_final, labels, str(output_path))
346 print("\n" + "="*50)
347 print("Pipeline complete!")
348 print(f"Output directory: {output_path}")
349 print("="*50)
351 return analyzer, sorting, qm, labels
354if __name__ == '__main__':
355 parser = argparse.ArgumentParser(
356 description='Neuropixels analysis pipeline (best practices)'
357 )
358 parser.add_argument('data_path', help='Path to SpikeGLX/OpenEphys recording')
359 parser.add_argument('output_path', help='Output directory')
360 parser.add_argument('--sorter', default='kilosort4',
361 choices=['kilosort4', 'kilosort3', 'spykingcircus2', 'mountainsort5'],
362 help='Spike sorter to use')
363 parser.add_argument('--stream', default='imec0.ap', help='Stream name')
364 parser.add_argument('--no-motion-correction', action='store_true',
365 help='Skip external motion correction; sorter defaults still apply')
366 parser.add_argument('--no-phase-shift', action='store_true', help='Skip ADC phase correction after verifying acquisition timing')
367 parser.add_argument('--curation', default='allen',
368 choices=['allen', 'ibl', 'strict'],
369 help='Curation method')
371 args = parser.parse_args()
373 run_pipeline(
374 args.data_path,
375 args.output_path,
376 sorter=args.sorter,
377 stream_name=args.stream,
378 apply_motion_correction=not args.no_motion_correction,
379 curation_method=args.curation,
380 apply_phase_shift=not args.no_phase_shift,
381 )