BUG: binary classification failed in MultinomialNB, +regression test

larsmans · larsmans · commit 02e673b7ee30 · 2011-06-09T12:03:04.000+02:00
diff --git a/scikits/learn/naive_bayes.py b/scikits/learn/naive_bayes.py
@@ -315,6 +315,9 @@ def fit(self, X, y, class_prior=None):
             self.intercept_ = np.zeros(n_classes) - np.log(n_classes)
 
         Y = LabelBinarizer().fit_transform(y)
+        if Y.shape[1] == 1:
+            Y = np.concatenate((1 - Y, Y), axis=1)
+
         N_c, N_c_i = self._count(X, Y)
 
         self.coef_ = (np.log(N_c_i + self.alpha)
diff --git a/scikits/learn/tests/test_naive_bayes.py b/scikits/learn/tests/test_naive_bayes.py
@@ -102,11 +102,11 @@ def test_mnnb_pickle():
 def test_mnnb_predict_proba():
     '''Test multinomial NB's probability scores'''
 
-    clf = naive_bayes.MultinomialNB().fit([[0, 1], [0, 1], [1, 0]], [0, 0, 1])
-    assert clf.predict([0, 1]) == 0
+    clf = naive_bayes.MultinomialNB().fit([[0, 1], [1, 3], [4, 0]], [0, 0, 2])
+    assert clf.predict([4, 1]) == 2
     assert clf.predict_proba([0, 1]).shape == (1, 2)
-    assert np.sum(clf.predict_proba([0, 1])) == 1
-    assert np.sum(clf.predict_proba([1, 0])) == 1
+    assert np.sum(clf.predict_proba([1, 5])) == 1
+    assert np.sum(clf.predict_proba([3, 0])) == 1
     assert np.sum(np.exp(clf.class_log_prior_)) == 1
     assert np.sum(np.exp(clf.intercept_)) == 1
 
