in causalml/inference/tree/models.py [0:0]
def tree_node_summary(self, treatment, y, min_samples_treatment=10, n_reg=100, parentNodeSummary=None):
'''
Tree node summary statistics.
Args
----
treatment : array-like, shape = [num_samples]
An array containing the treatment group for each unit.
y : array-like, shape = [num_samples]
An array containing the outcome of interest for each unit.
min_samples_treatment: int, optional (default=10)
The minimum number of samples required of the experiment group t be split at a leaf node.
n_reg : int, optional (default=10)
The regularization parameter defined in Rzepakowski et al. 2012,
the weight (in terms of sample size) of the parent node influence
on the child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
parentNodeSummary : dictionary
Node summary statistics of the parent tree node.
Returns
-------
nodeSummary : dictionary
The node summary of the current tree node.
'''
# returns {treatment_group: p(1)}
results = self.group_uniqueCounts(treatment, y)
# node Summary: {treatment_group: [p(1), size]}
nodeSummary = {}
# iterate treatment group
for r in results:
n1 = results[r][1]
ntot = results[r][0] + n1
if parentNodeSummary is None:
y_mean = n1 / ntot
elif ntot > min_samples_treatment:
y_mean = (n1 + parentNodeSummary[r][0] * n_reg) / (ntot + n_reg)
else:
y_mean = parentNodeSummary[r][0]
nodeSummary[r] = [y_mean, ntot]
return nodeSummary