Fix bug in MultinomialNB: output transposed

larsmans · larsmans · commit f228fdb3eab7 · 2011-06-08T16:36:40.000+02:00
Quick fix: .T before return. Needs a proper fix.
diff --git a/scikits/learn/naive_bayes.py b/scikits/learn/naive_bayes.py
@@ -408,4 +408,4 @@ def predict_log_proba(self, X):
         jll = self._joint_log_likelihood(X)
         # normalize by P(x) = P(f_1, ..., f_n)
         log_prob_x = np.logaddexp.reduce(jll[:, np.newaxis])
-        return jll - log_prob_x
+        return (jll - log_prob_x).T
diff --git a/scikits/learn/tests/test_naive_bayes.py b/scikits/learn/tests/test_naive_bayes.py
@@ -104,6 +104,7 @@ def test_mnnb_predict_proba():
 
     clf = naive_bayes.MultinomialNB().fit([[0,1], [0,1], [1,0]], [0,0,1])
     assert clf.predict([0,1]) == 0
+    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(np.exp(clf.class_log_prior_)) == 1
