Integrate SpikeInterface (DREDge motion correction) + Kilosort version selection via Apptainer

[tabedzki]

Summary

A user has requested the ability to use the DREDge motion-correction method (now integrated into SpikeInterface) instead of the drift correction that is currently baked into Kilosort. This issue scopes how SpikeInterface (SI) could be integrated into the pipeline, the two viable architectures, and the related question of letting users choose their Kilosort version via Apptainer/Singularity.

Background: how sorting works today

• U19-pipeline_python (DataJoint orchestration) generates per-job parameter JSON (preprocess_paramset_{id}.json, process_paramset_{id}.json) + chanmap_{id}.mat, writes a SLURM script, and scps everything to the cluster. It does no sorting itself.
• BrainCogsEphysSorters (this repo, compute layer) runs on the cluster via main_script.py in three stages:
  • Preprocess (preprocess_wrappers.py) → currently only CatGT.
  • Sort (sorter_wrappers.py) → dispatches on clustering_method: KS2/KS3 via MATLAB (subprocess), KS4 via Python (kilosort.run_kilosort).
  • Postprocess (postprocess_wrappers.py) → IBL atlas alignment.

Key finding: the "default Kilosort motion corrector" the user wants to replace is internal to Kilosort (KS2/3 ops.nblocks datashift, KS4 built-in registration). There is no separate motion-correction stage today — it is a sorter parameter, not a pluggable step. A 'SpikeInterface':'spike_interface' entry is already stubbed in config.sorters_names, but there is no implementation, no SI submodule, and SI is not in the dependencies.

DREDge is exposed in SI as a preprocessing step (spikeinterface.preprocessing.correct_motion(preset="dredge_ap")) that produces a motion-corrected recording fed to any sorter. This matches SI's model: recording → preprocess → motion correction → sorter → postprocess.

Option A — DREDge as a new preprocess step (recommended first step)

Add a dredge branch to the existing per-step loop in preprocess_wrappers.preprocess_main, mirroring CatGT. It reads the SpikeGLX run, runs correct_motion(preset="dredge_ap"), and writes a corrected binary that becomes the new_raw_data_directory handed to the existing Kilosort caller. Disable KS internal drift (nblocks=0) so motion isn't corrected twice.

• Touches: preprocess_wrappers.py (~40 lines), config.py (preproc_tools['dredge'] + add to preproc_tools_delete_post), the sorter call (nblocks=0), the preprocess JSON schema, and the DataJoint param generator.
• Untouched: sorter wrappers, IBL postprocessing, phy params.py shims, SLURM/transfer machinery.
• Cost: one new dependency (SI + torch); one extra full-size corrected copy on disk per probe (added to existing post-run cleanup).
• Difficulty: contained, ~1–2 days + cluster env validation.
• Does NOT require any container work. KS4 stays native exactly as today.

Option B — replace the compute layer with SpikeInterface end-to-end

SI owns the whole compute stage; main_script.py's three stages collapse into one SI script:
read_spikeglx → highpass/phase_shift/common_reference (replaces CatGT) → correct_motion(dredge_ap) → run_sorter(...) → quality metrics → export_to_phy → IBL atlas.

• Wins: drops the MATLAB dependency for KS2/3 (via containers, below); adds quality metrics; sorter-agnostic; one code path for any SI-supported sorter; unlocks per-version sorter selection via Apptainer.
• Main migration cost: SI's output layout differs from raw-KS, so the IBL atlas postprocessing must be re-pointed at the export_to_phy output. The phy params.py/win_params.py/mac_params.py cross-OS shim is replaced by SI's export (reapply the cup-path rewrite afterward).
• Difficulty: substantial but bounded, ~1–2 weeks. Risk is preprocessing-equivalence with CatGT + IBL rewiring, not the SI script.

Recommendation: do Option A first (delivers DREDge in days, de-risks the SI/torch/GPU env), then migrate to Option B as the deliberate architecture move. A is a strict subset of B's preprocessing, so it is not throwaway work.

Kilosort version selection (Apptainer/Singularity)

Two independent axes — do not conflate them:

1. MATLAB sorters (KS2/2.5/3) vs Python (KS4) — decides whether a container is needed at all. MATLAB ones need it (to escape the license + module load matlab); KS4 does not.
2. Which release of a given sorter (e.g. KS4 4.0.16 vs 4.1.0) — a version-pinning problem that exists for KS4 too. "Run KS4 native in the shared env" is unsafe: everyone gets whatever single version the env happens to have, and an admin upgrading it silently changes everyone's results.

So pin versions for both, but the mechanism differs:

• KS4 (Python): per-version immutable envs (frozen lockfile), selected by an allow-listed key. (Optionally containerize KS4 too for a single uniform mechanism.)
• KS2/2.5/3 (MATLAB): pre-pulled, pinned .sif images on shared storage — the only way to get version pinning and MATLAB-free execution.

