def test_probabilities(self):
inner_est = DummyProbEstimator(3)
est = Probabilities(inner_est)
# test attr wrap
self.assertEqual(est._coefs, inner_est._coefs)
self.assertRaises(AttributeError, getattr, est, "nope_not_attr")
preds = est.predict(self.data.values)
self.assertEqual(preds.shape, (10, 3))
def test_sklearn_probabilities(self):
# test multi-class
self.data["target"] = [0] * 5 + [1] * 3 + [2] * 2
inner_est = linear_model.LogisticRegression()
est = Probabilities(inner_est)
x = self.data[["a", "b"]]
est.fit(x, self.data.target)
preds = est.predict(x)
self.assertEqual(preds.shape, (10, 3))
# test binary, single output
self.data["target"] = [0] * 5 + [1] * 5
est = BinaryProbabilities(inner_est)
x = self.data[["a", "b"]]
est.fit(x, self.data.target)
preds = est.predict(x)
self.assertEqual(preds.shape, (10,))