Improve robustness of model classifier tests (#4447)

Author: chrishalcrow

src/spikeinterface/curation/tests/common.py

@@ -1,6 +1,6 @@
 import pytest
 
-from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer
+from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer, aggregate_units
 from spikeinterface.core.generate import inject_some_split_units
 from spikeinterface.curation import train_model
 from pathlib import Path
@@ -71,6 +71,19 @@ def sorting_analyzer_for_curation():
     return make_sorting_analyzer(sparse=True)
 
 
+@pytest.fixture(scope="module")
+def sorting_analyzer_for_unitrefine_curation():
+    """Makes an analyzer whose first 10 units are good normal units, and 10 which are noise. We make them
+    noise by using a spike trains which are uncorrelated with the recording for `sorting2`."""
+
+    recording, sorting_1 = generate_ground_truth_recording(num_channels=4, seed=1, num_units=10)
+    _, sorting_2 = generate_ground_truth_recording(num_channels=4, seed=2, num_units=10)
+    both_sortings = aggregate_units([sorting_1, sorting_2])
+    analyzer = create_sorting_analyzer(sorting=both_sortings, recording=recording)
+    analyzer.compute(["random_spikes", "noise_levels", "templates"])
+    return analyzer
+
+
 @pytest.fixture(scope="module")
 def sorting_analyzer_multi_segment_for_curation():
     return make_sorting_analyzer(sparse=True, durations=[50.0, 30.0])
@@ -83,7 +96,7 @@ def sorting_analyzer_with_splits():
 
 
 @pytest.fixture(scope="module")
-def trained_pipeline_path():
+def trained_pipeline_path(sorting_analyzer_for_unitrefine_curation):
     """
     Makes a model saved at "./trained_pipeline" which will be used by other tests in the module.
     If the model already exists, this function does nothing.
@@ -92,20 +105,22 @@ def trained_pipeline_path():
     if trained_model_folder.is_dir():
         yield trained_model_folder
     else:
-        analyzer = make_sorting_analyzer(sparse=True)
+        analyzer = sorting_analyzer_for_unitrefine_curation
         analyzer.compute(
             {
-                "quality_metrics": {"metric_names": ["snr", "num_spikes"]},
-                "template_metrics": {"metric_names": ["half_width"]},
+                "quality_metrics": {"metric_names": ["snr"]},
+                "template_metrics": {"metric_names": ["half_width", "peak_to_trough_duration", "number_of_peaks"]},
             }
         )
         train_model(
-            analyzers=[analyzer] * 5,
-            labels=[[1, 0, 1, 0, 1]] * 5,
+            analyzers=[analyzer],
             folder=trained_model_folder,
-            classifiers=["RandomForestClassifier"],
+            labels=[[1] * 10 + [0] * 10],
             imputation_strategies=["median"],
             scaling_techniques=["standard_scaler"],
+            classifiers=["RandomForestClassifier"],
+            overwrite=True,
+            search_kwargs={"cv": 3, "scoring": "balanced_accuracy", "n_iter": 2},
         )
         yield trained_model_folder
 

src/spikeinterface/curation/tests/test_model_based_curation.py

@@ -1,7 +1,7 @@
 import pytest
 from pathlib import Path
 
-from spikeinterface.curation.tests.common import sorting_analyzer_for_curation, trained_pipeline_path
+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
 
@@ -16,9 +16,9 @@
 
 @pytest.fixture
 def model(trained_pipeline_path):
-    """A toy model, created using the `sorting_analyzer_for_curation` from `spikeinterface.curation.tests.common`.
-    It has been trained locally and, when applied to `sorting_analyzer_for_curation` will label its 5 units with
-    the following labels: [1,0,1,0,1]."""
+    """A toy model, created using the `sorting_analyzer_for_unitrefine_curation` from `spikeinterface.curation.tests.common`.
+    It has been trained locally and, when applied to `sorting_analyzer_for_unitrefine_curation` will label its 10 units with
+    the following labels: [1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0]."""
 
     model = load_model(trained_pipeline_path, trusted=["numpy.dtype"])
     return model
@@ -30,32 +30,36 @@ def required_metrics_and_columns():
     return ["num_spikes", "snr", "half_width"], ["num_spikes", "snr", "trough_half_width", "peak_half_width"]
 
 
-def test_model_based_classification_init(sorting_analyzer_for_curation, model):
+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_curation, model[0])
-    assert model_based_classification.sorting_analyzer == sorting_analyzer_for_curation
+    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, 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_)
 
 
-def test_metric_ordering_independence(sorting_analyzer_for_curation, trained_pipeline_path):
+def test_metric_ordering_independence(sorting_analyzer_for_unitrefine_curation, trained_pipeline_path):
     """The function `model_based_label_units` needs the correct metrics to have been computed. However,
     it should be independent of the order of computation. We test this here."""
 
-    sorting_analyzer_for_curation.compute("template_metrics", metric_names=["half_width"])
-    sorting_analyzer_for_curation.compute("quality_metrics", metric_names=["num_spikes", "snr"])
+    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=["snr"])
 
     prediction_prob_dataframe_1 = model_based_label_units(
-        sorting_analyzer=sorting_analyzer_for_curation,
+        sorting_analyzer=sorting_analyzer_for_unitrefine_curation,
         model_folder=trained_pipeline_path,
         trusted=["numpy.dtype"],
     )
 
-    sorting_analyzer_for_curation.compute("quality_metrics", metric_names=["snr", "num_spikes"])
+    sorting_analyzer_for_unitrefine_curation.compute(
+        "template_metrics", metric_names=["peak_to_trough_duration", "half_width", "number_of_peaks"]
+    )
 
     prediction_prob_dataframe_2 = model_based_label_units(
-        sorting_analyzer=sorting_analyzer_for_curation,
+        sorting_analyzer=sorting_analyzer_for_unitrefine_curation,
         model_folder=trained_pipeline_path,
         trusted=["numpy.dtype"],
     )
@@ -64,40 +68,40 @@ def test_metric_ordering_independence(sorting_analyzer_for_curation, trained_pip
 
 
 def test_model_based_classification_get_metrics_for_classification(
-    sorting_analyzer_for_curation, model, required_metrics_and_columns
+    sorting_analyzer_for_unitrefine_curation, model, required_metrics_and_columns
 ):
     """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_curation.delete_extension("quality_metrics")
-    sorting_analyzer_for_curation.delete_extension("template_metrics")
+    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
 
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_curation, model[0])
+    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0])
 
     # Compute some (but not all) of the required metrics in sorting_analyzer, should still error
