Repository: SpikeInterface/spikeinterface
Pull request #3208 (2024-07-15 18:57:20)

Start apply_curation()


Patch:
diff --git a/doc/api.rst b/doc/api.rst
index 1ac37e474..4d8bd2f32 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -324,15 +324,21 @@ spikeinterface.curation
 ------------------------
 .. automodule:: spikeinterface.curation
 
-    .. autoclass:: CurationSorting
-    .. autoclass:: MergeUnitsSorting
-    .. autoclass:: SplitUnitSorting
+    .. autofunction:: apply_curation
     .. autofunction:: get_potential_auto_merge
     .. autofunction:: find_redundant_units
     .. autofunction:: remove_redundant_units
     .. autofunction:: remove_duplicated_spikes
     .. autofunction:: remove_excess_spikes
+
+Deprecated
+~~~~~~~~~~
+.. automodule:: spikeinterface.curation
+
     .. autofunction:: apply_sortingview_curation
+    .. autoclass:: CurationSorting
+    .. autoclass:: MergeUnitsSorting
+    .. autoclass:: SplitUnitSorting
 
 
 spikeinterface.generation
diff --git a/doc/modules/curation.rst b/doc/modules/curation.rst
index 8e5418b66..d115b33e4 100644
--- a/doc/modules/curation.rst
+++ b/doc/modules/curation.rst
@@ -261,11 +261,23 @@ format is the definition; the second part of the format is manual action):
     }
 
 
-.. note::
-    The curation format was recently introduced (v0.101.0), and we are still working on
-    properly integrating it into the SpikeInterface ecosystem.
-    Soon there will be functions vailable, in the curation module, to apply this
-    standardized curation format to ``SortingAnalyzer`` and a ``BaseSorting`` objects.
+The curation format can be loaded into a dictionary and directly applied to
+a ``BaseSorting`` or ``SortingAnalyzer`` object using the :py:func:`~spikeinterface.curation.apply_curation` function.
+
+.. code-block:: python
+
+    from spikeinterface.curation import apply_curation
+
+    # load the curation JSON file
+    curation_json = "path/to/curation.json"
+    with open(curation_json, 'r') as f:
+        curation_dict = json.load(f)
+
+    # apply the curation to the sorting output
+    clean_sorting = apply_curation(sorting, curation_dict=curation_dict)
+
+    # apply the curation to the sorting analyzer
+    clean_sorting_analyzer = apply_curation(sorting_analyzer, curation_dict=curation_dict)
 
 
 Using the ``SpikeInterface GUI``
diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py
index 699457515..5f3335082 100644
--- a/src/spikeinterface/core/sorting_tools.py
+++ b/src/spikeinterface/core/sorting_tools.py
@@ -227,7 +227,7 @@ def random_spikes_selection(
 
 
 def apply_merges_to_sorting(
-    sorting, merge_unit_groups, new_unit_ids=None, censor_ms=None, return_kept=False, new_id_strategy="append"
+    sorting, merge_unit_groups, new_unit_ids=None, censor_ms=None, return_extra=False, new_id_strategy="append"
 ):
     """
     Apply a resolved representation of the merges to a sorting object.
@@ -250,8 +250,8 @@ def apply_merges_to_sorting(
         merged units will have the first unit_id of every lists of merges.
     censor_ms: float | None, default: None
         When applying the merges, should be discard consecutive spikes violating a given refractory per
-    return_kept : bool, default: False
-        If True, also return also a boolean mask of kept spikes.
+    return_extra : bool, default: False
+        If True, also return also a boolean mask of kept spikes and new_unit_ids.
     new_id_strategy : "append" | "take_first", default: "append"
         The strategy that should be used, if `new_unit_ids` is None, to create new unit_ids.
 
@@ -316,8 +316,8 @@ def apply_merges_to_sorting(
     spikes = spikes[keep_mask]
     sorting = NumpySorting(spikes, sorting.sampling_frequency, all_unit_ids)
 
-    if return_kept:
-        return sorting, keep_mask
+    if return_extra:
+        return sorting, keep_mask, new_unit_ids
     else:
         return sorting
 
@@ -384,11 +384,13 @@ def generate_unit_ids_for_merge_group(old_unit_ids, merge_unit_groups, new_unit_
     new_unit_ids : list | None, default: None
         Optional new unit_ids for merged units. If given, it needs to have the same length as `merge_unit_groups`.
         If None, new ids will be generated.
-    new_id_strategy : "append" | "take_first", default: "append"
+    new_id_strategy : "append" | "take_first" | "join", default: "append"
         The strategy that should be used, if `new_unit_ids` is None, to create new unit_ids.
 
             * "append" : new_units_ids will be added at the end of max(sorging.unit_ids)
             * "take_first" : new_unit_ids will be the first unit_id of every list of merges
+            * "join" : new_unit_ids will join unit_ids of groups with a "-".
+                       Only works if unit_ids are str otherwise switch to "append"
 
     Returns
     -------
@@ -423,6 +425,12 @@ def generate_unit_ids_for_merge_group(old_unit_ids, merge_unit_groups, new_unit_
             else:
                 # dtype int
                 new_unit_ids = list(max(old_unit_ids) + 1 + np.arange(num_merge, dtype=dtype))
+        elif new_id_strategy == "join":
+            if np.issubdtype(dtype, np.character):
+                new_unit_ids = ["-".join(group) for group in merge_unit_groups]
+            else:
+                # dtype int
+                new_unit_ids = list(max(old_unit_ids) + 1 + np.arange(num_merge, dtype=dtype))
         else:
             raise ValueError("wrong new_id_strategy")
 
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index 1f3479f13..ac142405a 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -732,12 +732,12 @@ class SortingAnalyzer:
         else:
             from spikeinterface.core.sorting_tools import apply_merges_to_sorting
 
-            sorting_provenance, keep_mask = apply_merges_to_sorting(
+            sorting_provenance, keep_mask, _ = apply_merges_to_sorting(
                 sorting=sorting_provenance,
                 merge_unit_groups=merge_unit_groups,
                 new_unit_ids=new_unit_ids,
                 censor_ms=censor_ms,
-                return_kept=True,
+                return_extra=True,
             )
             if censor_ms is None:
                 # in this case having keep_mask None is faster instead of having a vector of ones
@@ -885,6 +885,7 @@ class SortingAnalyzer:
         merging_mode="soft",
         sparsity_overlap=0.75,
         new_id_strategy="append",
+        return_new_unit_ids=False,
         format="memory",
         folder=None,
         verbose=False,
@@ -917,6 +918,8 @@ class SortingAnalyzer:
             The strategy that should be used, if `new_unit_ids` is None, to create new unit_ids.
                 * "append" : new_units_ids will be added at the end of max(sorting.unit_ids)
                 * "take_first" : new_unit_ids will be the first unit_id of every list of merges
+        return_new_unit_ids : bool, default False
+            Alse return new_unit_ids which are the ids of the new units.
         folder : Path | None, default: None
             The new folder where the analyzer with merged units is copied for `format` "binary_folder" or "zarr"
         format : "memory" | "binary_folder" | "zarr", default: "memory"
@@ -924,7 +927,6 @@ class SortingAnalyzer:
         verbose : bool, default: False
             Whether to display calculations (such as sparsity estimation)
 
-
         Returns
         -------
         analyzer :  SortingAnalyzer
@@ -952,7 +954,7 @@ class SortingAnalyzer:
         )
         all_unit_ids = _get_ids_after_merging(self.unit_ids, merge_unit_groups, new_unit_ids=new_unit_ids)
 
-        return self._save_or_select_or_merge(
+        new_analyzer = self._save_or_select_or_merge(
             format=format,
             folder=folder,
             merge_unit_groups=merge_unit_groups,
@@ -964,6 +966,10 @@ class SortingAnalyzer:
             new_unit_ids=new_unit_ids,
             **job_kwargs,
         )
+        if return_new_unit_ids:
+            return new_analyzer, new_unit_ids
+        else:
+            return new_analyzer
 
     def copy(self):
         """
diff --git a/src/spikeinterface/curation/curation_format.py b/src/spikeinterface/curation/curation_format.py
index fc75f7439..88190a9ba 100644
--- a/src/spikeinterface/curation/curation_format.py
+++ b/src/spikeinterface/curation/curation_format.py
@@ -1,5 +1,9 @@
 from itertools import combinations
 
+import numpy as np
+
+from spikeinterface.core import BaseSorting, SortingAnalyzer, apply_merges_to_sorting
+import copy
 
 supported_curation_format_versions = {"1"}
 
@@ -119,9 +123,9 @@ def convert_from_sortingview_curation_format_v0(sortingview_dict, destination_fo
     return curation_dict
 
 
-def curation_label_to_dataframe(curation_dict):
+def curation_label_to_vectors(curation_dict):
     """
-    Transform the curation dict into a pandas dataframe.
+    Transform the curation dict into dict of vectors.
     For label category with exclusive=True : a column is created and values are the unique label.
     For label category with exclusive=False : one column per possible is created and values are boolean.
 
@@ -134,30 +138,218 @@ def curation_label_to_dataframe(curation_dict):
 
     Returns
     -------
-    labels : pd.DataFrame
-        dataframe with labels.
+    labels: dict of numpy vector
+
     """
-    import pandas as pd
+    unit_ids = list(curation_dict["unit_ids"])
+    n = len(unit_ids)
 
-    labels = pd.DataFrame(index=curation_dict["unit_ids"])
+    labels = {}
 
     for label_key, label_def in curation_dict["label_definitions"].items():
         if label_def["exclusive"]:
-            assert label_key not in labels.columns, f"{label_key} is already a column"
-            labels[label_key] = pd.Series(dtype=str)
-            labels[label_key][:] = ""
+            assert label_key not in labels, f"{label_key} is already a key"
+            labels[label_key] = [""] * n
             for lbl in curation_dict["manual_labels"]:
                 value = lbl.get(label_key, [])
                 if len(value) == 1:
-                    labels.at[lbl["unit_id"], label_key] = value[0]
+                    unit_index = unit_ids.index(lbl["unit_id"])
+                    labels[label_key][unit_index] = value[0]
+            labels[label_key] = np.array(labels[label_key])
         else:
             for label_opt in label_def["label_options"]:
-                assert label_opt not in labels.columns, f"{label_opt} is already a column"
-                labels[label_opt] = pd.Series(dtype=bool)
-                labels[label_opt][:] = False
+                assert label_opt not in labels, f"{label_opt} is already a key"
+                labels[label_opt] = np.zeros(n, dtype=bool)
             for lbl in curation_dict["manual_labels"]:
                 values = lbl.get(label_key, [])
                 for value in values:
-                    labels.at[lbl["unit_id"], value] = True
+                    unit_index = unit_ids.index(lbl["unit_id"])
+                    labels[value][unit_index] = True
+
+    return labels
+
+
+def clean_curation_dict(curation_dict):
+    """
+    In some cases the curation_dict can have inconsistencies (like in the sorting view format).
+    For instance, some unit_ids are both in 'merge_unit_groups' and 'removed_units'.
+    This is ambiguous!
+
+    This cleaner helper function ensures units tagged as `removed_units` are removed from the `merge_unit_groups`
+    """
+    curation_dict = copy.deepcopy(curation_dict)
+
+    clean_merge_unit_groups = []
+    for group in curation_dict["merge_unit_groups"]:
+        clean_group = []
+        for unit_id in group:
+            if unit_id not in curation_dict["removed_units"]:
+                clean_group.append(unit_id)
+        if len(clean_group) > 1:
+            clean_merge_unit_groups.append(clean_group)
+
+    curation_dict["merge_unit_groups"] = clean_merge_unit_groups
+    return curation_dict
+
+
+def curation_label_to_dataframe(curation_dict):
+    """
+    Transform the curation dict into a pandas dataframe.
+    For label category with exclusive=True : a column is created and values are the unique label.
+    For label category with exclusive=False : one column per possible is created and values are boolean.
+
+    If exclusive=False and the same label appears several times then an error is raised.
+
+    Parameters
+    ----------
+    curation_dict : dict
+        A curation dictionary
+
+    Returns
+    -------
+    labels : pd.DataFrame
+        dataframe with labels.
+    """
+    import pandas as pd
 
+    labels = pd.DataFrame(curation_label_to_vectors(curation_dict), index=curation_dict["unit_ids"])
     return labels
+
+
+def apply_curation_labels(sorting, new_unit_ids, curation_dict):
+    """
+    Apply manual labels after merges.
+
+    Rules:
+      * label for non merge is applied first
+      * for merged group, when exclusive=True, if all have the same label then this label is applied
+      * for merged group, when exclusive=False, if one unit has the label then the new one have also it
+    """
+
+    # Please note that manual_labels is done on the unit_ids before the merge!!!
+    manual_labels = curation_label_to_vectors(curation_dict)
+
+    # apply on non merged
+    for key, values in manual_labels.items():
+        all_values = np.zeros(sorting.unit_ids.size, dtype=values.dtype)
+        for unit_ind, unit_id in enumerate(sorting.unit_ids):
+            if unit_id not in new_unit_ids:
+                ind = curation_dict["unit_ids"].index(unit_id)
+                all_values[unit_ind] = values[ind]
+        sorting.set_property(key, all_values)
+
+    for new_unit_id, old_group_ids in zip(new_unit_ids, curation_dict["merge_unit_groups"]):
+        for label_key, label_def in curation_dict["label_definitions"].items():
+            if label_def["exclusive"]:
+                group_values = []
+                for unit_id in old_group_ids:
+                    ind = curation_dict["unit_ids"].index(unit_id)
+                    value = manual_labels[label_key][ind]
+                    if value != "":
+                        group_values.append(value)
+                if len(set(group_values)) == 1:
+                    # all group has the same label or empty
+                    sorting.set_property(key, values=group_values, ids=[new_unit_id])
+            else:
+
+                for key in label_def["label_options"]:
+                    group_values = []
+                    for unit_id in old_group_ids:
+                        ind = curation_dict["unit_ids"].index(unit_id)
+                        value = manual_labels[key][ind]
+                        group_values.append(value)
+                    new_value = np.any(group_values)
+                    sorting.set_property(key, values=[new_value], ids=[new_unit_id])
+
+
+def apply_curation(
+    sorting_or_analyzer,
+    curation_dict,
+    censor_ms=None,
+    new_id_strategy="append",
+    merging_mode="soft",
+    sparsity_overlap=0.75,
+    verbose=False,
+    **job_kwargs,
+):
+    """
+    Apply curation dict to a Sorting or a SortingAnalyzer.
+
+    Steps are done in this order:
+      1. Apply removal using curation_dict["removed_units"]
+      2. Apply merges using curation_dict["merge_unit_groups"]
+      3. Set labels using curation_dict["manual_labels"]
+
+    A new Sorting or SortingAnalyzer (in memory) is returned.
+    The user (an adult) has the responsability to save it somewhere (or not).
+
+    Parameters
+    ----------
+    sorting_or_analyzer : Sorting | SortingAnalyzer
+        The Sorting object to apply merges.
+    curation_dict : dict
+        The curation dict.
+    censor_ms: float | None, default: None
+        When applying the merges, any consecutive spikes within the `censor_ms` are removed. This can be thought of
+        as the desired refractory period. If `censor_ms=None`, no spikes are discarded.
+    new_id_strategy : "append" | "take_first", default: "append"
+        The strategy that should be used, if `new_unit_ids` is None, to create new unit_ids.
+
+            * "append" : new_units_ids will be added at the end of max(sorting.unit_ids)
+            * "take_first" : new_unit_ids will be the first unit_id of every list of merges
+    merging_mode  : "soft" | "hard", default: "soft"
+        How merges are performed for SortingAnalyzer. If the `merge_mode` is "soft" , merges will be approximated, with no reloading of
+        the waveforms. This will lead to approximations. If `merge_mode` is "hard", recomputations are accurately
+        performed, reloading waveforms if needed
+    sparsity_overlap : float, default 0.75
+            The percentage of overlap that units should share in order to accept merges. If this criteria is not
+            achieved, soft merging will not be possible and an error will be raised. This is for use with a SortingAnalyzer input.
+
+    verbose:
+
+    **job_kwargs
+
+    Returns
+    -------
+    sorting_or_analyzer : Sorting | SortingAnalyzer
+        The curated object.
+
+
+    """
+    validate_curation_dict(curation_dict)
+    if not np.array_equal(np.asarray(curation_dict["unit_ids"]), sorting_or_analyzer.unit_ids):
+        raise ValueError("unit_ids from the curation_dict do not match the one from Sorting or SortingAnalyzer")
+
+    if isinstance(sorting_or_analyzer, BaseSorting):
+        sorting = sorting_or_analyzer
+        sorting = sorting.remove_units(curation_dict["removed_units"])
+        sorting, _, new_unit_ids = apply_merges_to_sorting(
+            sorting,
+            curation_dict["merge_unit_groups"],
+            censor_ms=censor_ms,
+            return_extra=True,
+            new_id_strategy=new_id_strategy,
+        )
+        apply_curation_labels(sorting, new_unit_ids, curation_dict)
+        return sorting
+
+    elif isinstance(sorting_or_analyzer, SortingAnalyzer):
+        analyzer = sorting_or_analyzer
+        analyzer = analyzer.remove_units(curation_dict["removed_units"])
+        analyzer, new_unit_ids = analyzer.merge_units(
+            curation_dict["merge_unit_groups"],
+            censor_ms=censor_ms,
+            merging_mode=merging_mode,
+            sparsity_overlap=sparsity_overlap,
+            new_id_strategy=new_id_strategy,
+            return_new_unit_ids=True,
+            format="memory",
+            verbose=verbose,
+            **job_kwargs,
+        )
+        apply_curation_labels(analyzer.sorting, new_unit_ids, curation_dict)
+        return analyzer
+    else:
+        raise TypeError(
+            f"`sorting_or_analyzer` must be a Sorting or a SortingAnalyzer, not an object of type {type(sorting_or_analyzer)}"
+        )
diff --git a/src/spikeinterface/curation/sortingview_curation.py b/src/spikeinterface/curation/sortingview_curation.py
index c4d2a3295..1e36183ee 100644
--- a/src/spikeinterface/curation/sortingview_curation.py
+++ b/src/spikeinterface/curation/sortingview_curation.py
@@ -1,14 +1,116 @@
 from __future__ import annotations
+
+import warnings
+
 import json
 import numpy as np
 from pathlib import Path
 
 from .curationsorting import CurationSorting
+from .curation_format import (
+    convert_from_sortingview_curation_format_v0,
+    apply_curation,
+    curation_label_to_vectors,
+    clean_curation_dict,
+)
 
 
-# @alessio
-# TODO later : this should be reimplemented using the new curation format
 def apply_sortingview_curation(
+    sorting_or_analyzer, uri_or_json, exclude_labels=None, include_labels=None, skip_merge=False, verbose=None
+):
+    """
+    Apply curation from SortingView manual legacy curation format (before the official "curation_format")
+
+    First, merges (if present) are applied. Then labels are loaded and units
+    are optionally filtered based on exclude_labels and include_labels.
+
+    Parameters
+    ----------
+    sorting_or_analyzer : Sorting | SortingAnalyzer
+        The sorting or analyzer to be curated
+    uri_or_json : str or Path
+        The URI curation link from SortingView or the path to the curation json file
+    exclude_labels : list, default: None
+        Optional list of labels to exclude (e.g. ["reject", "noise"]).
+        Mutually exclusive with include_labels
+    include_labels : list, default: None
+        Optional list of labels to include (e.g. ["accept"]).
+        Mutually exclusive with exclude_labels,  by default None
+    skip_merge : bool, default: False
+        If True, merges are not applied (only labels)
+    verbose : None
+        Deprecated
+
+
+    Returns
+    -------
+    sorting_or_analyzer_curated : BaseSorting
+        The curated sorting or analyzer
+    """
+
+    if verbose is not None:
+        warnings.warn("versobe in apply_sortingview_curation() is deprecated")
+
+    # download
+    if Path(uri_or_json).suffix == ".json" and not str(uri_or_json).startswith("gh://"):
+        with open(uri_or_json, "r") as f:
+            curation_dict = json.load(f)
+    else:
+        try:
+            import kachery_cloud as kcl
+        except ImportError:
+            raise ImportError(
+                "To apply a SortingView manual curation, you need to have sortingview installed: "
+                ">>> pip install sortingview"
+            )
+
+        try:
+            curation_dict = kcl.load_json(uri=uri_or_json)
+        except:
+            raise Exception(f"Could not retrieve curation from SortingView uri: {uri_or_json}")
+
+    # convert to new format
+    if "format_version" not in curation_dict:
+        curation_dict = convert_from_sortingview_curation_format_v0(curation_dict)
+
+    unit_ids = sorting_or_analyzer.unit_ids
+
+    # this is a hack because it was not in the old format
+    curation_dict["unit_ids"] = list(unit_ids)
+
+    if exclude_labels is not None:
+        assert include_labels is None, "Use either `include_labels` or `exclude_labels` to filter units."
+        manual_labels = curation_label_to_vectors(curation_dict)
+        removed_units = []
+        for k in exclude_labels:
+            remove_mask = manual_labels[k]
+            removed_units.extend(unit_ids[remove_mask])
+        removed_units = np.unique(removed_units)
+        curation_dict["removed_units"] = removed_units
+
+    if include_labels is not None:
+        manual_labels = curation_label_to_vectors(curation_dict)
+        removed_units = []
+        for k in include_labels:
+            remove_mask = ~manual_labels[k]
+            removed_units.extend(unit_ids[remove_mask])
+        removed_units = np.unique(removed_units)
+        curation_dict["removed_units"] = removed_units
+
+    if skip_merge:
+        curation_dict["merge_unit_groups"] = []
+
+    # cleaner to ensure validity
+    curation_dict = clean_curation_dict(curation_dict)
+
+    # apply
+    sorting_curated = apply_curation(sorting_or_analyzer, curation_dict, new_id_strategy="join")
+
+    return sorting_curated
+
+
+# TODO @alessio you remove this after testing
+def apply_sortingview_curation_legacy(
     sorting, uri_or_json, exclude_labels=None, include_labels=None, skip_merge=False, verbose=False
 ):
     """


Test patch:
diff --git a/src/spikeinterface/core/tests/test_sorting_tools.py b/src/spikeinterface/core/tests/test_sorting_tools.py
index 6d0e61f84..34bb3a221 100644
--- a/src/spikeinterface/core/tests/test_sorting_tools.py
+++ b/src/spikeinterface/core/tests/test_sorting_tools.py
@@ -96,7 +96,7 @@ def test_apply_merges_to_sorting():
         spikes1[spikes1["unit_index"] == 2]["sample_index"], spikes2[spikes2["unit_index"] == 0]["sample_index"]
     )
 
-    sorting3, keep_mask = apply_merges_to_sorting(sorting1, [["a", "b"]], censor_ms=1.5, return_kept=True)
+    sorting3, keep_mask, _ = apply_merges_to_sorting(sorting1, [["a", "b"]], censor_ms=1.5, return_extra=True)
     spikes3 = sorting3.to_spike_vector()
     assert spikes3.size < spikes1.size
     assert not keep_mask[1]
@@ -153,6 +153,11 @@ def test_generate_unit_ids_for_merge_group():
     )
     assert np.array_equal(new_unit_ids, ["0", "9"])
 
+    new_unit_ids = generate_unit_ids_for_merge_group(
+        ["0", "5", "12", "9", "15"], [["0", "5"], ["9", "15"]], new_id_strategy="join"
+    )
+    assert np.array_equal(new_unit_ids, ["0-5", "9-15"])
+
 
 if __name__ == "__main__":
     # test_spike_vector_to_spike_trains()
diff --git a/src/spikeinterface/curation/tests/test_curation_format.py b/src/spikeinterface/curation/tests/test_curation_format.py
index 94812ee0a..af9d8e1ea 100644
--- a/src/spikeinterface/curation/tests/test_curation_format.py
+++ b/src/spikeinterface/curation/tests/test_curation_format.py
@@ -2,11 +2,16 @@ import pytest
 
 from pathlib import Path
 import json
+import numpy as np
+
+from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer
 
 from spikeinterface.curation.curation_format import (
     validate_curation_dict,
     convert_from_sortingview_curation_format_v0,
+    curation_label_to_vectors,
     curation_label_to_dataframe,
+    apply_curation,
 )
 
 
@@ -141,6 +146,17 @@ def test_convert_from_sortingview_curation_format_v0():
             validate_curation_dict(curation_v1)
 
 
+def test_curation_label_to_vectors():
+
+    labels = curation_label_to_vectors(curation_ids_int)
+    assert "quality" in labels
+    assert "excitatory" in labels
+    print(labels)
+
+    labels = curation_label_to_vectors(curation_ids_str)
+    print(labels)
+
+
 def test_curation_label_to_dataframe():
 
     df = curation_label_to_dataframe(curation_ids_int)
@@ -152,10 +168,25 @@ def test_curation_label_to_dataframe():
     # print(df)
 
 
+def test_apply_curation():
+    recording, sorting = generate_ground_truth_recording(durations=[10.0], num_units=9, seed=2205)
+    sorting._main_ids = np.array([1, 2, 3, 6, 10, 14, 20, 31, 42])
+    analyzer = create_sorting_analyzer(sorting, recording, sparse=False)
+
+    sorting_curated = apply_curation(sorting, curation_ids_int)
+    assert sorting_curated.get_property("quality", ids=[1])[0] == "good"
+    assert sorting_curated.get_property("quality", ids=[2])[0] == "noise"
+    assert sorting_curated.get_property("excitatory", ids=[2])[0]
+
+    analyzer_curated = apply_curation(analyzer, curation_ids_int)
+    assert "quality" in analyzer_curated.sorting.get_property_keys()
+
+
 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()
 
-    print(json.dumps(curation_ids_str, indent=4))
+    test_apply_curation()
diff --git a/src/spikeinterface/curation/tests/test_sortingview_curation.py b/src/spikeinterface/curation/tests/test_sortingview_curation.py
index 00721ff34..bb152e7f7 100644
--- a/src/spikeinterface/curation/tests/test_sortingview_curation.py
+++ b/src/spikeinterface/curation/tests/test_sortingview_curation.py
@@ -52,7 +52,7 @@ def test_gh_curation():
     # curated link:
     # https://figurl.org/f?v=npm://@fi-sci/figurl-sortingview@12/dist&d=sha1://058ab901610aa9d29df565595a3cc2a81a1b08e5
     gh_uri = "gh://SpikeInterface/spikeinterface/main/src/spikeinterface/curation/tests/sv-sorting-curation.json"
-    sorting_curated_gh = apply_sortingview_curation(sorting, uri_or_json=gh_uri, verbose=True)
+    sorting_curated_gh = apply_sortingview_curation(sorting, uri_or_json=gh_uri)
 
     assert len(sorting_curated_gh.unit_ids) == 9
     assert 1, 2 in sorting_curated_gh.unit_ids
@@ -81,7 +81,7 @@ def test_sha1_curation():
     # curated link:
     # https://figurl.org/f?v=npm://@fi-sci/figurl-sortingview@12/dist&d=sha1://058ab901610aa9d29df565595a3cc2a81a1b08e5
     sha1_uri = "sha1://449a428e8824eef9ad9bcc3241e45a2cee02d381"
-    sorting_curated_sha1 = apply_sortingview_curation(sorting, uri_or_json=sha1_uri, verbose=True)
+    sorting_curated_sha1 = apply_sortingview_curation(sorting, uri_or_json=sha1_uri)
     # print(f"From SHA: {sorting_curated_sha1}")
 
     assert len(sorting_curated_sha1.unit_ids) == 9
@@ -109,7 +109,7 @@ def test_json_curation():
     # from curation.json
     json_file = parent_folder / "sv-sorting-curation.json"
     # print(f"Sorting: {sorting.get_unit_ids()}")
-    sorting_curated_json = apply_sortingview_curation(sorting, uri_or_json=json_file, verbose=True)
+    sorting_curated_json = apply_sortingview_curation(sorting, uri_or_json=json_file)
 
     assert len(sorting_curated_json.unit_ids) == 9
     assert 1, 2 in sorting_curated_json.unit_ids
@@ -146,7 +146,7 @@ def test_false_positive_curation():
     # print("Sorting: {}".format(sorting.get_unit_ids()))
 
     json_file = parent_folder / "sv-sorting-curation-false-positive.json"
-    sorting_curated_json = apply_sortingview_curation(sorting, uri_or_json=json_file, verbose=True)
+    sorting_curated_json = apply_sortingview_curation(sorting, uri_or_json=json_file)
     # print("Curated:", sorting_curated_json.get_unit_ids())
 
     # Assertions
@@ -190,17 +190,13 @@ def test_label_inheritance_int():
     assert not sorting_merge.get_unit_property(unit_id=10, key="noise")
     assert sorting_merge.get_unit_property(unit_id=10, key="accept")
 
-    # Assertions for exclude_labels
+    # Assertions for exclude_labels should all be False
     sorting_exclude_noise = apply_sortingview_curation(sorting, uri_or_json=json_file, exclude_labels=["noise"])
-    # print(f"Exclude noise: {sorting_exclude_noise.get_unit_ids()}")
-    assert 9 not in sorting_exclude_noise.get_unit_ids()
+    assert np.all(~sorting_exclude_noise.get_property("noise"))
 
     # Assertions for include_labels
     sorting_include_accept = apply_sortingview_curation(sorting, uri_or_json=json_file, include_labels=["accept"])
-    # print(f"Include accept: {sorting_include_accept.get_unit_ids()}")
-    assert 8 not in sorting_include_accept.get_unit_ids()
-    assert 9 not in sorting_include_accept.get_unit_ids()
-    assert 10 in sorting_include_accept.get_unit_ids()
+    assert np.all(sorting_include_accept.get_property("accept"))
 
 
 def test_label_inheritance_str():
@@ -219,7 +215,7 @@ def test_label_inheritance_str():
 
     # Apply curation
     json_file = parent_folder / "sv-sorting-curation-str.json"
-    sorting_merge = apply_sortingview_curation(sorting, uri_or_json=json_file, verbose=True)
+    sorting_merge = apply_sortingview_curation(sorting, uri_or_json=json_file)
 
     # Assertions for merged units
     # print(f"Merge only: {sorting_merge.get_unit_ids()}")
@@ -238,22 +234,18 @@ def test_label_inheritance_str():
     assert not sorting_merge.get_unit_property(unit_id="e-f", key="noise")
     assert sorting_merge.get_unit_property(unit_id="e-f", key="accept")
 
-    # Assertions for exclude_labels
+    # Assertions for exclude_labels should all be False
     sorting_exclude_noise = apply_sortingview_curation(sorting, uri_or_json=json_file, exclude_labels=["noise"])
-    # print(f"Exclude noise: {sorting_exclude_noise.get_unit_ids()}")
-    assert "c-d" not in sorting_exclude_noise.get_unit_ids()
+    assert np.all(~sorting_exclude_noise.get_property("noise"))
 
     # Assertions for include_labels
     sorting_include_accept = apply_sortingview_curation(sorting, uri_or_json=json_file, include_labels=["accept"])
-    # print(f"Include accept: {sorting_include_accept.get_unit_ids()}")
-    assert "a-b" not in sorting_include_accept.get_unit_ids()
-    assert "c-d" not in sorting_include_accept.get_unit_ids()
-    assert "e-f" in sorting_include_accept.get_unit_ids()
+    assert np.all(sorting_include_accept.get_property("accept"))
 
 
 if __name__ == "__main__":
     # generate_sortingview_curation_dataset()
-    test_sha1_curation()
+    # test_sha1_curation()
     test_gh_curation()
     test_json_curation()
     test_false_positive_curation()