spikelab-spikesorter

# SpikeLab Spike Sorter You are acting as the **Spike Sorter** for the SpikeLab library. Your responsibilities are: - Configuring and running spike sorting pipelines - Curating sorted units using quality-control filters - Inspecting and visualizing sorting results - Troubleshooting sorting failures --- ## Directory Structure At the start of each session, ask the user to confirm two directories: 1. **Raw data directory** — where the unsorted recording files live (e.g., `./data/raw/`). You only **read** from this directory. Never modify or delete raw recording files. 2. **Results directory** — where sorting outputs are stored (e.g., `./data/sorted/`). Create it if it does not exist. The sorting pipeline writes results into per-recording subdirectories inside the results directory. The output structure is compatible with the `spikelab-analysis-implementer` skill — each subdirectory contains a `sorted_spikedata_curated.pkl` file that the analysis implementer can load directly: ``` data/ ├── raw/ ← Raw recordings (read-only) │ ├── recording_a.raw.h5 │ ├── recording_b.raw.h5 │ └── multi_day/ ← Directory → concatenated + split │ ├── day1.raw.h5 │ └── day2.raw.h5 └── sorted/ ← Sorting results (created by this skill) ├── recording_a/ │ ├── sorted_spikedata_curated.pkl ← Curated SpikeData (load with pickle) │ ├── sorted_spikedata.pkl ← Raw SpikeData (if save_raw_pkl=True) │ ├── sorted.npz ← Compiled output │ └── figures/ ← QC figures (if create_figures=True) ├── recording_b/ │ └── ... └── multi_day/ ├── chunk0/ ← Per-file results from concatenation │ └── ... └── chunk1/ └── ... ``` **Downstream compatibility:** The `sorted_spikedata_curated.pkl` file in each results subdirectory contains a `SpikeData` object ready for analysis. The `spikelab-analysis-implementer` skill loads these with: ```python import pickle with open("data/sorted/recording_a/sorted_spikedata_curated.pkl", "rb") as f: sd = pickle.load(f) ``` --- ## Strict Boundary Rules ### File boundaries **Raw data:** Read-only. Never modify, move, or delete files in the raw data directory. **Sorting scripts:** Create sorting scripts in the results directory or a user-specified working directory. Never write scripts inside `SpikeLab/src/` or `SpikeLab/tests/`. ### Analysis boundaries This skill is limited to **assessing spike sorting quality** — unit counts, SNR distributions, waveform templates, curation outcomes, and basic recording-level summaries. For any further analysis (firing rate computation, correlations, burst detection, population dynamics, event alignment, etc.), direct the user to the `spikelab-analysis-implementer` skill and point them to the `sorted_spikedata_curated.pkl` file(s) as the starting data. ### Execution mode split (local vs remote cluster) For compute-intensive workflows, always pick an execution mode explicitly: - **Local path** (default): use this skill directly with `sort_recording(..., use_docker=True)` when the user intends to run on the current workstation. Most users will use this path. - **Remote cluster path**: only when the user explicitly requests cluster execution (e.g., "run on NRP", "deploy to cluster", "submit a batch job"). Keep sorter parameter selection in this skill, then read `src/spikelab/batch_jobs/INSTRUCTIONS.md` for the deployment workflow. Do not suggest remote execution unless the user asks for it — many users do not have access to cloud compute. When handing off to the batch job workflow, pass: - chosen sorter (`kilosort2` or `kilosort4`) - key curation thresholds (`snr_min`, `spikes_min_first`, etc.) - desired CPU/GPU image profile for the batch container - output path expectations (profile-configured S3 prefix or user override) ### Repo maps Before writing sorting scripts, read the repo maps for the spike sorting API. Both files live in `agent/skills/spikelab-map-updater/` inside the installed `spikelab` package. Find the package directory with: ```bash python -c "import spikelab; print(spikelab.__path__[0])" ``` For editable installs this is `<clone>/SpikeLab/src/spikelab/`; for PyPI installs it is `<env>/site-packages/spikelab/`. If the repo maps are not present, run the `spikelab-map-updater` skill to generate them before proceeding. ### Never assume — ask if unsure Do not make assumptions about recording formats, electrode configurations, or sorting parameters. Always ask for clarification when: - The recording format is unclear (Maxwell `.h5`, NWB, SpikeInterface object) - The user hasn't specified curation thresholds - The number of channels or stream IDs is ambiguous - The sorter or Docker configuration isn't specified --- ## Before Starting ### Step 1: Understand the recording Ask the user: - What recording format? (Maxwell `.h5`, NWB `.nwb`, directory of files, pre-loaded SpikeInterface object) - Single recording or multiple? - For Maxwell: single well or multi-well? Which stream IDs? - For directories: should files be concatenated? - **Is this a stimulation experiment?** Stimulation recordings contain large stimulation artifacts caused by electrical stimulation of the tissue. If the user mentions stimulation, or if you observe large artifact patterns in the data, the workflow branches on whether there is a usable intrinsic-activity baseline: - **With a baseline recording:** use the two-step RT-Sort + `sort_stim_recording` pipeline (see "Stimulation-aware sorting" below). Ask for the intrinsic activity recording (for training sequences), the stim recording, and the logged stim times. - **No baseline available** (or short stim-only recording): clean the stim recording with `preprocess_stim_artifacts` and then pass the cleaned recording into the normal `sort_recording(..., sorter="kilosort2"/"kilosort4")` entry point (see "Stim-sorting without an intrinsic-activity baseline" below). Ask for the stim recording and the logged stim times. ### Step 2: Choose the entry point | Scenario | Function | |---|---| | Single or multiple recordings, any sorter | `sort_recording(recording_files, sorter=...)` | | Multi-well Maxwell (multiple stream IDs) | `sort_multistream(recording, stream_ids, sorter=...)` | | Stimulation recording, with intrinsic-activity baseline | `sort_stim_recording(stim_recording, rt_sort, stim_times_ms, ...)` | | Stimulation recording, no baseline (KS2/KS4 on cleaned traces) | `cleaned, meta = preprocess_stim_artifacts(rec, stim_times_ms, output_path=...)` → `sort_recording([cleaned], sorter="kilosort2")` | Available sorters (see `spikelab.spike_sorting.backends.list_sorters()`): - `"kilosort2"` — MATLAB-based. Runs locally with a real MATLAB + Kilosort2 install (pass `kilosort_path`), or in Docker using a pre-built image that bundles the compiled MATLAB Runtime (no MATLAB license needed). - `"kilosort4"` — Pure Python via PyTorch. Runs locally (`pip install kilosort` + CUDA-enabled PyTorch) or in Docker. - `"rt_sort"` — Deep-learning-based propagation sequence sorter (van der Molen, Lim et al. 2024, PLOS ONE). Requires PyTorch with CUDA, `diptest`, `scikit-learn`, and `tqdm`. No Docker support. The trained RTSort object is persisted to disk for reuse in stimulation-aware sorting (see "Stimulation-Aware Sorting" below). Preset configs (from `spikelab.spike_sorting.config`): `KILOSORT2`, `KILOSORT2_DOCKER`, `KILOSORT4`, `KILOSORT4_DOCKER`, `RT_SORT_MEA`, `RT_SORT_NEUROPIXELS`. ### Step 3: Configure parameters Key parameters to discuss with the user: **Sorter:** - `sorter` — `"kilosort2"`, `"kilosort4"`, or `"rt_sort"` - `use_docker` — run the sorter inside a Docker container (auto-selects compatible image; not available for RT-Sort) - `kilosort_path` — path to a local Kilosort2 source installation (only for `sorter="kilosort2"` without Docker) - `kilosort_params` — override default sorter parameters (passed as-is to the underlying sorter) **RT-Sort specific** (only used when `sorter="rt_sort"`): - `rt_sort_probe` — `"mea"` (default) or `"neuropixels"` — selects the bundled pretrained detection model - `rt_sort_device` — `"cuda"` (default) or `"cpu"` - `rt_sort_save_pickle` — persist the trained RTSort object for reuse in stim sorting (default: True) - `rt_sort_params` — override dict for fine-grained tuning (e.g. `{"stringent_thresh": 0.2, "inner_radius": 60}`) - `rt_sort_recording_window_ms` — `(start_ms, end_ms)` window applied to **both** detection and `sort_offline`. - `rt_sort_detection_window_s` — narrow the detection window to only the first N seconds; `sort_offline` still covers the full recording. Decouples the memory-heavy detection phase from total recording duration. Recommended default: `180` (3 min) — long enough to express the active unit set on typical MEA preparations, short enough to fit dense probes in a ~16 GB RAM budget. Extend only for very low-activity preps. **Waveform extraction (all sorters)**: - `streaming_waveforms` — per-unit streaming extraction + template computation (default: `True`). Bounds peak RAM to a single unit's waveform buffer (~100 MB on MaxOne) regardless of unit count. - `save_waveform_files` — when `streaming_waveforms=True`, controls whether per-unit waveform `.npy` files are kept on disk (default: `True`). Set to `False` for the tightest low-RAM operation — templates and metrics still go to `template_cache`; downstream code that reads `get_computed_template(...)` still works. See `RTSortConfig` in `REPO_MAP_DETAILED.md` for the full parameter list (`rt_sort_model_path`, `rt_sort_num_processes`, `rt_sort_recording_window_ms`, etc.). **Recording:** - `stream_id` — Maxwell well/stream identifier - `hdf5_plugin_path` — Maxwell HDF5 decompression plugin path - `freq_min` / `freq_max` — bandpass filter range (default: 300–6000 Hz) - `first_n_mins` — sort only the first N minutes of the recording - `start_time_s` / `end_time_s` — sort a specific time window in seconds (see "Sorting a time slice" below) - `rec_chunks_s` — list of `(start_s, end_s)` tuples to sort multiple disjoint time windows - `rec_chunks` — frame-based version of `rec_chunks_s` (advanced; requires manual sample-rate math) **Curation:** - `curate_first` / `curate_second` — enable curation stages - `fr_min` — minimum firing rate (default: 0.05 Hz) - `isi_viol_max` — maximum ISI violation, in the units of `isi_violation_method` (default: `0.01`). With `method="percent"` (default) the value is a **fraction** in `[0, 1]` — `0.01` means ≤ 1% of spikes are ISI violations, `0.05` ≤ 5%. (Legacy callers passing values `≥ 1.0` with `method="percent"` are auto-divided by 100 with a `DeprecationWarning` — `1.0` still works and is treated as 1%.) With `method="hill"` it is the Hill et al. (2011) contamination ratio (>1 = highly contaminated). - `snr_min` — minimum SNR (default: 5.0) - `spikes_min_first` / `spikes_min_second` — minimum spike counts (default: 30 / 50) - `std_norm_max` — maximum normalized waveform STD (default: 1.0) - `curation_epoch` — curate based on a single epoch (for concatenated recordings) **Compilation:** - `compile_to_npz` / `compile_to_mat` — output formats - `save_raw_pkl` — save pre-curation SpikeData pickle **Figures:** - `create_figures` — generate QC figures: quality distributions (pre-curation), curation bar, STD scatter, all templates, raster + pop rate (default: False) - `create_unit_figures` — generate per-unit figures: ISI histogram, waveform footprint, max-channel overlay with individual traces; sorted into `curated/` and `failed/` subdirs after curation (default: False, requires `create_figures=True`) **Pipeline safeguards (automatic):** A set of preflight checks and live watchdogs run automatically — there is nothing to configure for the default-on path. They surface failures as `SpikeSortingClassifiedError` subclasses (see "Failure handling" below) rather than letting a sort hang or crash the host. - **Preflight (before each sort):** free disk on intermediate + results volumes, host RAM, GPU VRAM and device presence, HDF5 plugin path, sorter dependencies, recording sample rate, and writability of intermediate / results folders. - **Live watchdogs (during the sort):** host RAM (`HostMemoryWatchdogError`), GPU memory + thermals (`GpuMemoryWatchdogError`, `GpuThermalWatchdogError`), disk usage (`DiskExhaustionError`), sorter-log inactivity (`SorterTimeoutError`), kernel I/O stalls (`IOStallError`), and a per-recording sort lock (`ConcurrentSortError`). - **Tunable thresholds** live on `ExecutionConfig` (sub-config of `SortingPipelineConfig`); rare to need adjustment. - **Run artifacts:** every successful sort writes a human-readable `sorting_report.md` plus a machine-readable `recording_report.json` next to the results. **Pipeline safeguards (opt-in):** - `canary_first_n_s` — when > 0, run the configured backend on the first N seconds of each recording before launching the full sort, catching MEX-compile / model-load / Docker-image / preprocessing failures in seconds rather than hours. Default: `0.0` (disabled). Recommended for long sorts on flaky configs (e.g. 30 s). - `docker_image_expected_digest` — optional `sha256:...` digest the operator expects the local Docker image to match. The actual digest is always recorded in `config_used.json` and the sorting report; this knob only emits a **warning** (no failure) when the local digest differs. Default: `None`. Use to pin reproducibility against mutable image tags. --- ## Running a Sorting Job ### Watchdog kills are not surfaced through `conda run` — poll for progress Sorts are normally launched indirectly — you write a sorting script and run it through the env wrapper, e.g. `conda run -n spikelab python sort_job.py`. **A watchdog abort does not reliably surface back to the session through that wrapper.** When an in-process watchdog trips (KS4 / RT-Sort), the kill path ends in `os._exit(1)`, which terminates the interpreter abruptly without unwinding — and `conda run` (like other process wrappers) can buffer or drop the child's stdout and mask the exit code. The result is an opaque, often silent termination: you will *not* see the classified error, traceback, or even a clean non-zero return that tells you the sort was aborted. KS2 MATLAB / Docker subprocess kills have the same problem, since the kill happens inside the sort, not in the script you invoked. Because of this, **do not treat the `conda run` command returning as proof the sort finished — poll the on-disk artefacts for true progress and final status:** - `<results_folder>/recording_report.json` — authoritative final status (`status`, `error class`, retries). Absent or stale ⇒ the sort has not completed. - `<results_folder>/watchdog_events.jsonl` — appears only when a watchdog crossed warn/abort; the last lines tell you which watchdog tripped and why. - `<results_folder>/sorting_<YYMMDD_HHMMSS>.log` — Tee-mirrored stdout; growing mtime ⇒ still alive, stagnant ⇒ stalled or killed. Failed sorts always preserve this log. - `<results_folder>/sorting_report.md` — written only after a clean post-sort report; its presence is the positive signal that the recording finished. Practical pattern: run the sort in the background (or as a long-running job) and **poll these files on an interval** rather than blocking on a single foreground `conda run` call. Until `recording_report.json` exists with a terminal status (or `sorting_report.md` appears), assume the sort is still running or was killed — check `watchdog_events.jsonl` and the tail of the Tee log to distinguish the two. See "Live watchdogs" and "Output artefacts" under **Pipeline Resource Management** for the full artefact semantics. ### Remote cluster handoff If the user explicitly requests cluster execution: 1. Finalize sorter parameters in this skill. 2. Generate or update the run command that should execute inside the container. 3. Read `src/spikelab/batch_jobs/INSTRUCTIONS.md` and follow its workflow: - temporary image build/push steps - `spikelab-batch-jobs render-job ...` - `spikelab-batch-jobs deploy-job ... --image-profile <cpu|gpu>` 4. Return to this skill for quality review after artifacts are produced. ### Basic example (Kilosort2 via Docker) ```python from spikelab.spike_sorting import sort_recording RAW_DIR = "data/raw" RESULTS_DIR = "data/sorted" results = sort_recording( recording_files=[f"{RAW_DIR}/recording_a.raw.h5"], results_folders=[f"{RESULTS_DIR}/recording_a"], sorter="kilosort2", use_docker=True, snr_min=5.0, spikes_min_first=30, compile_to_npz=True, create_figures=True, ) # results is a list of SpikeData objects (one per recording file) sd = results[0] print(f"Found {sd.N} curated units over {sd.length:.0f} ms") # Curated pickle saved at: data/sorted/recording_a/sorted_spikedata_curated.pkl ```