diff --git a/doc/api.rst b/doc/api.rst index ce850d1291..ecc209a9cc 100755 --- a/doc/api.rst +++ b/doc/api.rst @@ -411,6 +411,8 @@ spikeinterface.curation .. autofunction:: bombcell_label_units .. autofunction:: bombcell_get_default_thresholds .. autofunction:: model_based_label_units + .. autofunction:: get_required_metrics_from_model + .. autofunction:: check_required_metrics_are_present .. autofunction:: load_model .. autofunction:: train_model .. autofunction:: unitrefine_label_units diff --git a/src/spikeinterface/curation/__init__.py b/src/spikeinterface/curation/__init__.py index 16bdddd870..9d3405ef62 100644 --- a/src/spikeinterface/curation/__init__.py +++ b/src/spikeinterface/curation/__init__.py @@ -22,7 +22,12 @@ # automated curation from .curation_tools import get_labeling_summary from .threshold_metrics_curation import threshold_metrics_label_units -from .model_based_curation import model_based_label_units, load_model, auto_label_units +from .model_based_curation import ( + model_based_label_units, + load_model, + get_required_metrics_from_model, + check_required_metrics_are_present, +) from .train_manual_curation import train_model, get_default_classifier_search_spaces from .unitrefine_curation import unitrefine_label_units from .bombcell_curation import ( diff --git a/src/spikeinterface/curation/model_based_curation.py b/src/spikeinterface/curation/model_based_curation.py index 992f3ecd12..bc1192e7ce 100644 --- a/src/spikeinterface/curation/model_based_curation.py +++ b/src/spikeinterface/curation/model_based_curation.py @@ -12,6 +12,13 @@ _format_metric_dataframe, ) +# Map old metric column names to new metric column names for backwards compatibility +_BACKWARD_COMPATIBILITY_MAP = { + "peak_to_valley": {"name": "peak_to_trough_duration"}, + "peak_trough_ratio": {"name": "peak_after_to_trough_ratio", "flip_sign": True}, + "half_width": {"name": "trough_half_width"}, +} + class ModelBasedClassification: """ @@ -30,6 +37,8 @@ class ModelBasedClassification: The sorting analyzer object containing the spike sorting data. pipeline : Pipeline The pipeline object representing the trained classification model. + unit_ids : list | None + The list of unit IDs to consider for classification. If None, all units in the sorting analyzer are used. Methods ------- @@ -37,18 +46,31 @@ class ModelBasedClassification: Predicts the labels for the spike sorting data using the trained model. """ - def __init__(self, sorting_analyzer: SortingAnalyzer, pipeline): + def __init__(self, sorting_analyzer: SortingAnalyzer, pipeline=None, unit_ids: list[str | int] | None = None): from sklearn.pipeline import Pipeline if not isinstance(pipeline, Pipeline): raise ValueError("The `pipeline` must be an instance of sklearn.pipeline.Pipeline") + if sorting_analyzer is None: + raise ValueError("`sorting_analyzer` must be provided.") self.sorting_analyzer = sorting_analyzer self.pipeline = pipeline self.required_metrics = pipeline.feature_names_in_ + if unit_ids is None: + unit_ids = sorting_analyzer.unit_ids + self.unit_ids = unit_ids def predict_labels( - self, label_conversion=None, input_data=None, export_to_phy=False, model_info=None, enforce_metric_params=False + self, + label_conversion: dict[int, str] | None = None, + metrics=None, + export_to_phy: bool = False, + phy_folder: Path | None = None, + model_info: dict | None = None, + enforce_metric_params: bool = False, + set_predictions_as_properties: bool = True, + input_data=None, ): """ Predicts the labels for the spike sorting data using the trained model. @@ -61,13 +83,22 @@ def predict_labels( label_conversion : dict or None, default: None A dictionary for converting the predicted labels (which are integers) to custom labels. If None, tries to find in `model_info` file. The dictionary should have the format {old_label: new_label}. - input_data : pandas.DataFrame or None, default: None - The input data for classification. If not provided, the method will extract metrics stored in the sorting analyzer. export_to_phy : bool, default: False. Whether to export the classified units to Phy format. Default is False. + phy_folder : Path or None, default: None + The path to the Phy folder where the classified units will be exported. If None, + the Phy folder will be inferred from the sorting object. If the sorting object does not have a Phy folder, + the esport will be skipped. + model_info : dict or None, default: None + Dictionary of model info containing provenance of the model. enforce_metric_params : bool, default: False If True and the parameters used to compute the metrics in `sorting_analyzer` are different than the parmeters used to compute the metrics used to train the model, this function will raise an error. Otherwise, a warning is raised. + set_predictions_as_properties : bool, default: True + Whether to set the predictions as properties in the sorting object. + If True, the predicted labels and probabilities will be stored in the 'classifier_label' and 'classifier_probability' properties of the sorting object. + input_data : deprecated, default: None + Deprecated parameter. Use `metrics` instead. Returns ------- @@ -77,17 +108,28 @@ def predict_labels( """ import pandas as pd + if input_data is not None: + warnings.warn( + "`input_data` is deprecated. Use the `metrics` argument instead.", FutureWarning, stacklevel=2 + ) + if metrics is None: + metrics = input_data + # Get metrics DataFrame for classification - if input_data is None: - input_data = self.sorting_analyzer.get_metrics_extension_data() + if metrics is None: + metrics = self.sorting_analyzer.get_metrics_extension_data() else: - if not isinstance(input_data, pd.DataFrame): + if not isinstance(metrics, pd.DataFrame): raise ValueError("Input data must be a pandas DataFrame") - input_data = self.handle_backwards_compatibility_in_metrics(input_data, model_info=model_info) - input_data = self._check_required_metrics_are_present(input_data) + # Restrict metrics to the unit_ids of the sorting_analyzer if available + if len(metrics) > len(self.unit_ids): + metrics = metrics.loc[self.unit_ids] + + metrics = _handle_backwards_compatibility_in_metrics(metrics, model_info=model_info) + metrics = check_required_metrics_are_present(self.required_metrics, metrics) - if model_info is not None: + if model_info is not None and self.sorting_analyzer is not None: self._check_params_for_classification(enforce_metric_params, model_info=model_info) if model_info is not None and label_conversion is None: @@ -100,83 +142,34 @@ def predict_labels( except: warnings.warn("Could not find `label_conversion` key in `model_info.json` file") - input_data = _format_metric_dataframe(input_data) + metrics = _format_metric_dataframe(metrics) # Apply classifier - predictions = self.pipeline.predict(input_data) - probabilities = self.pipeline.predict_proba(input_data) + predictions = self.pipeline.predict(metrics) + probabilities = self.pipeline.predict_proba(metrics) probabilities = np.max(probabilities, axis=1) if isinstance(label_conversion, dict): - if set(predictions).issubset(set(label_conversion.keys())) is False: raise ValueError("Labels in predictions do not match those in label_conversion") predictions = [label_conversion[label] for label in predictions] classified_units = pd.DataFrame( - zip(predictions, probabilities), columns=["prediction", "probability"], index=self.sorting_analyzer.unit_ids + zip(predictions, probabilities), columns=["prediction", "probability"], index=self.unit_ids ) # Set predictions and probability as sorting properties - self.sorting_analyzer.set_sorting_property("classifier_label", predictions) - self.sorting_analyzer.set_sorting_property("classifier_probability", probabilities) + if set_predictions_as_properties: + self.sorting_analyzer.set_sorting_property("classifier_label", predictions) + self.sorting_analyzer.set_sorting_property("classifier_probability", probabilities) if export_to_phy: - self._export_to_phy(classified_units) + if phy_folder is None: + raise ValueError("Phy folder must be provided using the `phy_folder` parameter.") + classified_units.to_csv(f"{phy_folder}/cluster_prediction.tsv", sep="\t", index_label="cluster_id") return classified_units - def handle_backwards_compatibility_in_metrics(self, calculated_metrics, model_info): - """ - Handles backwards compatibility in metric names for models trained with older versions of SpikeInterface. - In recent versions, some metric names have been changed for clarity. In addition, the sign of some metrics - has been inverted to maintain consistency. - - Parameters - ---------- - calculated_metrics : pd.DataFrame - The DataFrame containing the calculated metrics. - model_info : dict or None - Dictionary of model info containing provenance of the model. - - Returns - ------- - pd.DataFrame - The DataFrame with updated metric names for compatibility. - """ - if model_info is None: - return calculated_metrics - si_version = model_info["requirements"].get("spikeinterface", None) - if si_version is not None and parse(si_version) < parse("0.103.2"): - # if the model was trained with SI version < 0.103.2, we need to rename some metrics - calculated_metrics = calculated_metrics.copy() - # peak_to_trough_duration was named peak_to_valley - if "peak_to_trough_duration" in calculated_metrics.columns: - calculated_metrics = calculated_metrics.rename(columns={"peak_to_trough_duration": "peak_to_valley"}) - # peak_after_to_trough_ratio was named peak_trough_ratio and had inverted sign - if "peak_after_to_trough_ratio" in calculated_metrics.columns: - calculated_metrics = calculated_metrics.rename( - columns={"peak_after_to_trough_ratio": "peak_trough_ratio"} - ) - calculated_metrics["peak_trough_ratio"] = -1 * calculated_metrics["peak_trough_ratio"] - # trough_half_width was named half_width - if "trough_half_width" in calculated_metrics.columns: - calculated_metrics = calculated_metrics.rename(columns={"trough_half_width": "half_width"}) - return calculated_metrics - - def _check_required_metrics_are_present(self, calculated_metrics): - # Check all the required metrics have been calculated - required_metrics = set(self.required_metrics) - if required_metrics.issubset(set(calculated_metrics)): - input_data = calculated_metrics[self.required_metrics] - else: - raise ValueError( - "Input data does not contain all required metrics for classification", - f"Missing metrics: {required_metrics.difference(calculated_metrics)}", - ) - - return input_data - def _check_params_for_classification(self, enforce_metric_params=False, model_info=None): """ Check that quality and template metrics parameters match those used to train the model @@ -226,22 +219,10 @@ def _check_params_for_classification(self, enforce_metric_params=False, model_in else: warnings.warn(warning_message) - def _export_to_phy(self, classified_df): - """Export the classified units to Phy as cluster_prediction.tsv file""" - - # Export to Phy format - try: - sorting_path = self.sorting_analyzer.sorting.get_annotation("phy_folder") - assert sorting_path is not None - assert Path(sorting_path).is_dir() - except AssertionError: - raise ValueError("Phy folder not found in sorting annotations, or is not a directory") - - classified_df.to_csv(f"{sorting_path}/cluster_prediction.tsv", sep="\t", index_label="cluster_id") - def model_based_label_units( sorting_analyzer: SortingAnalyzer, + unit_ids=None, model_folder=None, repo_id=None, model_name=None, @@ -250,6 +231,7 @@ def model_based_label_units( trusted=None, export_to_phy=False, enforce_metric_params=False, + set_predictions_as_properties=True, ): """ Automatically labels units based on a model-based classification, either from a model @@ -263,13 +245,15 @@ def model_based_label_units( ---------- sorting_analyzer : SortingAnalyzer The sorting analyzer object containing the spike sorting results. + unit_ids : list[str | int] | None, default: None + A list of unit IDs to consider. If None, all units in the sorting_analyzer will be used. model_folder : str or Path, default: None The path to the folder containing the model repo_id : str, default: None Hugging face repo id which contains the model e.g. 'username/model' model_name: str, default: None Filename of model e.g. 'my_model.skops'. If None, uses first model found. - label_conversion : dic | None, default: None + label_conversion : dict | None, default: None A dictionary for converting the predicted labels (which are integers) to custom labels. If None, tries to extract from `model_info.json` file. The dictionary should have the format {old_label: new_label}. export_to_phy : bool, default: False @@ -282,6 +266,9 @@ def model_based_label_units( enforce_metric_params : bool, default: False If True and the parameters used to compute the metrics in `sorting_analyzer` are different than the parmeters used to compute the metrics used to train the model, this function will raise an error. Otherwise, a warning is raised. + set_predictions_as_properties : bool, default: True + Whether to set the predictions as properties in the sorting object. + If True, the predicted labels and probabilities will be stored in the 'classifier_label' and 'classifier_probability' properties of the sorting object. Returns @@ -305,29 +292,59 @@ def model_based_label_units( if not isinstance(model, Pipeline): raise ValueError("The model must be an instance of sklearn.pipeline.Pipeline") - model_based_classification = ModelBasedClassification(sorting_analyzer, model) + model_based_classification = ModelBasedClassification( + sorting_analyzer=sorting_analyzer, unit_ids=unit_ids, pipeline=model + ) classified_units = model_based_classification.predict_labels( label_conversion=label_conversion, export_to_phy=export_to_phy, model_info=model_info, enforce_metric_params=enforce_metric_params, + set_predictions_as_properties=set_predictions_as_properties, ) return classified_units -def auto_label_units(*args, **kwargs): +def get_required_metrics_from_model( + model_folder=None, repo_id=None, model_name=None, model=None, trust_model=False, trusted=None +): """ - Deprecated function. Please use `model_based_label_units` instead. + Returns the required metrics for a model, either from a model hosted on HuggingFaceHub or one available in a local folder. + + Parameters + ---------- + model_folder : str or Path, default: None + The path to the folder containing the model + repo_id : str, default: None + Hugging face repo id which contains the model e.g. 'username/model' + model_name: str, default: None + Filename of model e.g. 'my_model.skops'. If None, uses first model found. + model : sklearn.pipeline.Pipeline, default: None + A trained sklearn pipeline model. If provided, the required metrics will be extracted from this model + trust_model : bool, default: False + Whether to trust the model. If True, the `trusted` parameter that is passed to `skops.load` to load the model will be + automatically inferred. If False, the `trusted` parameter must be provided to indicate the trusted objects. + trusted : list of str, default: None + Passed to skops.load. The object will be loaded only if there are only trusted objects and objects of types listed in trusted in the dumped file. + + Returns + ------- + required_metrics : list of str + A list of required metrics for the model. """ - warnings.warn( - "`auto_label_units` is deprecated and will be removed in v0.105.0. " - "Please use `model_based_label_units` instead.", - FutureWarning, - stacklevel=2, - ) - return model_based_label_units(*args, **kwargs) + from sklearn.pipeline import Pipeline + + if model is None: + model, _ = load_model( + model_folder=model_folder, repo_id=repo_id, model_name=model_name, trust_model=trust_model, trusted=trusted + ) + + if not isinstance(model, Pipeline): + raise ValueError("The model must be an instance of sklearn.pipeline.Pipeline") + + return _handle_backwards_compatibility_in_metric_names(list(model.feature_names_in_)) def load_model(model_folder=None, repo_id=None, model_name=None, trust_model=False, trusted=None): @@ -400,6 +417,26 @@ def _load_model_from_huggingface(repo_id=None, model_name=None, trust_model=Fals return model, model_info +def _patch_sklearn_imputer_compatibility(model): + """Fix SimpleImputer attribute rename from _fill_dtype (sklearn<1.5) to _fit_dtype (sklearn>=1.5).""" + from sklearn.impute import SimpleImputer + + steps_to_check = [] + if hasattr(model, "steps"): + steps_to_check = [step for _, step in model.steps] + elif hasattr(model, "estimators_"): + steps_to_check = list(model.estimators_) + + for step in steps_to_check: + if hasattr(step, "steps"): + _patch_sklearn_imputer_compatibility(step) + elif isinstance(step, SimpleImputer): + if hasattr(step, "_fill_dtype") and not hasattr(step, "_fit_dtype"): + step._fit_dtype = step._fill_dtype + elif hasattr(step, "_fit_dtype") and not hasattr(step, "_fill_dtype"): + step._fill_dtype = step._fit_dtype + + def _load_model_from_folder(model_folder=None, model_name=None, trust_model=False, trusted=None): """ Loads a model and model_info from a folder @@ -444,6 +481,7 @@ def _load_model_from_folder(model_folder=None, model_name=None, trust_model=Fals trusted = [list_item for list_item in string_list.split("'") if len(list_item) > 2] model = skio.load(skops_file, trusted=trusted) + _patch_sklearn_imputer_compatibility(model) model_info_path = folder / "model_info.json" if not model_info_path.is_file(): @@ -452,13 +490,16 @@ def _load_model_from_folder(model_folder=None, model_name=None, trust_model=Fals else: model_info = json.load(open(model_info_path)) - model_info = handle_backwards_compatibility_metric_params(model_info) + model_info = _handle_backwards_compatibility_metric_params(model_info) return model, model_info -def handle_backwards_compatibility_metric_params(model_info): - +def _handle_backwards_compatibility_metric_params(model_info): + """ + Handles backwards compatibility in metric parameters for models trained with older versions of SpikeInterface. + In recent versions, some metric parameters have been changed for clarity. + """ if ( model_info.get("metric_params") is not None and model_info.get("metric_params").get("quality_metric_params") is not None @@ -479,3 +520,74 @@ def handle_backwards_compatibility_metric_params(model_info): del model_info["metric_params"]["template_metric_params"]["metrics_kwargs"] return model_info + + +def _handle_backwards_compatibility_in_metrics(calculated_metrics, model_info): + """ + Handles backwards compatibility in metric names for models trained with older versions of SpikeInterface. + In recent versions, some metric names have been changed for clarity. In addition, the sign of some metrics + has been inverted to maintain consistency. + + Parameters + ---------- + calculated_metrics : pd.DataFrame + The DataFrame containing the calculated metrics. + model_info : dict or None + Dictionary of model info containing provenance of the model. + + Returns + ------- + pd.DataFrame + The DataFrame with updated metric names for compatibility. + """ + if model_info is None: + return calculated_metrics + si_version = model_info["requirements"].get("spikeinterface", None) + if si_version is not None and parse(si_version) < parse("0.103.2"): + # If the model was trained with SI version < 0.103.2, we need to rename some metrics + calculated_metrics = calculated_metrics.copy() + # We need to rename the metrics and flip sign when needed + for old_name, updated_dict in _BACKWARD_COMPATIBILITY_MAP.items(): + updated_name = updated_dict["name"] + if updated_name in calculated_metrics.columns: + calculated_metrics = calculated_metrics.rename(columns={updated_name: old_name}) + if updated_dict.get("flip_sign", False): + calculated_metrics[old_name] = -1 * calculated_metrics[old_name] + return calculated_metrics + + +def _handle_backwards_compatibility_in_metric_names(model_metric_names): + """ + Handles backwards compatibility in metric names for models trained with older versions of SpikeInterface. + In recent versions, some metric names have been changed for clarity. + + Parameters + ---------- + model_metric_names : list of str + The list of metric names used in the model. + + Returns + ------- + list of str + The list of updated metric names for compatibility. + """ + updated_metric_names = [] + for metric_name in model_metric_names: + if metric_name in _BACKWARD_COMPATIBILITY_MAP: + updated_metric_names.append(_BACKWARD_COMPATIBILITY_MAP[metric_name]["name"]) + else: + updated_metric_names.append(metric_name) + return updated_metric_names + + +def check_required_metrics_are_present(required_metrics, calculated_metrics): + # Check all the required metrics have been calculated, preserving the order expected by the pipeline + if set(required_metrics).issubset(set(calculated_metrics.columns)): + input_data = calculated_metrics[list(required_metrics)] + else: + raise ValueError( + "Input data does not contain all required metrics for classification", + f"Missing metrics: {set(required_metrics).difference(calculated_metrics.columns)}", + ) + + return input_data diff --git a/src/spikeinterface/curation/tests/test_model_based_curation.py b/src/spikeinterface/curation/tests/test_model_based_curation.py index 94b98418bb..a42ba2250d 100644 --- a/src/spikeinterface/curation/tests/test_model_based_curation.py +++ b/src/spikeinterface/curation/tests/test_model_based_curation.py @@ -3,7 +3,12 @@ from spikeinterface.curation.tests.common import sorting_analyzer_for_unitrefine_curation, trained_pipeline_path from spikeinterface.curation.model_based_curation import ModelBasedClassification -from spikeinterface.curation import model_based_label_units, load_model +from spikeinterface.curation import ( + model_based_label_units, + load_model, + get_required_metrics_from_model, + check_required_metrics_are_present, +) import numpy as np @@ -25,15 +30,25 @@ def model(trained_pipeline_path): @pytest.fixture -def required_metrics_and_columns(): +def required_metrics(): """These are the metrics which `model` are trained on.""" - return ["num_spikes", "snr", "half_width"], ["num_spikes", "snr", "trough_half_width", "peak_half_width"] + from spikeinterface.metrics import ComputeQualityMetrics, ComputeTemplateMetrics + + all_metric_names = ["snr", "half_width", "peak_to_trough_duration", "number_of_peaks"] + quality_metric_names = ["snr"] + template_metric_names = ["half_width", "peak_to_trough_duration", "number_of_peaks"] + all_metric_columns = ComputeQualityMetrics.get_metric_columns( + quality_metric_names + ) + ComputeTemplateMetrics.get_metric_columns(template_metric_names) + return all_metric_names, all_metric_columns, quality_metric_names, template_metric_names def test_model_based_classification_init(sorting_analyzer_for_unitrefine_curation, model): """Test that the ModelBasedClassification attributes are correctly initialised""" - model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0]) + model_based_classification = ModelBasedClassification( + sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0] + ) assert model_based_classification.sorting_analyzer == sorting_analyzer_for_unitrefine_curation assert model_based_classification.pipeline == model[0] assert np.all(model_based_classification.required_metrics == model_based_classification.pipeline.feature_names_in_) @@ -68,52 +83,34 @@ def test_metric_ordering_independence(sorting_analyzer_for_unitrefine_curation, def test_model_based_classification_get_metrics_for_classification( - sorting_analyzer_for_unitrefine_curation, model, required_metrics_and_columns + sorting_analyzer_for_unitrefine_curation, model, required_metrics ): """If the user has not computed the required metrics, an error should be returned. This test checks that an error occurs when the required metrics have not been computed, and that no error is returned when the required metrics have been computed. """ - sorting_analyzer_for_unitrefine_curation.delete_extension("quality_metrics") sorting_analyzer_for_unitrefine_curation.delete_extension("template_metrics") - required_metric_names, required_metric_columns = required_metrics_and_columns + all_metric_names, all_metric_columns, qm_names, tm_names = required_metrics - model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0]) + model_based_classification = ModelBasedClassification( + sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0] + ) # Compute some (but not all) of the required metrics in sorting_analyzer, should still error - sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=[required_metric_names[0]]) + sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=[all_metric_names[0]]) computed_metrics = sorting_analyzer_for_unitrefine_curation.get_metrics_extension_data() with pytest.raises(ValueError): - model_based_classification._check_required_metrics_are_present(computed_metrics) + check_required_metrics_are_present(all_metric_columns, computed_metrics) # Compute all of the required metrics in sorting_analyzer, no more error - sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=required_metric_names[0:2]) - sorting_analyzer_for_unitrefine_curation.compute("template_metrics", metric_names=[required_metric_names[2]]) + sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=qm_names) + sorting_analyzer_for_unitrefine_curation.compute("template_metrics", metric_names=tm_names) metrics_data = sorting_analyzer_for_unitrefine_curation.get_metrics_extension_data() - assert metrics_data.shape[0] == len(sorting_analyzer_for_unitrefine_curation.sorting.get_unit_ids()) - assert set(metrics_data.columns.to_list()) == set(required_metric_columns) - - -def test_model_based_classification_export_to_phy(sorting_analyzer_for_unitrefine_curation, model): - import pandas as pd - - # Test the _export_to_phy() method of ModelBasedClassification - model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0]) - - classified_units = pd.DataFrame.from_dict({0: (1, 0.5), 1: (0, 0.5), 2: (1, 0.5), 3: (0, 0.5), 4: (1, 0.5)}) - # Function should fail here - with pytest.raises(ValueError): - model_based_classification._export_to_phy(classified_units) - # Make temp output folder and set as phy_folder - phy_folder = cache_folder / "phy_folder" - phy_folder.mkdir(parents=True, exist_ok=True) - - model_based_classification.sorting_analyzer.sorting.annotate(phy_folder=phy_folder) - model_based_classification._export_to_phy(classified_units) - assert (phy_folder / "cluster_prediction.tsv").exists() + assert len(metrics_data) == len(sorting_analyzer_for_unitrefine_curation.unit_ids) + assert set(metrics_data.columns.to_list()) == set(all_metric_columns) def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefine_curation, model): @@ -128,7 +125,9 @@ def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefi sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=["num_spikes", "snr"]) # Test the predict_labels() method of ModelBasedClassification - model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0]) + model_based_classification = ModelBasedClassification( + sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0] + ) classified_units = model_based_classification.predict_labels() predictions = classified_units["prediction"].values @@ -142,6 +141,45 @@ def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefi assert np.all(predictions_labelled == expected_result_converted) +def test_predict_labels_with_phy_export(sorting_analyzer_for_unitrefine_curation, model): + """Test that the predict_labels() method of ModelBasedClassification correctly exports to Phy format when requested.""" + + sorting_analyzer_for_unitrefine_curation.compute( + "template_metrics", metric_names=["half_width", "peak_to_trough_duration", "number_of_peaks"] + ) + sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=["num_spikes", "snr"]) + + phy_folder = cache_folder / "phy_export" + phy_folder.mkdir(parents=True, exist_ok=True) + + model_based_classification = ModelBasedClassification( + sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0] + ) + classified_units = model_based_classification.predict_labels(export_to_phy=True, phy_folder=phy_folder) + + # Check that the cluster_prediction.tsv file was created in the specified phy_folder + assert (phy_folder / "cluster_prediction.tsv").exists() + + # Using export_to_phy=True without providing a phy_folder should raise a ValueError + with pytest.raises(ValueError): + model_based_classification.predict_labels(export_to_phy=True, phy_folder=None) + + +def test_get_required_metrics_from_model(model, required_metrics): + """Test that the get_required_metrics_from_model function returns the correct required metrics and columns.""" + + required_from_model = get_required_metrics_from_model(model=model[0]) + + _, all_metric_columns, _, _ = required_metrics + assert set(all_metric_columns) == set(required_from_model) + + # from HF + required_metrics_from_model_hf = get_required_metrics_from_model( + repo_id="SpikeInterface/UnitRefine_sua_mua_classifier", trust_model=True + ) + assert set(all_metric_columns) != set(required_metrics_from_model_hf[0]) + + @pytest.mark.skip(reason="We need to retrain the model to reflect any changes in metric computation") def test_exception_raised_when_metric_params_not_equal(sorting_analyzer_for_unitrefine_curation, trained_pipeline_path): """We track whether the metric parameters used to compute the metrics used to train diff --git a/src/spikeinterface/curation/unitrefine_curation.py b/src/spikeinterface/curation/unitrefine_curation.py index 309e924bb8..c619b7372a 100644 --- a/src/spikeinterface/curation/unitrefine_curation.py +++ b/src/spikeinterface/curation/unitrefine_curation.py @@ -6,7 +6,7 @@ def unitrefine_label_units( - sorting_analyzer: SortingAnalyzer, + sorting_analyzer: SortingAnalyzer | None = None, noise_neural_classifier: str | Path | None = None, sua_mua_classifier: str | Path | None = None, ): @@ -18,7 +18,7 @@ def unitrefine_label_units( Parameters ---------- - sorting_analyzer : SortingAnalyzer + sorting_analyzer : SortingAnalyzer or None, default: None The sorting analyzer object containing the spike sorting results. noise_neural_classifier : str or Path or None, default: None The path to the folder containing the model, a full path to a model (".skops") @@ -57,6 +57,7 @@ def unitrefine_label_units( noise_neuron_labels = model_based_label_units( sorting_analyzer=sorting_analyzer, trust_model=True, + set_predictions_as_properties=False, **get_model_based_classification_kwargs(noise_neural_classifier), ) if set(noise_neuron_labels["prediction"]) != {"noise", "neural"}: @@ -64,20 +65,21 @@ def unitrefine_label_units( "The noise/neural classifier did not return the expected labels 'noise' and 'neural'. " "Please check the model used for classification." ) - noise_units = noise_neuron_labels[noise_neuron_labels["prediction"] == "noise"] - sorting_analyzer_neural = sorting_analyzer.remove_units(noise_units.index) + unit_ids_neural = noise_neuron_labels[noise_neuron_labels["prediction"] != "noise"].index else: - sorting_analyzer_neural = sorting_analyzer - noise_units = pd.DataFrame(columns=["prediction", "probability"]) + noise_neuron_labels = pd.DataFrame(index=sorting_analyzer.unit_ids, columns=["prediction", "probability"]) + unit_ids_neural = sorting_analyzer.unit_ids if sua_mua_classifier is not None: # 2. apply the sua/mua classification and aggregate results - if len(sorting_analyzer.unit_ids) > len(noise_units): + if len(unit_ids_neural) > 0: with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=InconsistentVersionWarning) sua_mua_labels = model_based_label_units( - sorting_analyzer=sorting_analyzer_neural, + sorting_analyzer=sorting_analyzer, + unit_ids=unit_ids_neural, trust_model=True, + set_predictions_as_properties=False, **get_model_based_classification_kwargs(sua_mua_classifier), ) if set(sua_mua_labels["prediction"]) != {"sua", "mua"}: @@ -85,9 +87,10 @@ def unitrefine_label_units( "The sua/mua classifier did not return the expected labels 'sua' and 'mua'. " "Please check the model used for classification." ) - all_labels = pd.concat([sua_mua_labels, noise_units]).sort_index() + noise_labels = noise_neuron_labels[noise_neuron_labels["prediction"] == "noise"] + all_labels = pd.concat([sua_mua_labels, noise_labels]).reindex(sorting_analyzer.unit_ids) else: - all_labels = noise_units + all_labels = noise_neuron_labels else: all_labels = noise_neuron_labels