Skip to content

Update tridesclous2 #2860

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

Merged
samuelgarcia merged 12 commits into from
May 24, 2024
Merged

Update tridesclous2 #2860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
602e37a
Add "apply_motion_correction" in tridesclous2
samuelgarcia May 14, 2024
1517cc9
improve benchmark matching with collision option
samuelgarcia May 15, 2024
1d03eec
benchmark clustering return fig
samuelgarcia May 15, 2024
340dd12
resolve conflict with main
samuelgarcia May 16, 2024
6a2f00d
Merge branch 'main' of github.com:SpikeInterface/spikeinterface into …
samuelgarcia May 17, 2024
44b28dd
handle colors in benchmark components
samuelgarcia May 17, 2024
9a03744
Merge branch 'main' of github.com:SpikeInterface/spikeinterface into …
samuelgarcia May 22, 2024
ef0bba9
clean
samuelgarcia May 22, 2024
210d913
yep
samuelgarcia May 22, 2024
3ab27d4
tdc update and clean
samuelgarcia May 24, 2024
88d2dfd
yep
samuelgarcia May 24, 2024
023baba
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 24, 2024
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
137 changes: 64 additions & 73 deletions src/spikeinterface/sorters/internal/tridesclous2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from spikeinterface.core import (
get_noise_levels,
NumpySorting,
get_channel_distances,
estimate_templates_with_accumulator,
Templates,
compute_sparsity,
Expand All @@ -18,29 +17,26 @@
from spikeinterface.preprocessing import bandpass_filter, common_reference, zscore, whiten
from spikeinterface.core.basesorting import minimum_spike_dtype

from spikeinterface.sortingcomponents.tools import extract_waveform_at_max_channel, cache_preprocessing
from spikeinterface.sortingcomponents.tools import cache_preprocessing

# from spikeinterface.qualitymetrics import compute_snrs

import numpy as np

import pickle
import json


class Tridesclous2Sorter(ComponentsBasedSorter):
sorter_name = "tridesclous2"

_default_params = {
"apply_preprocessing": True,
"apply_motion_correction": False,
"motion_correction": {"preset": "nonrigid_fast_and_accurate"},
"cache_preprocessing": {"mode": "memory", "memory_limit": 0.5, "delete_cache": True},
"waveforms": {
"ms_before": 0.5,
"ms_after": 1.5,
"radius_um": 120.0,
},
"filtering": {"freq_min": 300.0, "freq_max": 12000.0},
"filtering": {"freq_min": 300.0, "freq_max": 8000.0},
"detection": {"peak_sign": "neg", "detect_threshold": 5, "exclude_sweep_ms": 1.5, "radius_um": 150.0},
"selection": {"n_peaks_per_channel": 5000, "min_n_peaks": 20000},
"svd": {"n_components": 6},
Expand All @@ -53,7 +49,7 @@ class Tridesclous2Sorter(ComponentsBasedSorter):
"ms_before": 2.0,
"ms_after": 3.0,
"max_spikes_per_unit": 400,
"sparsity_threshold": 2.0,
"sparsity_threshold": 1.5,
# "peak_shift_ms": 0.2,
},
# "matching": {"method": "tridesclous", "method_kwargs": {"peak_shift_ms": 0.2, "radius_um": 100.0}},
Expand Down Expand Up @@ -86,31 +82,18 @@ def get_sorter_version(cls):

@classmethod
def _run_from_folder(cls, sorter_output_folder, params, verbose):
job_kwargs = params["job_kwargs"].copy()
job_kwargs = fix_job_kwargs(job_kwargs)
job_kwargs["progress_bar"] = verbose

from spikeinterface.sortingcomponents.matching import find_spikes_from_templates
from spikeinterface.core.node_pipeline import (
run_node_pipeline,
ExtractDenseWaveforms,
ExtractSparseWaveforms,
PeakRetriever,
)
from spikeinterface.sortingcomponents.peak_detection import detect_peaks, DetectPeakLocallyExclusive
from spikeinterface.sortingcomponents.peak_detection import detect_peaks
from spikeinterface.sortingcomponents.peak_selection import select_peaks
from spikeinterface.sortingcomponents.peak_localization import LocalizeCenterOfMass, LocalizeGridConvolution
from spikeinterface.sortingcomponents.waveforms.temporal_pca import TemporalPCAProjection

from spikeinterface.sortingcomponents.clustering.split import split_clusters
from spikeinterface.sortingcomponents.clustering.merge import merge_clusters
from spikeinterface.sortingcomponents.clustering.tools import compute_template_from_sparse
from spikeinterface.sortingcomponents.clustering.main import find_cluster_from_peaks
from spikeinterface.sortingcomponents.tools import remove_empty_templates
from spikeinterface.preprocessing import correct_motion
from spikeinterface.sortingcomponents.motion_interpolation import InterpolateMotionRecording

from sklearn.decomposition import TruncatedSVD

import hdbscan
job_kwargs = params["job_kwargs"].copy()
job_kwargs = fix_job_kwargs(job_kwargs)
job_kwargs["progress_bar"] = verbose

recording_raw = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False)

