Skip to content

perf: reuse TDC peeler distance work and bound temporal neighbor searches #4874

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JESUSROYETH wants to merge 1 commit into
base: main
Choose a base branch
from
+164 −4
Open

perf: reuse TDC peeler distance work and bound temporal neighbor searches #4874

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 119 additions & 4 deletions src/spikeinterface/sortingcomponents/matching/tdc_peeler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -767,10 +767,21 @@ def get_most_probable_cluster(
wf_short = traces[s0:s1, :]

## numba with cluster+channel spasity
union_channels = np.any(template_sparsity_mask[possible_clusters, :], axis=0)
distances = numba_sparse_distance(
wf_short, sparse_templates_array, template_sparsity_mask, union_channels, possible_clusters
)
selected_sparsity = template_sparsity_mask[possible_clusters, :]
union_channels = np.any(selected_sparsity, axis=0)
num_union_channels = np.count_nonzero(union_channels)
num_zero_channel_comparisons = possible_clusters.size * num_union_channels - np.count_nonzero(selected_sparsity)
if (
wf_short.dtype.kind == sparse_templates_array.dtype.kind == "f"
and num_zero_channel_comparisons * (wf_short.shape[0] - 1) > num_union_channels * wf_short.shape[0]
):
distances = _numba_sparse_distance_reused_energy(
wf_short, sparse_templates_array, template_sparsity_mask, union_channels, possible_clusters
)
else:
distances = numba_sparse_distance(
wf_short, sparse_templates_array, template_sparsity_mask, union_channels, possible_clusters
)

ind = np.argmin(distances)
cluster_index = possible_clusters[ind]
Expand All @@ -779,6 +790,46 @@ def get_most_probable_cluster(


def get_neighbors_spikes(sample_inds, chan_inds, delta_sample, near_chan_mask):
# Bound temporal candidates first, keeping the original index order for
# the joint amplitude fit. Other input types retain the existing behavior.
if (
type(sample_inds) is np.ndarray
and sample_inds.dtype == np.int64
and sample_inds.ndim == 1
and sample_inds.size > 1
and type(chan_inds) is np.ndarray
and chan_inds.dtype == np.int64
and chan_inds.shape == sample_inds.shape
and type(near_chan_mask) is np.ndarray
and near_chan_mask.dtype == bool
and near_chan_mask.ndim == 2
and isinstance(delta_sample, (int, np.signedinteger))
and delta_sample >= 0
):
integer_delta = int(delta_sample)
first_sample = int(np.min(sample_inds))
last_sample = int(np.max(sample_inds))
limits = np.iinfo(np.int64)
if (
first_sample >= limits.min + integer_delta
and last_sample <= limits.max - integer_delta
and last_sample - first_sample <= limits.max
and np.min(chan_inds) >= 0
and np.max(chan_inds) < min(near_chan_mask.shape)
):
sample_order = np.argsort(sample_inds)
samples_sorted = sample_inds[sample_order]
left_bounds = np.searchsorted(samples_sorted, sample_inds - integer_delta, side="right")
right_bounds = np.searchsorted(samples_sorted, sample_inds + integer_delta, side="left")
neighbors_spikes_indices = []
for spike_index in range(sample_inds.size):
candidate_indices = sample_order[left_bounds[spike_index] : right_bounds[spike_index]]
if candidate_indices.size:
nearby = near_chan_mask[chan_inds[spike_index], chan_inds[candidate_indices]]
nearby = nearby & (candidate_indices != spike_index)
candidate_indices = np.sort(candidate_indices[nearby])
neighbors_spikes_indices.append(list(candidate_indices))
return neighbors_spikes_indices

neighbors_spikes_inds = []
for i in range(sample_inds.size):
Expand Down Expand Up @@ -965,6 +1016,70 @@ def numba_sparse_distance(
distances[i] = sum_dist
return distances

@jit(nopython=True, nogil=True)
def _numba_sparse_distance_reused_energy(
waveform, sparse_templates_array, template_sparsity_mask, wanted_channel_mask, possible_clusters
):
"""Compute distances with the union of all candidate support channels."""
width, total_channels = waveform.shape
channel_indices = np.flatnonzero(wanted_channel_mask)
zero_energy = np.zeros(total_channels, dtype=np.float64)
for channel_index in channel_indices:
for sample_index in range(width):
zero_value = np.float64(waveform[sample_index, channel_index])
zero_energy[channel_index] += zero_value**2

distances = np.zeros(possible_clusters.size, dtype=np.float32)
# Both nonnegative sums have at most 2N additions on any path.
# Bound their separation before float32 rounding; ambiguous cases keep
# the original summation order.
roundoff = 8.0 * waveform.size * np.finfo(np.float64).eps
for candidate_index in range(possible_clusters.size):
cluster_index = possible_clusters[candidate_index]
total = 0.0
sparse_channel_index = 0
for channel_index in channel_indices:
if template_sparsity_mask[cluster_index, channel_index]:
# Independent accumulators shorten the reduction dependency chain.
partial0 = partial1 = partial2 = partial3 = 0.0
grouped_width = width // 4 * 4
for sample_index in range(0, grouped_width, 4):
value0 = waveform[sample_index + 0, channel_index]
template0 = sparse_templates_array[cluster_index, sample_index + 0, sparse_channel_index]
partial0 += (value0 - template0) ** 2
value1 = waveform[sample_index + 1, channel_index]
template1 = sparse_templates_array[cluster_index, sample_index + 1, sparse_channel_index]
partial1 += (value1 - template1) ** 2
value2 = waveform[sample_index + 2, channel_index]
template2 = sparse_templates_array[cluster_index, sample_index + 2, sparse_channel_index]
partial2 += (value2 - template2) ** 2
value3 = waveform[sample_index + 3, channel_index]
template3 = sparse_templates_array[cluster_index, sample_index + 3, sparse_channel_index]
partial3 += (value3 - template3) ** 2
total += (partial0 + partial1) + (partial2 + partial3)
for sample_index in range(grouped_width, width):
value = waveform[sample_index, channel_index]
template_value = sparse_templates_array[cluster_index, sample_index, sparse_channel_index]
total += (value - template_value) ** 2
sparse_channel_index += 1
else:
total += zero_energy[channel_index]

if np.isfinite(total) and roundoff < 0.5:
lower = np.float32(np.nextafter(total * (1.0 - roundoff), -np.inf))
upper = np.float32(np.nextafter(total * (1.0 + roundoff), np.inf))
if lower == upper and np.isfinite(lower):
distances[candidate_index] = np.float32(total)
continue
distances[candidate_index] = numba_sparse_distance(
waveform,
sparse_templates_array,
template_sparsity_mask,
wanted_channel_mask,
possible_clusters[candidate_index : candidate_index + 1],
)[0]
return distances

@jit(nopython=True, nogil=True)
def numba_best_shift_sparse(
traces, sparse_template, sample_index, nbefore, possible_shifts, distances_shift, chan_sparsity
Expand Down
45 changes: 45 additions & 0 deletions src/spikeinterface/sortingcomponents/matching/tests/test_template_matching.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,51 @@ def test_find_spikes_from_templates(method, sorting_analyzer):
# plt.show()


@pytest.mark.parametrize("dtype", ["float32", "float64"])
def test_sparse_distance_preserves_rounding(dtype, monkeypatch):
import spikeinterface.sortingcomponents.matching.tdc_peeler as peeler

# Sequential additions lose the 2**-54 terms after the float32 midpoint;
# grouping the last channel first would round these candidates upwards.
waveform = np.array([[1, 2**-12], [0, 2**-27], [0, 2**-27], [0, 2**-27]], dtype=dtype)
templates = np.zeros((4, 4, 1), dtype=dtype)
templates[[0, 2], 0, 0] = 2
templates[[1, 3], :, 0] = waveform[:, 1]
sparsity = np.array([[True, False], [False, True], [True, False], [False, True]])
candidates = np.array([2, 3, 0, 1])
calls = []
original = peeler._numba_sparse_distance_reused_energy

def checked(*args):
calls.append(True)
return original(*args)

monkeypatch.setattr(peeler, "_numba_sparse_distance_reused_energy", checked)
cluster_index = peeler.get_most_probable_cluster(waveform, templates, candidates, 2, 2, 2, sparsity)
assert cluster_index == 2
assert len(calls) == 1


def test_temporal_neighbors_preserve_order(monkeypatch):
import spikeinterface.sortingcomponents.matching.tdc_peeler as peeler

sample_indices = np.array([10, 0, 8, 10, 11, 9], dtype="int64")
channel_indices = np.array([1, 0, 1, 0, 1, 1], dtype="int64")
near_channels = np.array([[True, False], [True, True]])
original_samples = sample_indices.tobytes()
original_channels = channel_indices.tobytes()

def unexpected_full_scan(*args):
raise AssertionError("Full sample scan")

monkeypatch.setattr(peeler.np, "abs", unexpected_full_scan)
neighbors = peeler.get_neighbors_spikes(sample_indices, channel_indices, 2, near_channels)
assert neighbors == [[3, 4, 5], [], [5], [], [0, 3], [0, 2, 3]]
assert all(isinstance(index, np.integer) for row in neighbors for index in row)
assert sample_indices.tobytes() == original_samples
assert channel_indices.tobytes() == original_channels


if __name__ == "__main__":
sorting_analyzer = get_sorting_analyzer()
# method = "nearest"
Expand Down
Loading