Merge pull request #2674 from GaelVaroquaux/bug_ovo_string_y

GaelVaroquaux · GaelVaroquaux · commit 6ec2c8b2ae38 · 2013-12-18T22:55:43.000-08:00
BUG: OneVsOneClassifier was broken with string labels
diff --git a/doc/whats_new.rst b/doc/whats_new.rst
@@ -87,6 +87,9 @@ Changelog
    - Fixed bug in :class:`linear_model.stochastic_gradient` :
      ``l1_ratio`` was used as ``(1.0 - l1_ratio)`` .
 
+   - Fixed bug in :class:`multiclass.OneVsOneClassifier` with string
+     labels
+
 API changes summary
 -------------------
 
diff --git a/sklearn/multiclass.py b/sklearn/multiclass.py
@@ -304,10 +304,11 @@ def _fit_ovo_binary(estimator, X, y, i, j):
     """Fit a single binary estimator (one-vs-one)."""
     cond = np.logical_or(y == i, y == j)
     y = y[cond]
-    y[y == i] = 0
-    y[y == j] = 1
+    y_binary = np.empty(y.shape, np.int)
+    y_binary[y == i] = 0
+    y_binary[y == j] = 1
     ind = np.arange(X.shape[0])
-    return _fit_binary(estimator, X[ind[cond]], y, classes=[i, j])
+    return _fit_binary(estimator, X[ind[cond]], y_binary, classes=[i, j])
 
 
 def fit_ovo(estimator, X, y, n_jobs=1):
diff --git a/sklearn/tests/test_multiclass.py b/sklearn/tests/test_multiclass.py
@@ -312,6 +312,17 @@ def test_ovo_ties2():
         assert_equal(ovo_prediction[0], (1 + i) % 3)
 
 
+def test_ovo_string_y():
+    "Test that the OvO doesn't screw the encoding of string labels"
+    X = np.eye(4)
+    y = np.array(['a', 'b', 'c', 'd'])
+
+    svc = LinearSVC()
+    ovo = OneVsOneClassifier(svc)
+    ovo.fit(X, y)
+    assert_array_equal(y, ovo.predict(X))
+
+
 def test_ecoc_exceptions():
     ecoc = OutputCodeClassifier(LinearSVC(random_state=0))
     assert_raises(ValueError, ecoc.predict, [])
