Decision trees python

decision_tree.py

import pandas as pd

import math as math

def loop(df, target):
    if df.columns.size > 1:
        gains = info_gains(df, target)
        res = max(gains, key=gains.get)
        df.groupby(res).apply(lambda v: loop(v, target))

def log2(x):
    return math.log(x, 2)

def info_gains(df, target):
    res = {}
    for k in df: res[k] = gain_of_attr(df, k, target)
    res.pop(target, None)
    return res


def gain_of_attr(df, attr, target):
    grouped = df.groupby(attr)
    res = 0
    for k,v in grouped:
        p_t = grouped.indices[k].size/float(df[target].size)
        e_k = entropy_of(v[target])
        res += p_t * e_k

    return entropy_of(df[target]) - res

def entropy_of(col):
    props = proportion_of_attr(col)
    ixs = props.index
    xs = ixs.map(lambda ix: props[ix]*log2(props[ix]))
    return -sum(xs)

def proportion_of_attr(col):
    counts = col.size
    return col.value_counts().apply(lambda x: x/float(counts))