Expand All @@ -119,10 +102,44 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):

# preprocessing
if params["apply_preprocessing"]:
recording = bandpass_filter(recording_raw, **params["filtering"])
if params["apply_motion_correction"]:
rec_for_motion = recording_raw
if params["apply_preprocessing"]:
rec_for_motion = bandpass_filter(rec_for_motion, freq_min=300.0, freq_max=6000.0, dtype="float32")
rec_for_motion = common_reference(rec_for_motion)
if verbose:
print("Start correct_motion()")
_, motion_info = correct_motion(
rec_for_motion,
folder=sorter_output_folder / "motion",
output_motion_info=True,
**params["motion_correction"],
)
if verbose:
print("Done correct_motion()")

recording = bandpass_filter(recording_raw, **params["filtering"], dtype="float32")
recording = common_reference(recording)

if params["apply_motion_correction"]:
interpolate_motion_kwargs = dict(
direction=1,
border_mode="force_extrapolate",
spatial_interpolation_method="kriging",
sigma_um=20.0,
p=2,
)

recording = InterpolateMotionRecording(
recording,
motion_info["motion"],
motion_info["temporal_bins"],
motion_info["spatial_bins"],
**interpolate_motion_kwargs,
)

recording = zscore(recording, dtype="float32")
recording = whiten(recording, dtype="float32")
recording = whiten(recording, dtype="float32", mode="local", radius_um=100.0)

# used only if "folder" or "zarr"
cache_folder = sorter_output_folder / "cache_preprocessing"
Expand All @@ -141,7 +158,7 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
all_peaks = detect_peaks(recording, method="locally_exclusive", **detection_params, **job_kwargs)

if verbose:
print("We found %d peaks in total" % len(all_peaks))
print(f"detect_peaks(): {len(all_peaks)} peaks found")

# selection
selection_params = params["selection"].copy()
Expand All @@ -150,36 +167,38 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
peaks = select_peaks(all_peaks, method="uniform", n_peaks=n_peaks)

if verbose:
print("We kept %d peaks for clustering" % len(peaks))
print(f"select_peaks(): {len(peaks)} peaks kept for clustering")

clustering_kwargs = {}
clustering_kwargs["folder"] = sorter_output_folder
clustering_kwargs["waveforms"] = params["waveforms"].copy()
clustering_kwargs["clustering"] = params["clustering"].copy()

