in causalml/inference/tree/models.py [0:0]
def predict(self, X, full_output=False):
'''
Returns the recommended treatment group and predicted optimal
probability conditional on using the recommended treatment group.
Args
----
X : ndarray, shape = [num_samples, num_features]
An ndarray of the covariates used to train the uplift model.
full_output : bool, optional (default=False)
Whether the UpliftTree algorithm returns upliftScores, pred_nodes
alongside the recommended treatment group and p_hat in the treatment group.
Returns
-------
df_res : DataFrame, shape = [num_samples, (num_treatments + 1)]
A DataFrame containing the predicted delta in each treatment group,
the best treatment group and the maximum delta.
'''
df_res = pd.DataFrame()
y_pred_ensemble = dict()
y_pred_list = np.zeros((X.shape[0], len(self.classes_)))
# Make prediction by each tree
for tree_i in range(len(self.uplift_forest)):
_, _, _, y_pred_full = self.uplift_forest[tree_i].predict(X=X, full_output=True)
if tree_i == 0:
for treatment_group in y_pred_full:
y_pred_ensemble[treatment_group] = (
np.array(y_pred_full[treatment_group]) / len(self.uplift_forest)
)
else:
for treatment_group in y_pred_full:
y_pred_ensemble[treatment_group] = (
np.array(y_pred_ensemble[treatment_group])
+ np.array(y_pred_full[treatment_group]) / len(self.uplift_forest)
)
# Summarize results into dataframe
for treatment_group in y_pred_ensemble:
df_res[treatment_group] = y_pred_ensemble[treatment_group]
df_res['recommended_treatment'] = df_res.apply(np.argmax, axis=1)
# Calculate delta
delta_cols = []
for treatment_group in y_pred_ensemble:
if treatment_group != self.control_name:
delta_cols.append('delta_%s' % (treatment_group))
df_res['delta_%s' % (treatment_group)] = df_res[treatment_group] - df_res[self.control_name]
# Add deltas to results list
y_pred_list[:, self.classes_[treatment_group]] = df_res['delta_%s' % (treatment_group)].values
df_res['max_delta'] = df_res[delta_cols].max(axis=1)
return y_pred_list