Extend splitting-tests to multi-segment and mask labels

Author: alejoe91

src/spikeinterface/core/sorting_tools.py

@@ -469,16 +469,17 @@ def apply_splits_to_sorting(sorting, unit_splits, new_unit_ids=None, return_extr
     spike_vector_list = [spikes[s0:s1] for s0, s1 in segment_slices]
     spike_indices = spike_vector_to_indices(spike_vector_list, sorting.unit_ids, absolute_index=True)
 
-    # split_indices are a concatenation across segments
     for unit_id in sorting.unit_ids:
         if unit_id in unit_splits:
             split_indices = unit_splits[unit_id]
             new_split_ids = new_unit_ids[list(unit_splits.keys()).index(unit_id)]
 
             for split, new_unit_id in zip(split_indices, new_split_ids):
                 new_unit_index = all_unit_ids.index(new_unit_id)
+                # split_indices are a concatenation across segments with absolute indices
+                # so we need to concatenate the spike indices across segments
                 spike_indices_unit = np.concatenate(
-                    spike_indices[segment_index][unit_id] for segment_index in range(num_seg)
+                    [spike_indices[segment_index][unit_id] for segment_index in range(num_seg)]
                 )
                 spikes["unit_index"][spike_indices_unit[split]] = new_unit_index
         else:

src/spikeinterface/curation/tests/test_curation_format.py

@@ -127,7 +127,6 @@
 # Test dictionary format for merges with string IDs
 curation_ids_str_dict = {**curation_ids_str, "merges": {"u50": ["u3", "u6"], "u51": ["u10", "u14", "u20"]}}
 
-
 # This is a failure example with duplicated merge
 duplicate_merge = curation_ids_int.copy()
 duplicate_merge["merge_unit_groups"] = [[3, 6, 10], [10, 14, 20]]
@@ -292,11 +291,117 @@ def test_apply_curation_with_split():
         assert analyzer_curated.sorting.get_property("pyramidal", ids=[unit_id])[0]
 
 
+def test_apply_curation_with_split_multi_segment():
+    recording, sorting = generate_ground_truth_recording(durations=[10.0, 10.0], num_units=9, seed=2205)
+    sorting = sorting.rename_units(np.array([1, 2, 3, 6, 10, 14, 20, 31, 42]))
+    analyzer = create_sorting_analyzer(sorting, recording, sparse=False)
+    num_segments = sorting.get_num_segments()
+
+    curation_with_splits_multi_segment = curation_with_splits.copy()
+
+    # we make a split so that each subsplit will have all spikes from different segments
+    split_unit_id = curation_with_splits_multi_segment["splits"][0]["unit_id"]
+    sv = sorting.to_spike_vector()
+    unit_index = sorting.id_to_index(split_unit_id)
+    spikes_from_split_unit = sv[sv["unit_index"] == unit_index]
+
+    split_indices = []
+    cum_spikes = 0
+    for segment_index in range(num_segments):
+        spikes_in_segment = spikes_from_split_unit[spikes_from_split_unit["segment_index"] == segment_index]
+        split_indices.append(np.arange(0, len(spikes_in_segment)) + cum_spikes)
+        cum_spikes += len(spikes_in_segment)
+
+    curation_with_splits_multi_segment["splits"][0]["split_indices"] = split_indices
+
+    sorting_curated = apply_curation(sorting, curation_with_splits_multi_segment)
+
+    assert len(sorting_curated.unit_ids) == len(sorting.unit_ids) + 1
+    assert 2 not in sorting_curated.unit_ids
+    assert 43 in sorting_curated.unit_ids
+    assert 44 in sorting_curated.unit_ids
+
+    # check that spike trains are correctly split across segments
+    for seg_index in range(num_segments):
+        st_43 = sorting_curated.get_unit_spike_train(43, segment_index=seg_index)
+        st_44 = sorting_curated.get_unit_spike_train(44, segment_index=seg_index)
+        if seg_index == 0:
+            assert len(st_43) > 0
+            assert len(st_44) == 0
+        else:
+            assert len(st_43) == 0
+            assert len(st_44) > 0
+
+
+def test_apply_curation_splits_with_mask():
+    recording, sorting = generate_ground_truth_recording(durations=[10.0], num_units=9, seed=2205)
+    sorting = sorting.rename_units(np.array([1, 2, 3, 6, 10, 14, 20, 31, 42]))
+    analyzer = create_sorting_analyzer(sorting, recording, sparse=False)
+
+    # Get number of spikes for unit 2
+    num_spikes = sorting.count_num_spikes_per_unit()[2]
+
+    # Create split labels that assign spikes to 3 different clusters
+    split_labels = np.zeros(num_spikes, dtype=int)
+    split_labels[: num_spikes // 3] = 0  # First third to cluster 0
+    split_labels[num_spikes // 3 : 2 * num_spikes // 3] = 1  # Second third to cluster 1
+    split_labels[2 * num_spikes // 3 :] = 2  # Last third to cluster 2
+
+    curation_with_mask_split = {
+        "format_version": "2",
+        "unit_ids": [1, 2, 3, 6, 10, 14, 20, 31, 42],
+        "label_definitions": {
+            "quality": {"label_options": ["good", "noise", "MUA", "artifact"], "exclusive": True},
+            "putative_type": {
+                "label_options": ["excitatory", "inhibitory", "pyramidal", "mitral"],
+                "exclusive": False,
+            },
+        },
+        "manual_labels": [
+            {"unit_id": 2, "quality": ["good"], "putative_type": ["excitatory", "pyramidal"]},
+        ],
+        "splits": [
+            {
+                "unit_id": 2,
+                "split_mode": "labels",
+                "split_labels": split_labels.tolist(),
+                "split_new_unit_ids": [43, 44, 45],
+            }
+        ],
+    }
+
+    sorting_curated = apply_curation(sorting, curation_with_mask_split)
+
+    # Check results
+    assert len(sorting_curated.unit_ids) == len(sorting.unit_ids) + 2  # Original units - 1 (split) + 3 (new)
+    assert 2 not in sorting_curated.unit_ids  # Original unit should be removed
+
+    # Check new split units
+    split_unit_ids = [43, 44, 45]
+    for unit_id in split_unit_ids:
+        assert unit_id in sorting_curated.unit_ids
+        # Check properties are propagated
+        assert sorting_curated.get_property("quality", ids=[unit_id])[0] == "good"
+        assert sorting_curated.get_property("excitatory", ids=[unit_id])[0]
+        assert sorting_curated.get_property("pyramidal", ids=[unit_id])[0]
+
+    # Check analyzer
+    analyzer_curated = apply_curation(analyzer, curation_with_mask_split)
+    assert len(analyzer_curated.sorting.unit_ids) == len(analyzer.sorting.unit_ids) + 2
+
+    # Verify split sizes
+    spike_counts = analyzer_curated.sorting.count_num_spikes_per_unit()
+    assert spike_counts[43] == num_spikes // 3  # First third
+    assert spike_counts[44] == num_spikes // 3  # Second third
+    assert spike_counts[45] == num_spikes - 2 * (num_spikes // 3)  # Remainder
+
+
 if __name__ == "__main__":
-    # test_curation_format_validation()
-    # test_to_from_json()
-    # test_convert_from_sortingview_curation_format_v0()
-    # test_curation_label_to_vectors()
-    # test_curation_label_to_dataframe()
-    # test_apply_curation()
-    test_apply_curation_with_split()
+    test_curation_format_validation()
+    test_to_from_json()
+    test_convert_from_sortingview_curation_format_v0()
+    test_curation_label_to_vectors()
+    test_curation_label_to_dataframe()
+    test_apply_curation()
+    test_apply_curation_with_split_multi_segment()
+    test_apply_curation_splits_with_mask()