labels_set, post_clean_label, extra_out = find_cluster_from_peaks(
labels_set, clustering_label, extra_out = find_cluster_from_peaks(
recording, peaks, method="tdc_clustering", method_kwargs=clustering_kwargs, extra_outputs=True, **job_kwargs
)
peak_shifts = extra_out["peak_shifts"]
new_peaks = peaks.copy()
new_peaks["sample_index"] -= peak_shifts

mask = post_clean_label >= 0
mask = clustering_label >= 0
sorting_pre_peeler = NumpySorting.from_times_labels(
new_peaks["sample_index"][mask],
post_clean_label[mask],
clustering_label[mask],
sampling_frequency,
unit_ids=labels_set,
)
# sorting_pre_peeler = sorting_pre_peeler.save(folder=sorter_output_folder / "sorting_pre_peeler")

recording_w = whiten(recording, mode="local", radius_um=100.0)
if verbose:
print(f"find_cluster_from_peaks(): {sorting_pre_peeler.unit_ids.size} cluster found")

recording_for_peeler = recording

nbefore = int(params["templates"]["ms_before"] * sampling_frequency / 1000.0)
nafter = int(params["templates"]["ms_after"] * sampling_frequency / 1000.0)
sparsity_threshold = params["templates"]["sparsity_threshold"]

templates_array = estimate_templates_with_accumulator(
recording_w,
recording_for_peeler,
sorting_pre_peeler.to_spike_vector(),
sorting_pre_peeler.unit_ids,
nbefore,
Expand All @@ -192,61 +211,33 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
sampling_frequency=sampling_frequency,
nbefore=nbefore,
sparsity_mask=None,
probe=recording_w.get_probe(),
probe=recording_for_peeler.get_probe(),
is_scaled=False,
)

# TODO : try other methods for sparsity
sparsity_threshold = params["templates"]["sparsity_threshold"]
# sparsity = compute_sparsity(templates_dense, method="radius", radius_um=120.)
sparsity = compute_sparsity(templates_dense, noise_levels=noise_levels, threshold=sparsity_threshold)
templates = templates_dense.to_sparse(sparsity)
templates = remove_empty_templates(templates)

# snrs = compute_snrs(we, peak_sign=params["detection"]["peak_sign"], peak_mode="extremum")
# print(snrs)

# matching_params = params["matching"].copy()
# matching_params["noise_levels"] = noise_levels
# matching_params["peak_sign"] = params["detection"]["peak_sign"]
# matching_params["detect_threshold"] = params["detection"]["detect_threshold"]
# matching_params["radius_um"] = params["detection"]["radius_um"]

# spikes = find_spikes_from_templates(
# recording, method="tridesclous", method_kwargs=matching_params, **job_kwargs
# )

## peeler
matching_method = params["matching"]["method"]
matching_params = params["matching"]["method_kwargs"].copy()

matching_params["templates"] = templates
matching_params["noise_levels"] = noise_levels
# matching_params["peak_sign"] = params["detection"]["peak_sign"]
# matching_params["detect_threshold"] = params["detection"]["detect_threshold"]
# matching_params["radius_um"] = params["detection"]["radius_um"]

# spikes = find_spikes_from_templates(
# recording, method="tridesclous", method_kwargs=matching_params, **job_kwargs
# )
# )

# if matching_method == "circus-omp-svd":
# job_kwargs = job_kwargs.copy()
# for value in ["chunk_size", "chunk_memory", "total_memory", "chunk_duration"]:
# if value in job_kwargs:
# job_kwargs.pop(value)
# job_kwargs["chunk_duration"] = "100ms"

spikes = find_spikes_from_templates(
recording_w, method=matching_method, method_kwargs=matching_params, **job_kwargs
recording_for_peeler, method=matching_method, method_kwargs=matching_params, **job_kwargs
)

if params["save_array"]:
sorting_pre_peeler = sorting_pre_peeler.save(folder=sorter_output_folder / "sorting_pre_peeler")

np.save(sorter_output_folder / "noise_levels.npy", noise_levels)
np.save(sorter_output_folder / "all_peaks.npy", all_peaks)
# np.save(sorter_output_folder / "post_split_label.npy", post_split_label)
# np.save(sorter_output_folder / "split_count.npy", split_count)
# np.save(sorter_output_folder / "post_merge_label.npy", post_merge_label)
np.save(sorter_output_folder / "peaks.npy", peaks)
np.save(sorter_output_folder / "clustering_label.npy", clustering_label)
np.save(sorter_output_folder / "spikes.npy", spikes)

final_spikes = np.zeros(spikes.size, dtype=minimum_spike_dtype)
Expand Down
21 changes: 21 additions & 0 deletions src/spikeinterface/sortingcomponents/benchmark/benchmark_clustering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,8 @@ def plot_agreements(self, case_keys=None, figsize=(15, 15)):
ax.set_title(self.cases[key]["label"])
plot_agreement_matrix(self.get_result(key)["gt_comparison"], ax=ax)

return fig

def plot_performances_vs_snr(self, case_keys=None, figsize=(15, 15)):
if case_keys is None:
case_keys = list(self.cases.keys())
Expand All @@ -210,6 +212,8 @@ def plot_performances_vs_snr(self, case_keys=None, figsize=(15, 15)):
if count == 2:
ax.legend()

return fig

def plot_error_metrics(self, metric="cosine", case_keys=None, figsize=(15, 5)):

if case_keys is None:
Expand Down Expand Up @@ -244,6 +248,8 @@ def plot_error_metrics(self, metric="cosine", case_keys=None, figsize=(15, 5)):
label = self.cases[key]["label"]
axs[0, count].set_title(label)

return fig

def plot_metrics_vs_snr(self, metric="agreement", case_keys=None, figsize=(15, 5)):

if case_keys is None:
Expand Down Expand Up @@ -296,6 +302,8 @@ def plot_metrics_vs_snr(self, metric="agreement", case_keys=None, figsize=(15, 5
axs[0, count].set_title(label)
axs[0, count].legend()

return fig

def plot_metrics_vs_depth_and_snr(self, metric="agreement", case_keys=None, figsize=(15, 5)):

if case_keys is None:
Expand Down Expand Up @@ -354,6 +362,8 @@ def plot_metrics_vs_depth_and_snr(self, metric="agreement", case_keys=None, figs
axs[0, count].set_title(label)
# axs[0, count].legend()

return fig

def plot_unit_losses(self, case_before, case_after, metric="agreement", figsize=None):

fig, axs = plt.subplots(ncols=1, nrows=3, figsize=figsize)
Expand Down Expand Up @@ -384,6 +394,7 @@ def plot_unit_losses(self, case_before, case_after, metric="agreement", figsize=
fig.colorbar(im, ax=ax)
ax.set_title(k)
ax.set_ylabel("snr")
return fig

def plot_comparison_clustering(
self,
Expand Down Expand Up @@ -444,10 +455,13 @@ def plot_comparison_clustering(

plt.tight_layout(h_pad=0, w_pad=0)

return fig

def plot_some_over_merged(self, case_keys=None, overmerged_score=0.05, max_units=5, figsize=None):
if case_keys is None:
case_keys = list(self.cases.keys())

figs = []
for count, key in enumerate(case_keys):
label = self.cases[key]["label"]
comp = self.get_result(key)["gt_comparison"]
Expand Down Expand Up @@ -475,13 +489,17 @@ def plot_some_over_merged(self, case_keys=None, overmerged_score=0.05, max_units
ax.set_xticks([])

fig.suptitle(label)
figs.append(fig)
else:
print(key, "no overmerged")

return figs

def plot_some_over_splited(self, case_keys=None, oversplit_score=0.05, max_units=5, figsize=None):
if case_keys is None:
case_keys = list(self.cases.keys())

figs = []
for count, key in enumerate(case_keys):
label = self.cases[key]["label"]
comp = self.get_result(key)["gt_comparison"]
Expand Down Expand Up @@ -509,5 +527,8 @@ def plot_some_over_splited(self, case_keys=None, oversplit_score=0.05, max_units
ax.set_xticks([])

fig.suptitle(label)
figs.append(fig)
else:
print(key, "no over splited")

return figs
Loading