Fix motion awareness in split; debug split/merge ensemble and plots

Author: cwindolf

src/dartsort/cluster/density.py

@@ -72,6 +72,7 @@ def nearest_higher_density_neighbor(
     distances, indices = distances[:, 1:].copy(), indices[:, 1:].copy()
 
     # find lowest distance higher density neighbor
+    print(f"nhdn {density.shape=}")
     density_padded = np.pad(density, (0, 1), constant_values=np.inf)
     is_lower_density = density_padded[indices] <= density[:, None]
     distances[is_lower_density] = np.inf
@@ -151,6 +152,8 @@ def density_peaks_clustering(
     density = get_smoothed_densities(X, inliers=inliers, sigmas=sigmas)
     if sigma_regional is not None:
         density = density[0] / density[1]
+    else:
+        density = density[0]
 
     nhdn, distances, indices = nearest_higher_density_neighbor(
         kdtree,
@@ -182,6 +185,8 @@ def density_peaks_clustering(
     if remove_clusters_smaller_than:
         labels = decrumb(labels, min_size=remove_clusters_smaller_than)
 
+    print("dpc", np.unique(labels).size)
+
     if not return_extra:
         return labels
 

src/dartsort/cluster/ensemble_utils.py

@@ -32,6 +32,7 @@ def split_merge_ensemble(
             split_strategy=split_merge_config.split_strategy,
             recursive=split_merge_config.recursive_split,
             n_jobs=n_jobs_split,
+            motion_est=motion_est,
             show_progress=False
         )
         for sorting in tqdm(chunk_sortings, desc="Split within chunks")

src/dartsort/cluster/initial.py

@@ -226,10 +226,10 @@ def ensemble_chunks(
             recording,
             chunk_sortings,
             motion_est=motion_est,
-            split_merge_config=clustering_config.split_merge_config,
+            split_merge_config=clustering_config.split_merge_ensemble_config,
             n_jobs_split=computation_config.n_jobs_cpu,
             n_jobs_merge=computation_config.actual_n_jobs_gpu,
-            device=None,
+            device=computation_config.actual_device,
             show_progress=True,
         )
 

src/dartsort/cluster/merge.py

@@ -60,8 +60,9 @@ def merge_templates(
     -------
     A new DARTsortSorting
     """
+    print("merge input", np.unique(sorting.labels).size - 1)
     if template_data is None:
-        template_data = TemplateData.from_config(
+        template_data, sorting = TemplateData.from_config(
             recording,
             sorting,
             template_config,
@@ -71,6 +72,7 @@ def merge_templates(
             overwrite=overwrite_templates,
             device=device,
             save_npz_name=template_npz_filename,
+            return_realigned_sorting=True,
         )
 
     units, dists, shifts, template_snrs = calculate_merge_distances(
@@ -246,6 +248,10 @@ def calculate_merge_distances(
         show_progress=show_progress,
     )
     for res in dec_res_iter:
+        if res is None:
+            # all pairs in chunk were ignored for one reason or another
+            continue
+
         tixa = res.template_indices_a
         tixb = res.template_indices_b
         rms_ratio = res.deconv_resid_norms / res.template_a_norms
@@ -359,7 +365,12 @@ def recluster(
     pdist = dists[np.triu_indices(dists.shape[0], k=1)]
     # scipy hierarchical clustering only supports finite values, so let's just
     # drop in a huge value here
-    pdist[~np.isfinite(pdist)] = 1_000_000 + pdist[np.isfinite(pdist)].max()
+    finite = np.isfinite(pdist)
+    if not finite.any():
+        print("no merges")
+        return sorting
+
+    pdist[~finite] = 1_000_000 + pdist[finite].max()
     # complete linkage: max dist between all pairs across clusters.
     Z = complete(pdist)
     # extract flat clustering using our max dist threshold
@@ -378,6 +389,7 @@ def recluster(
     clust_inverse = {i: [] for i in new_labels}
     for orig_label, new_label in enumerate(new_labels):
         clust_inverse[new_label].append(orig_label)
+    print(sum(len(v) - 1 for v in clust_inverse.values()), "merges")
 
     # align to best snr unit
     for new_label, orig_labels in clust_inverse.items():
@@ -409,8 +421,8 @@ def cross_match(
     units_b,
     merge_distance_threshold=0.5,
 ):
-    assert np.array_equal(units_a, sorting_a.units)
-    assert np.array_equal(units_b, sorting_b.units)
+    assert np.array_equal(units_a, sorting_a.unit_ids)
+    assert np.array_equal(units_b, sorting_b.unit_ids)
 
     ia, ib = np.nonzero(dists <= merge_distance_threshold)
     weights = coo_array((-dists[ia, ib], (ia.astype(np.intc), ib.astype(np.intc))))

src/dartsort/cluster/split.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, replace
 from typing import Optional
+from pathlib import Path
 
 import h5py
 import numpy as np
@@ -24,6 +25,7 @@ def split_clusters(
     recursive=False,
     split_big=False,
     split_big_kw=dict(dz=40, dx=48, min_size_split=50),
+    motion_est=None,
     show_progress=True,
     n_jobs=0,
 ):
@@ -50,6 +52,11 @@ def split_clusters(
     labels_to_process = list(labels_to_process[labels_to_process > 0])
     cur_max_label = max(labels_to_process)
 
+    if split_strategy_kwargs is None:
+        split_strategy_kwargs = {}
+    if motion_est is not None:
+        split_strategy_kwargs["motion_est"] = motion_est
+
     n_jobs, Executor, context = get_pool(n_jobs)
     with Executor(
         max_workers=n_jobs,
@@ -667,6 +674,7 @@ def initialize_from_h5(
         amplitudes_dataset_name="denoised_ptp_amplitudes",
         amplitude_vectors_dataset_name="denoised_ptp_amplitude_vectors",
     ):
+        peeling_hdf5_filename = Path(peeling_hdf5_filename)
         h5 = h5py.File(peeling_hdf5_filename, "r")
         self.geom = h5["geom"][:]
         self.channel_index = h5["channel_index"][:]
@@ -704,6 +712,11 @@ def initialize_from_h5(
             self.tpca_features = h5[tpca_features_dataset_name]
             self.match_distance = pdist(self.geom).min() / 2
 
+        if peeling_featurization_pt is None:
+            mdir = peeling_hdf5_filename.parent / f"{peeling_hdf5_filename.stem}_models"
+            peeling_featurization_pt = mdir / "featurization_pipeline.pt"
+            assert peeling_featurization_pt.exists()
+
         if self.n_pca_features and self.relocated:
             # load up featurization pipeline for tpca inversion
             assert peeling_featurization_pt is not None

src/dartsort/config.py

@@ -206,7 +206,7 @@ class SplitMergeConfig:
     merge_template_config: TemplateConfig = TemplateConfig(superres_templates=False)
     merge_distance_threshold: float = 0.25
     cross_merge_distance_threshold: float = 0.5
-    min_spatial_cosine: float = 0.5
+    min_spatial_cosine: float = 0.0
 
 
 @dataclass(frozen=True)
@@ -235,15 +235,17 @@ class ClusteringConfig:
     # density peaks parameters
     sigma_local: float = 5.0
     sigma_regional: Optional[float] = None
-    n_neighbors_search: int = 10
+    n_neighbors_search: int = 20
     radius_search: float = 5.0
     remove_clusters_smaller_than: int = 10
     noise_density: float = 0.0
 
     # -- ensembling
     ensemble_strategy: Optional[str] = "forward_backward"
     chunk_size_s: float = 300.0
-    split_merge_ensemble_config: SplitMergeConfig = SplitMergeConfig()
+    split_merge_ensemble_config: SplitMergeConfig = SplitMergeConfig(
+        merge_template_config=TemplateConfig(superres_templates=False, realign_peaks=False,)
+    )
 
 
 @dataclass(frozen=True)
@@ -280,7 +282,7 @@ def __post_init__(self):
 
         if self.actual_device.type == "cuda":
             self.actual_n_jobs_gpu = self.n_jobs_gpu
-        if self.actual_device.type == "cuda":
+        else:
             self.actual_n_jobs_gpu = self.n_jobs_cpu
 
 

src/dartsort/main.py

@@ -196,6 +196,7 @@ def split_merge(
             split_strategy=split_merge_config.split_strategy,
             recursive=split_merge_config.recursive_split,
             n_jobs=n_jobs_split,
+            motion_est=motion_est,
         )
         if output_directory is not None:
             split_sorting.save(split_npz)
@@ -206,6 +207,7 @@ def split_merge(
         merge_sorting = merge_templates(
             split_sorting,
             recording,
+            motion_est=motion_est,
             template_config=split_merge_config.merge_template_config,
             merge_distance_threshold=split_merge_config.merge_distance_threshold,
             min_spatial_cosine=split_merge_config.min_spatial_cosine,

src/dartsort/templates/templates.py

@@ -108,6 +108,7 @@ def from_config(
         device=None,
         trough_offset_samples=42,
         spike_length_samples=121,
+        return_realigned_sorting=False,
     ):
         if save_folder is not None:
             save_folder = Path(save_folder)
@@ -181,7 +182,7 @@ def from_config(
 
         # handle superresolved templates
         if template_config.superres_templates:
-            unit_ids, sorting = superres_sorting(
+            unit_ids, superres_sorting = superres_sorting(
                 sorting,
                 sorting.times_seconds,
                 spike_depths_um,
@@ -194,6 +195,7 @@ def from_config(
                 adaptive_bin_size=template_config.adaptive_bin_size,
             )
         else:
+            superres_sorting = sorting
             # we don't skip empty units
             unit_ids = np.arange(sorting.labels.max() + 1)
 
@@ -203,7 +205,7 @@ def from_config(
         spike_counts[ix[ix >= 0]] = counts[ix >= 0]
 
         # main!
-        results = get_templates(recording, sorting, **kwargs)
+        results = get_templates(recording, superres_sorting, **kwargs)
 
         # handle registered templates
         if template_config.registered_templates and motion_est is not None:
@@ -233,4 +235,7 @@ def from_config(
         if save_folder is not None:
             obj.to_npz(npz_path)
 
+        if return_realigned_sorting:
+            return obj, sorting
+
         return obj

src/dartsort/vis/scatterplots.py

@@ -239,12 +239,12 @@ def scatter_time_vs_depth(
         if depths_um is None:
             depths_um = getattr(sorting, "point_source_localizations", None)
             if depths_um is not None:
-                depths_um = x[:, 2]
+                depths_um = depths_um[:, 2]
         if amplitudes is None:
             amplitudes = getattr(sorting, amplitudes_dataset_name, None)
         if hdf5_filename is None:
             hdf5_filename = sorting.parent_h5_path
-    
+
     needs_load = any(v is None for v in (times_s, depths_um, amplitudes, geom))
     if needs_load and hdf5_filename is not None:
         with h5py.File(hdf5_filename, "r") as h5:
@@ -473,6 +473,7 @@ def scatter_feature_vs_depth(
 
     if sorting is not None:
         labels = sorting.labels
+
     if labels is None:
         c = np.clip(amplitudes, 0, amplitude_color_cutoff)
         cmap = amplitude_cmap
@@ -490,7 +491,6 @@ def scatter_feature_vs_depth(
             rasterized=rasterized,
             **scatter_kw,
         )
-        to_show = to_show[kept]
     else:
         c = labels
         cmap = colorcet.m_glasbey_light