Rename CurationModel to Curation (#4421)

Author: alejoe91

doc/api.rst

@@ -390,7 +390,7 @@ such as merging and splitting units, as well as defining labels for units.
 
 .. automodule:: spikeinterface.curation.curation_model
 
-    .. autopydantic_model:: CurationModel
+    .. autopydantic_model:: Curation
     .. autopydantic_model:: Merge
     .. autopydantic_model:: Split
     .. autopydantic_model:: ManualLabel

src/spikeinterface/curation/curation_format.py

@@ -4,7 +4,7 @@
 from itertools import chain
 
 from spikeinterface.core import BaseSorting, SortingAnalyzer, apply_merges_to_sorting, apply_splits_to_sorting
-from spikeinterface.curation.curation_model import CurationModel, SequentialCuration
+from spikeinterface.curation.curation_model import Curation, SequentialCuration
 
 
 def validate_curation_dict(curation_dict: dict):
@@ -19,10 +19,10 @@ def validate_curation_dict(curation_dict: dict):
 
     """
     # this will validate the format of the curation_dict
-    CurationModel(**curation_dict)
+    Curation(**curation_dict)
 
 
-def curation_label_to_vectors(curation_dict_or_model: dict | CurationModel):
+def curation_label_to_vectors(curation_dict_or_model: dict | Curation):
     """
     Transform the curation dict into dict of vectors.
     For label category with exclusive=True : a column is created and values are the unique label.
@@ -32,7 +32,7 @@ def curation_label_to_vectors(curation_dict_or_model: dict | CurationModel):
 
     Parameters
     ----------
-    curation_dict : dict or CurationModel
+    curation_dict : dict or Curation
         A curation dictionary or model
 
     Returns
@@ -41,7 +41,7 @@ def curation_label_to_vectors(curation_dict_or_model: dict | CurationModel):
 
     """
     if isinstance(curation_dict_or_model, dict):
-        curation_model = CurationModel(**curation_dict_or_model)
+        curation_model = Curation(**curation_dict_or_model)
     else:
         curation_model = curation_dict_or_model
     unit_ids = list(curation_model.unit_ids)
@@ -71,7 +71,7 @@ def curation_label_to_vectors(curation_dict_or_model: dict | CurationModel):
     return labels
 
 
-def curation_label_to_dataframe(curation_dict_or_model: dict | CurationModel):
+def curation_label_to_dataframe(curation_dict_or_model: dict | Curation):
     """
     Transform the curation dict into a pandas dataframe.
     For label category with exclusive=True : a column is created and values are the unique label.
@@ -92,17 +92,15 @@ def curation_label_to_dataframe(curation_dict_or_model: dict | CurationModel):
     import pandas as pd
 
     if isinstance(curation_dict_or_model, dict):
-        curation_model = CurationModel(**curation_dict_or_model)
+        curation_model = Curation(**curation_dict_or_model)
     else:
         curation_model = curation_dict_or_model
 
     labels = pd.DataFrame(curation_label_to_vectors(curation_model), index=curation_model.unit_ids)
     return labels
 
 
-def apply_curation_labels(
-    sorting_or_analyzer: BaseSorting | SortingAnalyzer, curation_dict_or_model: dict | CurationModel
-):
+def apply_curation_labels(sorting_or_analyzer: BaseSorting | SortingAnalyzer, curation_dict_or_model: dict | Curation):
     """
     Apply manual labels after merges/splits.
 
@@ -113,7 +111,7 @@ def apply_curation_labels(
       * for split units, the original label is applied to all split units
     """
     if isinstance(curation_dict_or_model, dict):
-        curation_model = CurationModel(**curation_dict_or_model)
+        curation_model = Curation(**curation_dict_or_model)
     else:
         curation_model = curation_dict_or_model
 
@@ -136,7 +134,7 @@ def apply_curation_labels(
 
 def apply_curation(
     sorting_or_analyzer: BaseSorting | SortingAnalyzer,
-    curation_dict_or_model: dict | list | CurationModel | SequentialCuration,
+    curation_dict_or_model: dict | list | Curation | SequentialCuration,
     censor_ms: float | None = None,
     new_id_strategy: str = "append",
     merging_mode: str = "soft",
@@ -162,7 +160,7 @@ def apply_curation(
     ----------
     sorting_or_analyzer : Sorting | SortingAnalyzer
         The Sorting or SortingAnalyzer object to apply merges.
-    curation_dict : dict | CurationModel | SequentialCuration
+    curation_dict : dict | Curation | SequentialCuration
         The curation dict or model.
     censor_ms : float | None, default: None
         When applying the merges, any consecutive spikes within the `censor_ms` are removed. This can be thought of
@@ -197,10 +195,10 @@ def apply_curation(
         sorting_or_analyzer, (BaseSorting, SortingAnalyzer)
     ), f"`sorting_or_analyzer` must be a Sorting or a SortingAnalyzer, not an object of type {type(sorting_or_analyzer)}"
     assert isinstance(
-        curation_dict_or_model, (dict, list, CurationModel, SequentialCuration)
-    ), f"`curation_dict_or_model` must be a dict, CurationModel or a SequentialCuration not an object of type {type(curation_dict_or_model)}"
+        curation_dict_or_model, (dict, list, Curation, SequentialCuration)
+    ), f"`curation_dict_or_model` must be a dict, Curation or a SequentialCuration not an object of type {type(curation_dict_or_model)}"
     if isinstance(curation_dict_or_model, dict):
-        curation_model = CurationModel(**curation_dict_or_model)
+        curation_model = Curation(**curation_dict_or_model)
     elif isinstance(curation_dict_or_model, list):
         curation_model = SequentialCuration(curation_steps=curation_dict_or_model)
     else:
@@ -298,7 +296,7 @@ def apply_curation(
     return curated_sorting_or_analyzer
 
 
-def load_curation(curation_path: str | Path) -> CurationModel:
+def load_curation(curation_path: str | Path) -> Curation:
     """
     Loads a curation from a local json file.
 
@@ -309,9 +307,9 @@ def load_curation(curation_path: str | Path) -> CurationModel:
 
     Returns
     -------
-    curation_model : CurationModel
-        A CurationModel object
+    curation_model : Curation
+        A Curation object
     """
     with open(curation_path) as f:
         curation_dict = json.load(f)
-    return CurationModel(**curation_dict)
+    return Curation(**curation_dict)

src/spikeinterface/curation/curation_model.py

@@ -1,6 +1,8 @@
-from pydantic import BaseModel, Field, model_validator, field_validator, field_serializer
+import warnings
 from typing import Literal, List
 from itertools import chain, combinations
+from pydantic import BaseModel, Field, model_validator, field_validator, field_serializer
+
 import numpy as np
 
 from spikeinterface import BaseSorting
@@ -74,7 +76,7 @@ def get_full_spike_indices(self, sorting: BaseSorting):
         return full_spike_indices
 
 
-class CurationModel(BaseModel):
+class Curation(BaseModel):
     supported_versions: tuple[Literal["1"], Literal["2"]] = Field(
         default=("1", "2"), description="Supported versions of the curation format"
     )
@@ -477,14 +479,23 @@ def validate_curation_dict(self):
         return self
 
 
+def CurationModel(*args, **kwargs):
+    warnings.warn(
+        "`CurationModel` is deprecated and will be removed in 0.105.0. Use `Curation` instead",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return Curation(*args, **kwargs)
+
+
 class SequentialCuration(BaseModel):
     """
     A Pydantic model which defines a sequence of curation steps. If using sequential curations,
     we demand that each individual curation (except the final one) has manually defined new unit ids,
     and that these match the unit ids of the following curation.
     """
 
-    curation_steps: List[CurationModel] = Field(description="List of curation steps applied sequentially")
+    curation_steps: List[Curation] = Field(description="List of curation steps applied sequentially")
 
     @model_validator(mode="after")
     def validate_sequential_curation(self):

src/spikeinterface/curation/sortingview_curation.py

@@ -9,7 +9,7 @@
     apply_curation,
     curation_label_to_vectors,
 )
-from .curation_model import CurationModel, Merge
+from .curation_model import Curation, Merge
 
 
 def get_kachery():
@@ -83,7 +83,7 @@ def apply_sortingview_curation(
 
     unit_ids = sorting_or_analyzer.unit_ids
     curation_dict["unit_ids"] = unit_ids
-    curation_model = CurationModel(**curation_dict)
+    curation_model = Curation(**curation_dict)
 
     if skip_merge:
         curation_model.merges = []

src/spikeinterface/curation/tests/test_curation_model.py

@@ -3,19 +3,19 @@
 from pydantic import ValidationError
 import numpy as np
 
-from spikeinterface.curation.curation_model import CurationModel, SequentialCuration, LabelDefinition
+from spikeinterface.curation.curation_model import Curation, SequentialCuration, LabelDefinition
 
 
 # Test data for format version
 def test_format_version():
     # Valid format version
-    CurationModel(format_version="1", unit_ids=[1, 2, 3])
+    Curation(format_version="1", unit_ids=[1, 2, 3])
 
     # Invalid format version
     with pytest.raises(ValidationError):
-        CurationModel(format_version="3", unit_ids=[1, 2, 3])
+        Curation(format_version="3", unit_ids=[1, 2, 3])
     with pytest.raises(ValidationError):
-        CurationModel(format_version="0.1", unit_ids=[1, 2, 3])
+        Curation(format_version="0.1", unit_ids=[1, 2, 3])
 
 
 # Test data for label definitions
@@ -29,7 +29,7 @@ def test_label_definitions():
         },
     }
 
-    model = CurationModel(**valid_label_def)
+    model = Curation(**valid_label_def)
     assert "quality" in model.label_definitions
     assert model.label_definitions["quality"].name == "quality"
     assert model.label_definitions["quality"].exclusive is True
@@ -54,7 +54,7 @@ def test_manual_labels():
         ],
     }
 
-    model = CurationModel(**valid_labels)
+    model = Curation(**valid_labels)
     assert len(model.manual_labels) == 2
 
     # Test invalid unit ID
@@ -67,7 +67,7 @@ def test_manual_labels():
         "manual_labels": [{"unit_id": 4, "labels": {"quality": ["good"]}}],  # Non-existent unit
     }
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_unit)
+        Curation(**invalid_unit)
 
     # Test violation of exclusive label
     invalid_exclusive = {
@@ -81,7 +81,7 @@ def test_manual_labels():
         ],
     }
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_exclusive)
+        Curation(**invalid_exclusive)
 
 
 # Test merge functionality
