def tree_node_summary()

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