Run dartsort's template matching from external Kilosort 4 sorting

matching_from_ks4.py

import numpy as np
import dartsort
from pathlib import Path

# -- helpers

def load_ks4_sorting(recording, kilosort_dir):
    times = np.load(kilosort_dir / 'spike_times.npy')
    labels = np.load(kilosort_dir / 'spike_clusters.npy')
    pos = np.load(kilosort_dir / 'spike_positions.npy')
    amp = np.load(kilosort_dir / 'amplitudes.npy')

    # load the sorting, but use 0s for channels because KS
    # doesn't save them. we'll back them out in the next step.
    ks_st = dartsort.DARTsortSorting(
        times_samples=times,
        labels=labels,
        channels=np.zeros_like(labels),
        ephemeral_features=dict(amplitudes=amp)
    )

    # make quick templates (not suitable for matching)
    templates = dartsort.TemplateData.from_config(
        recording=recording,
        sorting=ks_st,
        template_cfg=dartsort.raw_template_cfg,
    )
    # set spike channels from template main channels
    main_channels = np.ptp(templates.templates, axis=1).argmax(1)
    # final spike train with channel info
    ks_st = dartsort.DARTsortSorting(
        times_samples=times,
        labels=labels,
        channels=main_channels[labels],
        ephemeral_features=dict(
            positions=pos,
            amplitudes=amp,
            times_seconds=recording.sample_index_to_time(times),
        ),
    )
    return ks_st
 
# -- main

# -- paths
# this should be the kilosort4/ folder with the .npy files inside
kilosort_dir = Path("...")

# output data will be saved to this folder
output_dir = Path("...")

# this recording should have time info which matches the recording
# that kilosort was run on
recording_sorted = ...

# this recording is the target of matching
recording_match = ...

# amount of realignment to do
realign_ms = 1.5

# -- load and estimate templates
# load ks outputs into a dartsort object
ks_st = load_ks4_sorting(recording_sorted, kilosort_dir)

# trough-align the spike train and estimate low-rank templates for matching
# they will be saved to template_data.npz in the output dir, and this code
# block will re-load them if that file already exists
ks_st_realigned, template_data = dartsort.estimate_template_library(
    recording=recording_sorted,
    sorting=ks_st,
    realign_cfg=dartsort.TemplateRealignmentConfig(
        realign_shift_ms=realign_ms,
    ),
    template_cfg=dartsort.default_template_cfg,
    template_npz_path=output_dir / 'template_data.npz',
)

# optional: save realigned spike times
if not (output_dir / 'ks_st_realigned.npz').exists():
    ks_st_realigned.save(output_dir / 'ks_st_realigned.npz')
else:
    ks_st_realigned = dartsort.load(output_dir / 'ks_st_realigned.npz')

# run the matcher
match_st = dartsort.match(
    recording=recording_match,
    template_data=template_data,
    matching_cfg=dartsort.MatchingConfig(
        whitening=dartsort.WhiteningConfig(strategy="none"),
    ),
)

# then...
# match_st has properties like:
# spike times
match_st.times_seconds
match_st.times_samples
# matched template index
match_st.labels 
# matching score (higher is better)
match_st.scores

# you can also export this to phy, if helpful, see https://dartsort.github.io/#outputs-and-exporting