feat: DREDge motion correction (CatGT→DREDge→KS4) + uv/py3.14 env

pyproject.toml

@@ -7,23 +7,32 @@ name = "u19-sorting"
 version = "0.1.1"
 description = "BRAINCoGS scripts to process electrophysiology data"
 readme = "README.md"
-license = ""
+license-files = [
+    "LICENSE*"
+    ]
 authors = [
     { name = "BRAINCoGS"},
     { name = "Alvaro Luna", email = "[email protected]" },
     { name = "Christian Tabedzki", email = "[email protected]" },
 ]
 dependencies = [
-    "ONE-api",
-    "ibllib",
+    "ONE-api>=3.4.1",
+    "ibllib>=4.0.1",
     "kilosort[gui]",
     "numpy",
     "pandas",
+    "spikeinterface>=0.104.8",
 ]
-requires-python = ">= 3.10"
+requires-python = ">= 3.14"
 
 
 [tool.hatch.build.targets.sdist]
 include = [
     "/u19_sorting",
 ]
+
+# Point Pyright / basedpyright (and Pylance) at the uv-managed virtualenv so imports of
+# spikeinterface, torch, kilosort, etc. resolve instead of being flagged as missing.
+[tool.pyright]
+venvPath = "."
+venv = ".venv"

requirements.txt

@@ -3,3 +3,4 @@ ibllib
 ONE-api
 pandas
 numpy
+spikeinterface[full]>=0.101

u19_sorting/config.py

@@ -54,8 +54,10 @@
 
 preproc_tools = {
     'catgt':      'catgt',
+    'dredge':     'dredge',
 }
 
 preproc_tools_delete_post = [
-    'catgt'
+    'catgt',
+    'dredge',
 ]

u19_sorting/matlab_scripts/kilosortbatch.m

