def uplift_tree_string()

in causalml/inference/tree/plot.py [0:0]

• 32 lines of code
• 6 McCabe index (conditional complexity)

def uplift_tree_string(decisionTree, x_names):
    '''
    Convert the tree to string for print.

    Args
    ----

    decisionTree : object
        object of DecisionTree class

    x_names : list
        List of feature names

    Returns
    -------
    A string representation of the tree.
    '''

    # Column Heading
    dcHeadings = {}
    for i, szY in enumerate(x_names + ['treatment_group_key']):
        szCol = 'Column %d' % i
        dcHeadings[szCol] = str(szY)

    def toString(decisionTree, indent=''):
        if decisionTree.results is not None:  # leaf node
            return str(decisionTree.results)
        else:
            szCol = 'Column %s' % decisionTree.col
            if szCol in dcHeadings:
                szCol = dcHeadings[szCol]
            if isinstance(decisionTree.value, int) or isinstance(decisionTree.value, float):
                decision = '%s >= %s?' % (szCol, decisionTree.value)
            else:
                decision = '%s == %s?' % (szCol, decisionTree.value)
            trueBranch = indent + 'yes -> ' + toString(decisionTree.trueBranch, indent + '\t\t')
            falseBranch = indent + 'no  -> ' + toString(decisionTree.falseBranch, indent + '\t\t')
            return (decision + '\n' + trueBranch + '\n' + falseBranch)

    print(toString(decisionTree))