sklearn.tree.plot_tree does not work with econml.grf.CausalForest

[4Freye]

Hello, I would like to plot a causal tree from a causal forest. I tried running the "1. Causal Forest: Heterogeneous causal effects with no unobserved confounders" section from this notebook but I receive an error when trying to plot a causal tree using plot_tree from sklearn.tree (see cell 8 in the notebook):

InvalidParameterError: The 'decision_tree' parameter of plot_tree must be an instance of 'sklearn.tree._classes.DecisionTreeClassifier' or an instance of 'sklearn.tree._classes.DecisionTreeRegressor'. Got GRFTree(criterion='het', max_features='auto', min_var_leaf_on_val=True,
random_state=1465622306) instead."

It seems that this worked in the past, perhaps this works with an older version of sklearn or econml? For reference, I am using sklearn version '1.4.2' and econml version '0.15.0'.

[kbattocchi]

Thanks for the report, and sorry for the inconvenience. Until we can address it on our end, I believe you should be able to work around it yourself by running

from sklearn.tree import DecisionTreeClassifier
from econml.grf._base_grftree import GRFTree
DecisionTreeClassifier.register(GRFTree)

before calling plot_tree. Could you see if this works for you?

[4Freye]

Yes that solved the issue, thank you!