@@ -96,15 +96,15 @@ def test_merge_units():
         ],
     }
 
-    model = CurationModel(**valid_merge)
+    model = Curation(**valid_merge)
     assert len(model.merges) == 2
     assert model.merges[0].new_unit_id == 5
     assert model.merges[1].new_unit_id == 6
 
     # Test dictionary format
     valid_merge_dict = {"format_version": "2", "unit_ids": [1, 2, 3, 4], "merges": {5: [1, 2], 6: [3, 4]}}
 
-    model = CurationModel(**valid_merge_dict)
+    model = Curation(**valid_merge_dict)
     assert len(model.merges) == 2
     merge_new_ids = {merge.new_unit_id for merge in model.merges}
     assert merge_new_ids == {5, 6}
@@ -115,7 +115,7 @@ def test_merge_units():
         "unit_ids": [1, 2, 3, 4],
         "merges": [[1, 2], [3, 4]],  # Merge each pair into a new unit
     }
-    model = CurationModel(**valid_merge_list)
+    model = Curation(**valid_merge_list)
     assert len(model.merges) == 2
 
     # Test invalid merge group (single unit)
@@ -125,7 +125,7 @@ def test_merge_units():
         "merges": [{"unit_ids": [1], "new_unit_id": 4}],
     }
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_merge_group)
+        Curation(**invalid_merge_group)
 
     # Test overlapping merge groups
     invalid_overlap = {
@@ -137,7 +137,7 @@ def test_merge_units():
         ],
     }
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_overlap)
+        Curation(**invalid_overlap)
 
 
 # Test split functionality
