Treatment featurizer in CausalForestDML

[mikelgda]

I’m trying to build a model for the heterogenous treatment effect using the CausalForestDML class where the nuisance functions are fitted by a CatBoost regressor. Although this works fine for a linear continuous treatment, I get an error from CatBoost when trying polynomial features for the treatment using the treatment_featurizer argument.

The CatBoost nuisance model for the treatment is receiving a multi-column target, which doesn’t seem to support. Is this the intended behavior? Can a model be fitted to each polynomial degree of the feature or do I need a multi-output regressor?

[immu4989]

Hi @mikelgda , the multi column target you're seeing is expected: treatment_featurizer produces a featurized 2D treatment, so the first stage model_t needs to predict every column. The fix is to wrap your CatBoost regressor with sklearn.multioutput.MultiOutputRegressor:

from sklearn.multioutput import MultiOutputRegressor
from catboost import CatBoostRegressor

est = CausalForestDML(
    model_t=MultiOutputRegressor(CatBoostRegressor(...)),
    treatment_featurizer=PolynomialFeatures(...),
    ...
)