How to get the outcome for Treatment=0 and Treatment=1 separately?

[AaronSh90]

Hi
I have seen some examples on how to use TLearner for Individual Treatment Effect (ITE) and the outputs represent those values. However, I am going to determine the Y(0)|X and Y(1)|X separately. I am wondering how I can have these two values?

[0trm]

predict() takes return_components=True. It then returns the effect plus two dicts keyed by treatment group, holding the control-model and treatment-model predictions:

from causalml.inference.meta import BaseTRegressor
from sklearn.ensemble import GradientBoostingRegressor

learner = BaseTRegressor(learner=GradientBoostingRegressor(), control_name=0)
learner.fit(X=X, treatment=treatment, y=y)

te, yhat_c, yhat_t = learner.predict(X, return_components=True)
g = learner.t_groups[0]   # the treatment label, e.g. 1
y0 = yhat_c[g]            # E[Y | X, T=0]
y1 = yhat_t[g]            # E[Y | X, T=g]
# te[:, 0] equals y1 - y0

With BaseTClassifier the components come from predict_proba, so they are P(Y=1 | X, T=0) and P(Y=1 | X, T=1). With several treatments there is one key per group, and yhat_c holds the same control prediction under every key. return_components can't be combined with return_ci=True. fit_predict() and the S-, X- and DR-learners take the same argument (tlearner.py). Checked on causalml 0.17.0.