@@ -83,7 +83,14 @@ function kilosortbatch(parameter_file, raw_directory, processed_directory, chann
     ops.fbinary = fullfile(raw_directory, fs(1).name);
 
     rez                = preprocessDataSub(ops);
-    rez                = datashift2(rez, 1);
+    % Drift correction. ops.nblocks controls registration: 0 turns it off (e.g. when DREDge
+    % already corrected motion in preprocessing), >0 enables datashift. datashift2's second arg
+    % is the do-correction flag, so nblocks==0 skips it.
+    if isfield(ops, 'nblocks') && ops.nblocks == 0
+        rez            = datashift2(rez, 0);
+    else
+        rez            = datashift2(rez, 1);
+    end
 
     [rez, st3, tF]     = extract_spikes(rez);
 

u19_sorting/preprocess_wrappers.py

@@ -29,9 +29,30 @@ def preprocess_main(recording_process_id, raw_data_directory, processed_data_dir
             catgt_output_dir = pathlib.Path(processed_data_directory, config.preproc_tools['catgt']+"_output")
             #pathlib.Path(catgt_output_dir).mkdir(parents=True, exist_ok=True)
             new_raw_data_directory = cat_gt.run_cat_gt(new_raw_data_directory, catgt_output_dir, this_preparam[config.preproc_tools['catgt']])
+        if config.preproc_tools['dredge'] in this_preparam:
+            dredge_output_dir = pathlib.Path(processed_data_directory, config.preproc_tools['dredge']+"_output")
+            new_raw_data_directory = dredge.run_dredge(new_raw_data_directory, dredge_output_dir, this_preparam[config.preproc_tools['dredge']])
 
     return new_raw_data_directory
 
+
+def preprocess_has_tool(recording_process_id, tool_key):
+    """ Return True if the given preproc tool (e.g. 'dredge') is part of this job's preprocess params.
+
+        Used by the sorter stage to decide whether to disable Kilosort's internal drift
+        correction (so motion is not corrected twice).
+    """
+
+    preprocess_parameter_filename = config.preprocess_parameter_file.format(recording_process_id)
+    if not pathlib.Path(preprocess_parameter_filename).is_file():
+        return False
+
+    with open(preprocess_parameter_filename, 'r') as preprocess_param_file:
+        preprocess_parameters = json.load(preprocess_param_file)
+
+    tool_name = config.preproc_tools[tool_key]
+    return any(tool_name in this_preparam for this_preparam in preprocess_parameters)
+
 def post_process_partial_results(recording_process_id, raw_data_directory, processed_data_directory):
 
     # Delete all unnecesary preprocessing tools results (to save storage)
@@ -192,4 +213,126 @@ def cat_gt_check_output(cat_gt_output_dir):
         if patterns_found:
             return 1
         else:
-            return 0
+            return 0
+
+
+class dredge():
+    """ DREDge motion / drift correction via SpikeInterface.
+
+        Runs after CatGT and before Kilosort. Reads the SpikeGLX ap.bin/ap.meta pair
+        produced by CatGT, estimates and applies motion correction with the SpikeInterface
+        `dredge` preset, and writes a corrected SpikeGLX-style ap.bin (+ copied ap.meta)
+        into `dredge_output`, so the sorter stage consumes it exactly like a CatGT output.
+
+        KS4's own drift corrector must be disabled downstream (nblocks=0) so correction is
+        applied once, here.
+    """
+
+    @staticmethod
+    def run_dredge(raw_data_directory, dredge_output_dir, dredge_params):
+        """ Estimate + apply motion correction and write a corrected ap.bin.
+
+            Args:
+                raw_data_directory  (Path): dir with the (CatGT) ap.bin/ap.meta to correct
+                dredge_output_dir   (Path): dir where the corrected ap.bin/ap.meta are written
+                dredge_params       (dict): {preset, device, motion_kwargs, job_kwargs}
+        """
+
+        dredge_output_dir = pathlib.Path(dredge_output_dir)
+
+        # Lazy-skip on restart / requeue if we already produced a corrected recording.
+        if dredge.dredge_check_output(dredge_output_dir):
+            print('dredge output already present, skipping', dredge_output_dir)
+            return dredge_output_dir
+
+        # Import SpikeInterface / torch lazily so catgt-only and KS2/KS3 runs never pay for it.
+        import shutil
+        import spikeinterface.full as si
+        from spikeinterface.preprocessing.motion import correct_motion
+        import torch
+
+        raw_data_directory = pathlib.Path(raw_data_directory)
+        dredge_output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Locate the SpikeGLX ap.meta produced by CatGT (flattened into the dir root).
+        meta_files = sorted(raw_data_directory.glob('*ap.meta'))
+        if not meta_files:
+            raise ValueError('No *ap.meta found in ' + raw_data_directory.as_posix())
+        src_meta = meta_files[0]
+        stem = src_meta.name[:-len('.ap.meta')]
+
+        # Probe index for the SpikeGLX stream id (imec<N>.ap), reusing catgt's regex idiom.
+        probe_match = re.search(r"imec([0-9]+)", src_meta.name)
+        stream_id = ("imec" + probe_match.group(1) + ".ap") if probe_match else "imec0.ap"
+
+        rec = si.read_spikeglx(folder_path=raw_data_directory.as_posix(), stream_id=stream_id)
+
+        # GPU for motion estimation unless overridden. DREDge does not auto-detect a GPU
+        # (unlike KS4), so we pass the device explicitly.
+        device = dredge_params.get('device')
+        if device is None or device == 'cuda':
+            if torch.cuda.is_available():
+                device = 'cuda'
+                print('dredge using CUDA device:', torch.cuda.get_device_name())
+            else:
+                device = 'cpu'
+                print('dredge: CUDA not available, falling back to CPU')
+
+        job_kwargs = dredge_params.get('job_kwargs', {})
+        motion_kwargs = dredge_params.get('motion_kwargs', {})
+
+        # output_motion_info=True (and output_motion=False) => returns (recording, motion_info).
+        motion_result = correct_motion(
+            rec,
+            preset=dredge_params.get('preset', 'dredge'),
+            folder=(dredge_output_dir / 'motion').as_posix(),
+            output_motion_info=True,
+            overwrite=True,
+            estimate_motion_kwargs={'device': device},
+            **motion_kwargs,
+            **job_kwargs,
+        )
+        rec_corr, motion_info = motion_result  # type: ignore[misc]  # union return; runtime is a 2-tuple
+
+        # Materialize a SpikeGLX-style ap.bin. write_binary_recording (not save()) writes a
+        # single flat binary; KS4's find_binary globs *.bin and prefers the 'ap.bin' tag.
+        corrected_bin = dredge_output_dir / (stem + '.ap.bin')
+        si.write_binary_recording(
+            rec_corr,
+            file_paths=[corrected_bin.as_posix()],
+            dtype='int16',
+            **job_kwargs,
+        )
+
+        # Copy the ap.meta alongside (provenance + lets read_spikeglx re-open on requeue).
+        # Motion interpolation preserves sample count, so the meta stays valid.
+        shutil.copy2(src_meta.as_posix(), (dredge_output_dir / (stem + '.ap.meta')).as_posix())
+
+        # Release GPU memory so the downstream KS4 stage gets the full card.
+        del rec_corr, rec, motion_info
+        if device == 'cuda':
+            torch.cuda.empty_cache()
+
+        return dredge_output_dir
+
+    @staticmethod
+    def dredge_check_output(dredge_output_dir):
+
+        file_patterns = ['/*ap.bin', '/*ap.meta']
+
+        child_dirs = [x[0] for x in os.walk(dredge_output_dir)]
+        patterns_found = 0
+        for dir in child_dirs:
+            for pat in file_patterns:
+                found_file = glob.glob(dir+pat)
+                if len(found_file) > 0:
+                    patterns_found = 1
+                    break
+
+            if patterns_found:
+                break
+
+        if patterns_found:
+            return 1
+        else:
+            return 0

u19_sorting/sorter_wrappers.py

@@ -5,6 +5,7 @@
 import subprocess
 import json
 import u19_sorting.config as config
+import u19_sorting.preprocess_wrappers as pw
 from u19_sorting.utils import write_file
 
 
@@ -28,6 +29,11 @@ def sorter_main(recording_process_id, raw_directory, processed_directory):
 
     sorter = config.sorters_names[process_parameters['clustering_method']]
 
+    # If DREDge ran as a preprocessing step, motion is already corrected. Disable Kilosort's
+    # own internal drift correction so motion is not corrected twice.
+    if pw.preprocess_has_tool(recording_process_id, 'dredge'):
+        process_parameters = disable_internal_drift(process_parameters, sorter)
+
 
     sorter_processed_directory = pathlib.Path(processed_directory, process_parameters['clustering_method']+'_output')
     pathlib.Path(sorter_processed_directory).mkdir(parents=True, exist_ok=True)
@@ -55,6 +61,29 @@ def sorter_main(recording_process_id, raw_directory, processed_directory):
     return sorter_processed_directory
 
 
+def disable_internal_drift(process_parameters, sorter):
+    """ Turn off a Kilosort sorter's built-in motion/drift correction.
+
+        Used when DREDge has already corrected motion in preprocessing. The mechanism differs
+        per Kilosort version:
+          - KS4 (python): nblocks=0 and do_correction=False in the settings dict.
+          - KS3 (matlab kilosortbatch): ops.nblocks=0 (kilosortbatch.m honors this to skip
+            datashift2).
+          - KS2 (matlab run_ks2): ops.reorder=0 (KS2 uses batch reordering, not nblocks).
+    """
+
+    if sorter == config.sorters_names['kilosort4']:
+        process_parameters['nblocks'] = 0
+        process_parameters['do_correction'] = False
+    elif sorter == config.sorters_names['kilosort3']:
+        process_parameters['nblocks'] = 0
+    elif sorter == config.sorters_names['kilosort2']:
+        process_parameters['reorder'] = 0
+
+    print('DREDge preprocessing detected: disabled internal drift correction for', sorter)
+    return process_parameters
+
+
 class Kilosort2():
     """ Kilosort2 caller functions """