MIPLearn v0.3

miplearn/classifiers/singleclass.py

@@ -0,0 +1,51 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+from typing import Callable, Optional
+
+import numpy as np
+import sklearn.base
+from sklearn.base import BaseEstimator
+from sklearn.utils.multiclass import unique_labels
+
+
+class SingleClassFix(BaseEstimator):
+    """
+    Some sklearn classifiers, such as logistic regression, have issues with datasets
+    that contain a single class. This meta-classifier fixes the issue. If the
+    training data contains a single class, this meta-classifier always returns that
+    class as a prediction. Otherwise, it fits the provided base classifier,
+    and returns its predictions instead.
+    """
+
+    def __init__(
+        self,
+        base_clf: BaseEstimator,
+        clone_fn: Callable = sklearn.base.clone,
+    ):
+        self.base_clf = base_clf
+        self.clf_: Optional[BaseEstimator] = None
+        self.constant_ = None
+        self.classes_ = None
+        self.clone_fn = clone_fn
+
+    def fit(self, x: np.ndarray, y: np.ndarray) -> None:
+        classes = unique_labels(y)
+        if len(classes) == 1:
+            assert classes[0] is not None
+            self.clf_ = None
+            self.constant_ = classes[0]
+            self.classes_ = classes
+        else:
+            self.clf_ = self.clone_fn(self.base_clf)
+            assert self.clf_ is not None
+            self.clf_.fit(x, y)
+            self.constant_ = None
+            self.classes_ = self.clf_.classes_
+
+    def predict(self, x: np.ndarray) -> np.ndarray:
+        if self.constant_ is not None:
+            return np.full(x.shape[0], self.constant_)
+        else:
+            assert self.clf_ is not None
+            return self.clf_.predict(x)