dev: fix tests

Author: cwindolf

src/dartsort/clustering/cluster_util.py

@@ -332,7 +332,6 @@ def reorder_by_depth(
     geom: np.ndarray | None = None,
     centroids: np.ndarray | None = None,
     in_place: bool = False,
-    is_flat: bool = False,
 ) -> tuple[DARTsortSorting, np.ndarray]:
     """Reorder cluster labels so that centroid depth is increasing
 
@@ -350,9 +349,7 @@ def reorder_by_depth(
     reorder: np.ndarray
         reorder[j] is the new label of original unit j.
     """
-
-    if not is_flat:
-        sorting = sorting.flatten(include_gmm_properties=True, in_place=in_place)
+    sorting = sorting.flatten(include_gmm_properties=True, in_place=in_place)
     assert sorting.labels is not None
 
     if geom is None and motion is not None:
@@ -540,9 +537,11 @@ def decrumb_labels(labels: np.ndarray, min_size: int = 5, in_place=False, flatte
         The (flattened) decrumbed labels.
     """
     units, counts, _ = pos_int_unique_and_counts(labels)
+    if not units.size:
+        return labels
     all_big = counts.min() >= min_size
     flat_ok = (not flatten) or np.array_equal(units, np.arange(len(units)))
-    if (not units.size) or (all_big and flat_ok):
+    if all_big and flat_ok:
         return labels
     remapping = np.full((units.max() + 1,), -1, dtype=labels.dtype)
     kept_units = units[counts >= min_size]
@@ -561,7 +560,11 @@ def decrumb(
 ) -> DARTsortSorting:
     assert sorting.labels is not None
     units, counts, _ = pos_int_unique_and_counts(sorting.labels)
-    if (not units.size) or (counts.min() >= min_size):
+    if not units.size:
+        return sorting
+    all_big = counts.min() >= min_size
+    flat_ok = (not flatten) or np.array_equal(units, np.arange(len(units)))
+    if all_big and flat_ok:
         return sorting
 
     remapping = np.full((units.max() + 1,), -1, dtype=sorting.labels.dtype)

src/dartsort/templates/postprocess_util.py

@@ -178,7 +178,7 @@ def estimate_template_library(
 
     # re-order along probe length
     if depth_order:
-        sorting, templates = reorder_by_depth(sorting, templates)
+        sorting, templates = reorder_templates_by_depth(sorting, templates)
 
     return sorting, ensure_save(templates, template_npz_path)
 
@@ -336,7 +336,7 @@ def snr_mask(template_data, min_n_spikes=50, min_template_snr=15.0):
     return good_templates
 
 
-def reorder_by_depth(sorting, template_data):
+def reorder_templates_by_depth(sorting, template_data):
     assert template_data.registered_geom is not None
     w = template_data.snrs_by_channel()
     w /= w.sum(axis=1, keepdims=True)

src/dartsort/util/data_util.py

@@ -936,11 +936,12 @@ def flatten(
         old_unique, _, _ = pos_int_unique_and_counts(self.labels)
         old_K = old_unique.max() + 1 if old_unique.shape[0] > 0 else 0
         if np.array_equal(old_unique, np.arange(old_K)):
-            return self
-
-        new_labels = self.labels if in_place else self.labels.copy()
-        remap = flatten_remapping(old_unique)
-        apply_label_remapping_in_place(new_labels, remap)
+            new_labels = self.labels
+            remap = np.arange(old_K)
+        else:
+            new_labels = self.labels if in_place else self.labels.copy()
+            remap = flatten_remapping(old_unique)
+            apply_label_remapping_in_place(new_labels, remap)
 
         keys = ("merged_candidates", "gmm_candidates")
         if not include_gmm_properties or not any(hasattr(self, k) for k in keys):
@@ -1038,7 +1039,6 @@ def has_dataset(self, dataset_name: str) -> bool:
             return dataset_name in h5
 
     def load_dataset(self, dataset_name: str, sl=()) -> np.ndarray:
-
         if dataset_name in self._ephemeral_features:
             return self._ephemeral_features[dataset_name][sl]
         if dataset_name in self._persistent_features:

tests/test_alignment.py

@@ -115,7 +115,7 @@ def test_denoiser_alignment(align_sim, align_templates):
     ci = dartsort.waveform_util.full_channel_index(1, to_torch=True)
     rec = align_sim["recording"]
     gt_st = align_sim["sorting"]
-    noise_wfs = gt_st._load_dataset("noise_waveforms")
+    noise_wfs = gt_st.load_dataset("noise_waveforms")
 
     # an optimal linear denoiser
     rolls = (-1, 0, 1)

tests/test_clustering.py

@@ -10,7 +10,7 @@
     refinement_strategies,
 )
 from dartsort.main import cluster
-from dartsort.templates.postprocess_util import reorder_by_depth
+from dartsort.templates.postprocess_util import reorder_templates_by_depth
 from dartsort.util.internal_config import (
     ClusteringConfig,
     ClusteringFeaturesConfig,
@@ -208,7 +208,7 @@ def test_reorder_by_depth(simulations, sim_name):
     sorting = sim["sorting"]
     template_data = sim["templates"]
 
-    sorting1, template_data1 = reorder_by_depth(sorting, template_data)
+    sorting1, template_data1 = reorder_templates_by_depth(sorting, template_data)
 
     assert np.array_equal(sorting.times_samples, sorting1.times_samples)
     assert np.array_equal(sorting.channels, sorting1.channels)

tests/test_matching.py

@@ -185,7 +185,7 @@ def test_no_crumbs(subtests, refractory_sim, method, cd_iter, channel_selection_
     np.testing.assert_allclose(gt_up_templates, match_up_templates, atol=2.5e-3)
 
     # difference between upsampling before going to multichan or after...
-    true_temps_up = gt_sorting._load_dataset("templates_up")
+    true_temps_up = gt_sorting.load_dataset("templates_up")
     true_temps_up = true_temps_up.astype(np.float32)
     np.testing.assert_allclose(gt_up_templates, true_temps_up, atol=1e-4)
     up_err = np.abs(gt_up_templates - true_temps_up).max() * (1 + 1e-5 + scaling)