Dylan
Dylan

Reputation: 913

Get a column of probability values for a certain class

When using randomforest in sklearn to do a binary classification. I know I can use clf.predict(X) to get the predicted class. And when I use clf.predict_proba(X), I get an array look like this: enter image description here

I think the first column indicates the probability of the prediction? How can I get a column that is the probability of the class being 1?

Upvotes: 2

Views: 5240

Answers (1)

cs95
cs95

Reputation: 402553

From the RandomForest.predict_proba docs:

predict_proba(X)

Predict class probabilities for X.

Returns:
p : [...] The class probabilities of the input samples. The order of the classes corresponds to that in the attribute classes_.

You can look at the clf.classes_ attribute, see at what index your class 1 appears, and then just access the probabilities as so:

prob_class_1 = clf.predict_proba(X)[:, i] 

Where i is the index of class 1 in clf.classes_.

Upvotes: 8

Related Questions