Add a classes parameter to LabelEncoder.transform

Hamzeh Alsalhi · Hamzeh Alsalhi · commit 9153fc072436 · 2014-09-02T13:02:55.000-04:00
diff --git a/sklearn/preprocessing/label.py b/sklearn/preprocessing/label.py
@@ -145,28 +145,35 @@ def fit_transform(self, y):
         self.classes_, y = np.unique(y, return_inverse=True)
         return y
 
-    def transform(self, y):
+    def transform(self, y, classes=None):
         """Transform labels to normalized encoding.
 
         Parameters
         ----------
         y : array-like of shape [n_samples]
             Target values.
 
+        classes : array-like, optional (default: None)
+            List of unique sorted labels to encode the target data against.
+            If None the LabelEncoder must have already been fit and the unique
+            labels from the fit will be used.
+
         Returns
         -------
         y : array-like of shape [n_samples]
         """
-        self._check_fitted()
+        if classes is None:
+            self._check_fitted()
+            classes = self.classes_
 
-        classes = np.unique(y)
-        _check_numpy_unicode_bug(classes)
-        if len(np.intersect1d(classes, self.classes_)) < len(classes):
+        y_classes = np.unique(y)
+        _check_numpy_unicode_bug(y_classes)
+        if len(np.intersect1d(y_classes, classes)) < len(y_classes):
             # Get the new labels
-            unseen = np.setdiff1d(classes, self.classes_)
+            unseen = np.setdiff1d(y_classes, classes)
 
             if type(self.new_labels) is int:
-                ret = np.searchsorted(self.classes_, y)
+                ret = np.searchsorted(classes, y)
                 ret[np.in1d(y, unseen)] = self.new_labels
                 return ret
             elif self.new_labels is None:
@@ -176,7 +183,7 @@ def transform(self, y):
                 raise ValueError("Value of argument `new_labels`={0} is "
                                  "unknown.".format(self.new_labels))
 
-        return np.searchsorted(self.classes_, y)
+        return np.searchsorted(classes, y)
 
     def inverse_transform(self, y):
         """Transform labels back to original encoding.
diff --git a/sklearn/preprocessing/tests/test_label.py b/sklearn/preprocessing/tests/test_label.py
@@ -220,6 +220,18 @@ def test_label_encoder_new_label_replace():
     assert_array_equal(le.transform(["b", "c", "d"]), [1, 2, -99])
 
 
+def test_label_encoder_transform_classes_parameter():
+    """Test LabelEncoder's transform using the classes parameter"""
+    le = LabelEncoder(new_labels=None)
+    le.fit(["a", "b", "b", "c"])
+    assert_array_equal(le.classes_, ["a", "b", "c"])
+    assert_array_equal(le.transform(["d", "f", "e", "e"],
+                                    classes=["d", "e", "f"]),
+                       [0, 2, 1, 1])
+    assert_array_equal(le.inverse_transform([2, 1, 0]), ["c", "b", "a"])
+    assert_raises(ValueError, le.transform, ["b", "c", "d"])
+
+
 def test_label_encoder_fit_transform():
     """Test fit_transform"""
     le = LabelEncoder()
