def divideSet()

in causalml/inference/tree/uplift.pyx [0:0]

• 31 lines of code
• 7 McCabe index (conditional complexity)

    def divideSet(X, treatment_idx, y, column, value, missing_go_to_left=1):
        '''
        Tree node split.

        Args
        ----
        X : ndarray, shape = [num_samples, num_features]
            An ndarray of the covariates used to train the uplift model.
        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.
        column : int
                The column used to split the data.
        value : float or int
                The value in the column for splitting the data.

        Returns
        -------
        (X_l, X_r, treatment_l, treatment_r, y_l, y_r) : list of ndarray
                The covariates, treatments and outcomes of left node and the right node.
        '''
        # for int and float values
        if isinstance(value, numbers.Number):
            filt = X[:, column] >= value
        else:
            filt = X[:, column] == value

        # Handle NaNs only for numeric columns
        if np.issubdtype(X[:, column].dtype, np.number):
            nan_mask = np.isnan(X[:, column])

            if missing_go_to_left:
                filt = filt | nan_mask
            else:
                filt = filt & ~nan_mask

        return X[filt], X[~filt], treatment_idx[filt], treatment_idx[~filt], y[filt], y[~filt]