fix: explicit types and '==' for lc2st classifier comparison (#1550)

* fix: replace 'in' operator with '==' for proper classifier comparison

* Update classifier argument to accept classifier classes

* Update classifier initialization in lc2st constructor method

* Remove clf_class parameter

* Change classifier type from ClassifierMixin to BaseEstimator

Author: abelaba

sbi/diagnostics/lc2st.py

@@ -1,7 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 import numpy as np
 import torch
@@ -24,10 +24,9 @@ def __init__(
         seed: int = 1,
         num_folds: int = 1,
         num_ensemble: int = 1,
-        classifier: str = "mlp",
+        classifier: Union[str, Type[BaseEstimator]] = MLPClassifier,
         z_score: bool = False,
-        clf_class: Optional[Any] = None,
-        clf_kwargs: Optional[Dict[str, Any]] = None,
+        classifier_kwargs: Optional[Dict[str, Any]] = None,
         num_trials_null: int = 100,
         permutation: bool = True,
     ) -> None:
@@ -71,10 +70,11 @@ def __init__(
             num_ensemble: Number of classifiers for ensembling, defaults to 1.
                 This is useful to reduce variance coming from the classifier.
             z_score: Whether to z-score to normalize the data, defaults to False.
-            classifier: Classification architecture to use,
-                possible values: "random_forest" or "mlp", defaults to "mlp".
-            clf_class: Custom sklearn classifier class, defaults to None.
-            clf_kwargs: Custom kwargs for the sklearn classifier, defaults to None.
+            classifier: Classification architecture to use, can be one of the following:
+                    - "random_forest" or "mlp", defaults to "mlp" or
+                    - A classifier class (e.g., RandomForestClassifier, MLPClassifier)
+            classifier_kwargs: Custom kwargs for the sklearn classifier,
+                defaults to None.
             num_trials_null: Number of trials to estimate the null distribution,
                 defaults to 100.
             permutation: Whether to use the permutation method for the null hypothesis,
@@ -111,10 +111,26 @@ def __init__(
         self.num_ensemble = num_ensemble
 
         # initialize classifier
-        if "mlp" in classifier.lower():
-            ndim = thetas.shape[-1]
-            self.clf_class = MLPClassifier
-            if clf_kwargs is None:
+        if isinstance(classifier, str):
+            if classifier.lower() == 'mlp':
+                classifier = MLPClassifier
+            elif classifier.lower() == 'random_forest':
+                classifier = RandomForestClassifier
+            else:
+                raise ValueError(
+                    f'Invalid classifier: "{classifier}".'
+                    'Expected "mlp", "random_forest", '
+                    'or a valid scikit-learn classifier class.'
+                )
+        assert issubclass(classifier, BaseEstimator), (
+            "classifier must be a subclass of sklearn's BaseEstimator"
+        )
+        self.clf_class = classifier
+
+        self.clf_kwargs = classifier_kwargs
+        if self.clf_kwargs is None:
+            if self.clf_class == MLPClassifier:
+                ndim = thetas.shape[-1]
                 self.clf_kwargs = {
                     "activation": "relu",
                     "hidden_layer_sizes": (10 * ndim, 10 * ndim),
@@ -123,19 +139,8 @@ def __init__(
                     "early_stopping": True,
                     "n_iter_no_change": 50,
                 }
-        elif "random_forest" in classifier.lower():
-            self.clf_class = RandomForestClassifier
-            if clf_kwargs is None:
+            else:
                 self.clf_kwargs = {}
-        elif "custom":
-            if clf_class is None or clf_kwargs is None:
-                raise ValueError(
-                    "Please provide a valid sklearn classifier class and kwargs."
-                )
-            self.clf_class = clf_class
-            self.clf_kwargs = clf_kwargs
-        else:
-            raise NotImplementedError
 
         # initialize classifiers, will be set after training
         self.trained_clfs = None

tests/lc2st_test.py

@@ -16,7 +16,7 @@
 
 
 @pytest.mark.parametrize("method", (LC2ST, LC2ST_NF))
-@pytest.mark.parametrize("classifier", ('mlp', 'random_forest', 'custom'))
+@pytest.mark.parametrize("classifier", ('mlp', 'random_forest', MLPClassifier))
 @pytest.mark.parametrize("cv_folds", (1, 2))
 @pytest.mark.parametrize("num_ensemble", (1, 3))
 @pytest.mark.parametrize("z_score", (True, False))
@@ -72,9 +72,6 @@ def test_running_lc2st(method, classifier, cv_folds, num_ensemble, z_score):
             "num_eval": num_eval,
         }
         kwargs_eval = {}
-    if classifier == "custom":
-        kwargs_test["clf_class"] = MLPClassifier
-        kwargs_test["clf_kwargs"] = {"alpha": 0.0, "max_iter": 2500}
     kwargs_test["classifier"] = classifier
 
     lc2st = method(