@@ -156,7 +156,7 @@ def test_split_units():
         ],
     }
 
-    model = CurationModel(**valid_split_indices)
+    model = Curation(**valid_split_indices)
     assert len(model.splits) == 1
     assert model.splits[0].mode == "indices"
     assert len(model.splits[0].indices) == 2
@@ -168,7 +168,7 @@ def test_split_units():
         "splits": [{"unit_id": 1, "mode": "labels", "labels": [0, 0, 1, 1, 0, 2], "new_unit_ids": [4, 5, 6]}],
     }
 
-    model = CurationModel(**valid_split_labels)
+    model = Curation(**valid_split_labels)
     assert len(model.splits) == 1
     assert model.splits[0].mode == "labels"
     assert len(set(model.splits[0].labels)) == 3
@@ -183,7 +183,7 @@ def test_split_units():
         },
     }
 
-    model = CurationModel(**valid_split_dict)
+    model = Curation(**valid_split_dict)
     assert len(model.splits) == 2
     assert all(split.mode == "indices" for split in model.splits)
 
@@ -194,7 +194,7 @@ def test_split_units():
         "splits": [{"unit_id": 4, "mode": "indices", "indices": [[0, 1], [2, 3]]}],  # Non-existent unit
     }
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_unit_id)
+        Curation(**invalid_unit_id)
 
     # Test invalid new unit IDs count for indices mode
     invalid_new_ids = {
@@ -210,20 +210,20 @@ def test_split_units():
         ],
     }
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_new_ids)
+        Curation(**invalid_new_ids)
 
 
 # Test removed units
 def test_removed_units():
     valid_remove = {"format_version": "2", "unit_ids": [1, 2, 3], "removed": [2]}
 
-    model = CurationModel(**valid_remove)
+    model = Curation(**valid_remove)
     assert len(model.removed) == 1
 
     # Test removing non-existent unit
     invalid_remove = {"format_version": "2", "unit_ids": [1, 2, 3], "removed": [4]}  # Non-existent unit
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_remove)
+        Curation(**invalid_remove)
 
     # Test conflict between merge and remove
     invalid_merge_remove = {
@@ -233,7 +233,7 @@ def test_removed_units():
         "removed": [1],  # Unit is both merged and removed
     }
     with pytest.raises(ValidationError):
-        CurationModel(**invalid_merge_remove)
+        Curation(**invalid_merge_remove)
 
 
 # Test complete model with multiple operations
@@ -251,7 +251,7 @@ def test_complete_model():
         "removed": [5],
     }
 
-    model = CurationModel(**complete_model)
+    model = Curation(**complete_model)
     assert model.format_version == "2"
     assert len(model.unit_ids) == 5
     assert len(model.label_definitions) == 2
@@ -274,7 +274,7 @@ def test_complete_model():
         "removed": [5],
     }
 
-    model = CurationModel(**complete_model_dict)
+    model = Curation(**complete_model_dict)
     assert model.format_version == "2"
     assert len(model.unit_ids) == 5
     assert len(model.label_definitions) == 2