SI's run_sorter(singularity_image=...) accepts a local .sif path, runs the sorter inside the container, reads results back — host only needs SI + Apptainer. SI maintains official per-version images (spikeinterface/kilosort{4,3,2_5}-compiled-base:<ver>); the MATLAB ones are MATLAB-compiled (no license needed).

Offline image pulls (compute nodes have no internet)

Never let run_sorter pull at runtime. Two-phase:

1. Pull once on the login node (has internet) into shared fast storage:

   export APPTAINER_CACHEDIR=/scratch/gpfs/BRAINCOGS/apptainer_cache
   apptainer pull /scratch/gpfs/BRAINCOGS/sif/kilosort3_0.1.0.sif \
       docker://spikeinterface/kilosort3-compiled-base:0.1.0

2. Jobs reference the local .sif by absolute path, never a docker:// URL — no network at run time.

Safety invariants (enforced in the DataJoint layer)

1. Closed allow-list, not free text. User picks a registry key; param generator validates against it and refuses unknown keys (no arbitrary env/path injection).
2. Resolved version stamped into job provenance (process JSON + ClusteringParamSet).
3. Registry is append-only / admin-gated. New version = admin builds env or pulls .sif on login node + adds one registry line. Existing entries never mutated.
4. Resolution happens once, server-side, before queueing. Compute node makes no version decisions and no network calls.

SORTER_REGISTRY = {
  "kilosort4@4.0.16": {"mode": "native",    "env": "/scratch/.../envs/ks4-4.0.16"},
  "kilosort4@4.1.0":  {"mode": "native",    "env": "/scratch/.../envs/ks4-4.1.0"},
  "kilosort3@0.1.0":  {"mode": "container", "sif": "/scratch/.../sif/ks3_0.1.0.sif"},
  "kilosort2_5@0.1.0":{"mode": "container", "sif": "/scratch/.../sif/ks2_5_0.1.0.sif"},
}

Environment manager: pixi vs uv

The compute repo currently locks with uv.lock (31 nvidia-*/cuda packages, torch, numba/llvmlite, kilosort). uv resolves PyPI wheels only, so the cluster's CUDA is an unlocked external assumption and conda-only SI sorters (e.g. mountainsort5) have no clean story.

pixi is conda-forge-based (and embeds uv for the PyPI side), giving: cuda-version as a first-class lock dimension; system/non-Python deps in the same lock; native multi-environment in one manifest (ks4-4.0.16, ks4-4.1.0, …) — which maps directly onto the per-version SORTER_REGISTRY.

Recommendation: use pixi for this compute repo (locked CUDA + native multi-env = the version registry); keep uv for U19-pipeline_python (pure-Python orchestration, no CUDA/multi-env). Caveat: pixi pins conda-forge torch/CUDA builds, not the PyPI torch wheels. Apptainer images sidestep this for containerized sorters either way.

[project]
channels = ["conda-forge"]
platforms = ["linux-64", "osx-arm64"]
[system-requirements]
cuda = "12.4"
[dependencies]
spikeinterface = "*"
python = ">=3.10"
[feature.ks4-4016.pypi-dependencies]
kilosort = "==4.0.16"
[feature.ks4-4100.pypi-dependencies]
kilosort = "==4.1.0"
[environments]
ks4-4016 = ["ks4-4016"]
ks4-4100 = ["ks4-4100"]

Job activates pixi run -e ks4-4016 python main_script.py, resolved from the allow-listed registry key.

Proposed plan

• Prototype offline: one session, read_spikeglx → correct_motion(dredge_ap) → run_kilosort4 on a tiger/spock GPU node. De-risks env/CUDA/disk.
• Verify cluster prerequisites: module avail cuda versions + GPU driver max CUDA; confirm Apptainer present on compute nodes with --nv GPU passthrough.
• Implement Option A: dredge preprocess step + JSON schema; add SI dep; add output dir to post-run cleanup; set KS nblocks=0 when active.
• Thread param through DataJoint (parameter_file_creator.py + preprocess-param tables).
• (Later) Option B: full SI compute script + re-point IBL postprocessing at export_to_phy.
• (Later) Version registry: SORTER_REGISTRY resolver + login-node pull/setup script + provenance stamping.
• (Decision) Adopt pixi for the compute repo with per-version environments.

Open questions

1. What CUDA versions do tiger/spock expose, and are GPU drivers new enough for CUDA 12.x?
2. Native vs containerized KS4: (A) KS4 native via pixi per-version envs + MATLAB sorters via Apptainer, or (B) everything via Apptainer including KS4?
3. Who maintains versions (admin-gated "add a version" step vs self-service)?

🤖 Generated with Claude Code

[tabedzki]

We should also consider doing a version where we load SI via this wrapper and do everything as part of 'processing' step.