BUG: make_pc_features can associate PC features with the wrong channel IDs when multiple final clusters share a detection template

[gwfolk]

Describe the issue:

I stumbled onto this behavior in make_pc_features while investigating why reducing nearest_chans results in zero-valued/collapsed PC features on many clusters' peak channels, visible as a line on channel N/M plots and a point on single-channel A/B plots in Phy's FeatureView. Issue #787 briefly discusses that. Raising nearest_chans prevents the all-zero peak channels in FeatureView, but make_pc_features nonetheless swaps the channel assignments of PC features in tF in a manner that may be unintended.

Suppose final clusters i and j are derived from the same template. While looping through final clusters, make_pc_features...

1. identifies the detection templates associated with final cluster i
2. passes the complete spike_templates and mutable tF arrays to get_data_cpu, thus selecting all spikes using those detection templates, including spikes assigned to cluster j
3. reorders and overwrites all selected tF rows in place
4. when cluster j is up in the loop, the rows belonging to cluster j are interpreted as though they were still in their original detection-template channel order, but they were already reordered in the previous iteration targeting cluster i.
   Note that this makes the output incidentally depend on the order of iteration through clusters i and j.

The result with high channel spacing is that during curation, FeatureView can show a cloud of points that don't meaningfully reflect the waveforms on that channel. This can make the plots misleading during curation but isn't catastrophic in my experience. However, if some users want to reduce nearest_chans because they find that it produces better sorting results in their setup, FeatureView becomes a whole lot less useful because irrelevant channels with zero-value PCs get swapped into clusters at a relatively high rate.

One possible fix would make channel selection and reordering cluster-specific (gwfolk@668370c). With the linked branch, reducing nearest_chans does not lead to aberrant zero-valued PC features on peak channels in Phy. The function header says it replaces some data in tF "so that features are associated with the final clusters instead of templates", but if the pooling of spikes across clusters that share a template is intentional for the purpose of channel ranking, an alternative fix could preserve an immutable source tF and only update rows belonging to the current cluster i so that subsequent clusters with the same associated template won't be double-reordered.

Reproduce the bug:

# Minimal reproduction with two final clusters on two channels with spikes derived from the same detection template

import numpy as np
import torch
from kilosort.postprocessing import make_pc_features

ops = {
    "iU": torch.tensor([0]),
    "iCC": torch.tensor([[0], [1]]),
    "xc": np.array([0.0, 0.0]),
    "yc": np.array([0.0, 20.0]),
    "nearest_chans": 2,  # All available channels are included.
    "dmin": 20,
    "dminx": 32,
}

detection_templates = np.zeros(4, dtype=np.int32)
final_clusters = np.array([0, 1, 0, 1], dtype=np.int32)

# Columns initially represent channels [0, 1].
source = np.array([
    [1.0, 20.0],
    [1.0, 10.0],
    [2.0, 18.0],
    [2.0, 8.0],
], dtype=np.float32)

pc_features, pc_feature_ind = make_pc_features(
    ops, detection_templates, final_clusters,
    torch.tensor(source).unsqueeze(-1),
)

np.testing.assert_array_equal(
    pc_feature_ind,
    np.array([[1, 0], [1, 0]], dtype=np.uint32),
)

# Reconstruct dense channel values using the channel IDs exported for Phy.
restored = np.zeros_like(source)
for spike, cluster in enumerate(final_clusters):
    restored[spike, pc_feature_ind[cluster]] = pc_features[spike, 0].numpy()

print(pc_feature_ind)
print(restored)
np.testing.assert_allclose(restored, source)

Error message:

(No error message)
In the minimal reproduction above, Kilosort 4.1.7 produces:

pc_feature_ind:
[[1 0]
 [0 1]]

restored:
[[ 1. 20.]
 [10.  1.]
 [ 2. 18.]
 [ 8.  2.]]

Despite both clusters in fact having higher PC values on channel 1 than channel 0.

Version information:

Kilosort 4.1.7
Python 3.11.15
pytorch 2.7.1+cu118
Windows 10

[jacobpennington]

Hi,
Can you please give an example of where you see this happening in real data? Your supposition:
"Suppose final clusters i and j are derived from the same template" should never happen to begin with.