-    sorting_analyzer_for_curation.compute("quality_metrics", metric_names=[required_metric_names[0]])
-    computed_metrics = sorting_analyzer_for_curation.get_metrics_extension_data()
+    sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=[required_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)
 
     # Compute all of the required metrics in sorting_analyzer, no more error
-    sorting_analyzer_for_curation.compute("quality_metrics", metric_names=required_metric_names[0:2])
-    sorting_analyzer_for_curation.compute("template_metrics", metric_names=[required_metric_names[2]])
+    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]])
 
-    metrics_data = sorting_analyzer_for_curation.get_metrics_extension_data()
-    assert metrics_data.shape[0] == len(sorting_analyzer_for_curation.sorting.get_unit_ids())
+    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_curation, model):
+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_curation, model[0])
+    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
@@ -112,42 +116,48 @@ def test_model_based_classification_export_to_phy(sorting_analyzer_for_curation,
     assert (phy_folder / "cluster_prediction.tsv").exists()
 
 
-def test_model_based_classification_predict_labels(sorting_analyzer_for_curation, model):
+def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefine_curation, model):
     """The model `model` has been trained on the `sorting_analyzer` used in this test with
     the labels `[1, 0, 1, 0, 1]`. Hence if we apply the model to this `sorting_analyzer`
     we expect these labels to be outputted. The test checks this, and also checks
     that label conversion works as expected."""
 
