timmens / causal-forest

Implements the Causal Forest algorithm formulated in Athey and Wager (2018).
MIT License
65 stars 12 forks source link

Add prediction function. #7

Closed timmens closed 4 years ago

timmens commented 4 years ago

For a fitted (causal) tree --stored as a pandas DataFrame-- implement a prediction function which finds the leaf in which new observations fall and returns the treatment effect in the given leaf.

timmens commented 4 years ago

In commit (61ed772f39e) I add a prediction method for multiple new observations (predict_causaltree), which is based on the internal function _predict_row_causaltree that predicts the estimated treatment effect for a single observation.

In commit (7aa3f9a91b92f) I add tests for both functions.