in causalml/inference/tree/uplift.pyx [0:0]
def tree_node_summary(self, treatment_idx, y, min_samples_treatment=10, n_reg=100, parentNodeSummary=None):
'''
Tree node summary statistics.
Args
----
treatment_idx : array-like, shape = [num_samples]
An array containing the treatment group index 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 : list of list
The positive probabilities and sample sizes of each of the control and treatment groups
in the parent node.
Returns
-------
nodeSummary : list of list
The positive probabilities and sample sizes of each of the control and treatment groups
in the current node.
'''
# counts: [[N(Y=0, T=0), N(Y=1, T=0)], [N(Y=0, T=1), N(Y=1, T=1)], ...]
counts = self.group_uniqueCounts(treatment_idx, y)
# nodeSummary: [[P(Y=1|T=0), N(T=0)], [P(Y=1|T=1), N(T=1)], ...]
nodeSummary = []
# Iterate the control and treatment groups
for i, count in enumerate(counts):
n_pos = count[1]
n = count[0] + n_pos
if parentNodeSummary is None:
p = n_pos / n if n > 0 else 0.
elif n > min_samples_treatment:
p = (n_pos + parentNodeSummary[i][0] * n_reg) / (n + n_reg)
else:
p = parentNodeSummary[i][0]
nodeSummary.append([p, n])
return nodeSummary