-    sorting_analyzer_for_curation.compute("template_metrics", metric_names=["half_width"])
-    sorting_analyzer_for_curation.compute("quality_metrics", metric_names=["num_spikes", "snr"])
+    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"])
 
     # Test the predict_labels() method of ModelBasedClassification
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_curation, model[0])
+    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0])
     classified_units = model_based_classification.predict_labels()
     predictions = classified_units["prediction"].values
 
-    assert np.all(predictions == np.array([1, 0, 1, 0, 1]))
+    expected_result = np.array([1] * 10 + [0] * 10)
+    assert np.all(predictions == expected_result)
 
     conversion = {0: "noise", 1: "good"}
+    expected_result_converted = np.array(["good"] * 10 + ["noise"] * 10)
     classified_units_labelled = model_based_classification.predict_labels(label_conversion=conversion)
     predictions_labelled = classified_units_labelled["prediction"]
-    assert np.all(predictions_labelled == ["good", "noise", "good", "noise", "good"])
+    assert np.all(predictions_labelled == expected_result_converted)
 
 
 @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_curation, trained_pipeline_path):
+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
     a model are the same as the parameters used to compute the metrics in the sorting
     analyzer which is being curated. If they are different, an error or warning will
     be raised depending on the `enforce_metric_params` kwarg. This behaviour is tested here."""
 
-    sorting_analyzer_for_curation.compute(
-        "quality_metrics", metric_names=["num_spikes", "snr"], metric_params={"snr": {"peak_mode": "peak_to_peak"}}
+    sorting_analyzer_for_unitrefine_curation.compute(
+        "quality_metrics", metric_names=["snr"], metric_params={"snr": {"peak_mode": "peak_to_peak"}}
+    )
+    sorting_analyzer_for_unitrefine_curation.compute(
+        "template_metrics", metric_names=["half_width", "peak_to_trough_duration", "number_of_peaks"]
     )
-    sorting_analyzer_for_curation.compute("template_metrics", metric_names=["half_width"])
 
     model, model_info = load_model(model_folder=trained_pipeline_path, trusted=["numpy.dtype"])
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_curation, model)
+    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model)
 
     # an error should be raised if `enforce_metric_params` is True
     with pytest.raises(Exception):
@@ -158,13 +168,15 @@ def test_exception_raised_when_metric_params_not_equal(sorting_analyzer_for_cura
         model_based_classification._check_params_for_classification(enforce_metric_params=False, model_info=model_info)
 
     # Now test the positive case. Recompute using the default parameters
-    sorting_analyzer_for_curation.compute(
+    sorting_analyzer_for_unitrefine_curation.compute(
         "quality_metrics",
-        metric_names=["num_spikes", "snr"],
+        metric_names=["snr"],
         metric_params={"snr": {"peak_sign": "neg", "peak_mode": "extremum"}},
     )
-    sorting_analyzer_for_curation.compute("template_metrics", metric_names=["half_width"])
+    sorting_analyzer_for_unitrefine_curation.compute(
+        "template_metrics", metric_names=["half_width", "peak_to_trough_duration"]
+    )
 
     model, model_info = load_model(model_folder=trained_pipeline_path, trusted=["numpy.dtype"])
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_curation, model)
+    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model)
     model_based_classification._check_params_for_classification(enforce_metric_params=True, model_info=model_info)

src/spikeinterface/curation/tests/test_unitrefine_curation.py

@@ -1,84 +1,91 @@
 import pytest
 
-from spikeinterface.curation.tests.common import sorting_analyzer_for_curation, trained_pipeline_path
+from spikeinterface.curation.tests.common import sorting_analyzer_for_unitrefine_curation, trained_pipeline_path
 from spikeinterface.curation import unitrefine_label_units
 
 
-def test_unitrefine_label_units_hf(sorting_analyzer_for_curation):
+def test_unitrefine_label_units_hf(sorting_analyzer_for_unitrefine_curation):
     """Test the `unitrefine_label_units` function."""
-    sorting_analyzer_for_curation.compute("template_metrics", include_multi_channel_metrics=True)
-    sorting_analyzer_for_curation.compute("quality_metrics")
+    sorting_analyzer_for_unitrefine_curation.compute(
+        {
+            "spike_amplitudes": {},
+            "template_metrics": {"include_multi_channel_metrics": True},
+            "quality_metrics": {},
+        }
+    )
 
     # test passing both classifiers
     labels = unitrefine_label_units(
-        sorting_analyzer_for_curation,
+        sorting_analyzer_for_unitrefine_curation,
         noise_neural_classifier="SpikeInterface/UnitRefine_noise_neural_classifier_lightweight",
         sua_mua_classifier="SpikeInterface/UnitRefine_sua_mua_classifier_lightweight",
     )
 
     assert "unitrefine_label" in labels.columns
     assert "unitrefine_probability" in labels.columns
-    assert labels.shape[0] == len(sorting_analyzer_for_curation.sorting.unit_ids)
+    assert labels.shape[0] == len(sorting_analyzer_for_unitrefine_curation.sorting.unit_ids)
 
     # test only noise neural classifier
     labels = unitrefine_label_units(
-        sorting_analyzer_for_curation,
+        sorting_analyzer_for_unitrefine_curation,
         noise_neural_classifier="SpikeInterface/UnitRefine_noise_neural_classifier_lightweight",
         sua_mua_classifier=None,
     )
 
     assert "unitrefine_label" in labels.columns
     assert "unitrefine_probability" in labels.columns
-    assert labels.shape[0] == len(sorting_analyzer_for_curation.sorting.unit_ids)
+    assert labels.shape[0] == len(sorting_analyzer_for_unitrefine_curation.sorting.unit_ids)
 
     # test only sua mua classifier
     labels = unitrefine_label_units(
-        sorting_analyzer_for_curation,
+        sorting_analyzer_for_unitrefine_curation,
         noise_neural_classifier=None,
         sua_mua_classifier="SpikeInterface/UnitRefine_sua_mua_classifier_lightweight",
     )
 
     assert "unitrefine_label" in labels.columns
     assert "unitrefine_probability" in labels.columns
-    assert labels.shape[0] == len(sorting_analyzer_for_curation.sorting.unit_ids)
+    assert labels.shape[0] == len(sorting_analyzer_for_unitrefine_curation.sorting.unit_ids)
 
     # test passing none
     with pytest.raises(ValueError):
         labels = unitrefine_label_units(
-            sorting_analyzer_for_curation,
+            sorting_analyzer_for_unitrefine_curation,
             noise_neural_classifier=None,
             sua_mua_classifier=None,
         )
 
     # test warnings when unexpected labels are returned
     with pytest.warns(UserWarning):
         labels = unitrefine_label_units(
-            sorting_analyzer_for_curation,
+            sorting_analyzer_for_unitrefine_curation,
             noise_neural_classifier="SpikeInterface/UnitRefine_sua_mua_classifier_lightweight",
             sua_mua_classifier=None,
         )
 
     with pytest.warns(UserWarning):
         labels = unitrefine_label_units(
-            sorting_analyzer_for_curation,
+            sorting_analyzer_for_unitrefine_curation,
             noise_neural_classifier=None,
             sua_mua_classifier="SpikeInterface/UnitRefine_noise_neural_classifier_lightweight",
         )
 
 
-def test_unitrefine_label_units_with_local_models(sorting_analyzer_for_curation, trained_pipeline_path):
+def test_unitrefine_label_units_with_local_models(sorting_analyzer_for_unitrefine_curation, trained_pipeline_path):
     # test with trained local models
-    sorting_analyzer_for_curation.compute("template_metrics", include_multi_channel_metrics=True)
-    sorting_analyzer_for_curation.compute("quality_metrics")
+    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")
 
     # test passing model folder
     labels = unitrefine_label_units(
-        sorting_analyzer_for_curation,
+        sorting_analyzer_for_unitrefine_curation,
         noise_neural_classifier=trained_pipeline_path,
     )
 
     # test passing model folder
     labels = unitrefine_label_units(
-        sorting_analyzer_for_curation,
+        sorting_analyzer_for_unitrefine_curation,
         noise_neural_classifier=trained_pipeline_path / "best_model.skops",
     )