《MachineLearningInAction》之绘制决策树

《MachineLearningInAction》（Peter Harrington）中的代码有点小问题，我重写了全书所有代码，分享于此。

Block Ⅰ

import matplotlib.pyplot as plt #用于调用绘图
import matplotlib #用于调用rcParams属性，设置绘图窗口风格

Block Ⅱ

定义全局变量

decisionNode #决策节边框样式

leafNode #叶子节点边框样式

arrow_args #箭头样式

Block Ⅲ

定义函数

retrieveTree #为简化问题及做函数测试，手动生成大小不一的Tree

getNumLeafs #获取Tree的叶子数

getTreeDepth #获取Tree的深度，即decisionNode的个数

plotNode #绘制节点，通过nodeType参数区分decisionNode及leafNode

plotMidText #annotate每一个dict的key

plotTree #迭代绘制决策树

Block Ⅴ

测试代码

当__name__ == '__main__'时，即作为主模块调用时执行

效果图

关键思路：

1、迭代生成整棵树，代码测试时从只有一个decisionNode开始（通过调用retrieveTree（0）获得）。

2、通过参数传递plot axis来在同一个轴上绘图。Peter通过在实时调用时给plotTree函数增加axis属性达到同样效果，稍显复杂。

3、整张图绘制在(0,0)，(1,1)围成的矩形区域内绘制，第一个decisionNode中心位于(0.5,1)，通过decisionNode及LeafNode的数目控制纵向及横向间距，此二者皆为定值。这一点不知道是否与Peter的思路一致，因为他的代码太令我眼花缭乱，没看，我全部重写的。

4、plotTree先绘制decisionNode及其于parentNode之间的箭头，特别地当中心坐标与父节点坐标相等时，系统函数不绘制箭头。

treePlotter源码如下：

# -*- coding: utf-8 -*-
"""
treePlotter.py
~~~~~~~~~~

A module with functions to plot decision tree.

Created on Thu Mar 23 17:26:57 2017

Run on Python 3.6

@author: Luo Shaozhuo

refer to 'MachineLearninginAction'

"""
#==============================================================================
# import
#==============================================================================
import matplotlib.pyplot as plt
import matplotlib


#==============================================================================
# Global variables
#==============================================================================
decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-") 


#==============================================================================
# functions
#==============================================================================
def retrieveTree(i=0):
    """
    return a predefined tree
    ~~~~~~~~~~
    i: must be 0 or 1. 1 for a taller tree
    ~~~~~~~~~~
    dictTree
    """
    listOfTrees =[{'no surfacing': {0: 'no', 1: 'yes'}},
                  {'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},
                  {'no surfacing': {0: 'no', 1: {'flippers': {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}
                  ]
    return listOfTrees[i]


def getNumLeafs(dictTree):
    """
    return the number of leafs
    ~~~~~~~~~~
    dictTree: a dictonary dipicting a decidion tree
    ~~~~~~~~~~
    nNumLeaf: number of leafs
    """
    nNumLeaf = 0
    for key in dictTree.keys():
        if type(dictTree[key]) == dict:
            nNumLeaf += getNumLeafs(dictTree[key])
        else:   nNumLeaf +=1
    return nNumLeaf


def getTreeDepth(dictTree):
    """
    return the tree depth
    ~~~~~~~~~~
    dictTree: a dictonary dipicting a decidion tree
    ~~~~~~~~~~
    nMaxDepth: tree depth
    """
    nMaxDepth = 0
    keys = list(dictTree.keys())[0]
    dictTrunk = dictTree[keys]
    for key in dictTrunk.keys():
        if type(dictTrunk[key]) == dict:
            nCurDepth = 1 + getTreeDepth(dictTrunk[key])
        else:
            nCurDepth = 1
        if nCurDepth > nMaxDepth:
            nMaxDepth = nCurDepth
    return nMaxDepth


def plotNode(pltAxis,strNodeTxt, tplCntrPt, tplPrntPt, nodeType):
    """
    plot a decision node or a leaf node depend on nodeType.
    ~~~~~~~~~~
    pltAxis: plot axis
    strNodeTxt: text in node box
    tplCntrPt: center coordinates of box
    tplPrntPt: starting coordinates of arrow
    nodeType: leafNode or decisionNode
    ~~~~~~~~~~
    N/A
    """
    pltAxis.annotate(strNodeTxt, xy=tplPrntPt, xycoords='axes fraction',
    xytext=tplCntrPt, textcoords='axes fraction',
    va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)


def plotMidText(pltAxis, cntrPt, parentPt, txtString):
    """
    add feature value in the middle of arrow
    ~~~~~~~~~~
    cntrPt:
    parentPt:
    txtString:
    ~~~~~~~~~~
    N/A
    """
    xMid = (parentPt[0]+cntrPt[0])/2.0
    yMid = (parentPt[1]+cntrPt[1])/2.0
    pltAxis.text(xMid, yMid, txtString)


def plotTree(dictTree, pltAxis, fTrunkLen, fBrchLen, tplCntrPt, tplPrntPt, strNodeTxt):
    """
    plot tree recursivly
    ~~~~~~~~~~
    dictTree: decision tree
    pltAxis: axis used for plotting
    fTrunkLen: difference of y coordinates between two decision nodes
    fBrchLen: difference of y coordinates between two leafs
    tplCntrPt: coordinates of parent node  
    strNodeTxt: text in node box
    ~~~~~~~~~~
    N/A
    """
    #plot root node
    plotNode(pltAxis, strNodeTxt, tplCntrPt, tplPrntPt, decisionNode)
    #plot branch node
    tplPrntPt = tplCntrPt
    nNumKey = len(dictTree.keys())
    fMean = sum([x for x in range(nNumKey)])/nNumKey
    for i,key in enumerate(dictTree.keys()):
        tplCntrPt = (tplPrntPt[0]+(i-fMean)*fBrchLen, tplPrntPt[1]-fTrunkLen)
        plotMidText(pltAxis, tplCntrPt, tplPrntPt, key)
        if type(dictTree[key]) == dict:
            strNodeTxt = list(dictTree[key].keys())[0]
            plotTree(dictTree[key][strNodeTxt], pltAxis, fTrunkLen, fBrchLen, tplCntrPt, tplPrntPt,strNodeTxt)
        else:
            strNodeTxt = dictTree[key]
            plotNode(pltAxis,strNodeTxt, tplCntrPt, tplPrntPt, leafNode)


if __name__ == '__main__':
    dictTree = retrieveTree(2)
    matplotlib.rcParams['toolbar'] = 'none'
    pltAxis = plt.subplot(111, frameon=False,xticks=[], yticks=[])
    fBrchLen = 1/getNumLeafs(dictTree)
    fTrunkLen= 1/getTreeDepth(dictTree)
    tplCntrPt = (0.5,1)
    tplPrntPt = tplCntrPt
    strNodeTxt = list(dictTree.keys())[0]
    plotTree(dictTree[strNodeTxt], pltAxis, fTrunkLen, fBrchLen, tplCntrPt, tplPrntPt,strNodeTxt)