plt.pcolormesh()绘制分类图

本文主要介绍如何利用plt.pcolormesh来绘制如下的分类图

plt.pcolormesh的作用在于能够直观表现出分类边界。如果只是单纯的绘制散点图，效果如下：

那么我们就看不出分类的边界。

下面将以鸢尾花数据集为例说明如何使用plt.pcolormesh，该数据集一共包含3类鸢尾花的数据

首先引入必要的库

   
   
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier
   
   

然后读取鸢尾花数据集，并对数据做一定的处理

   
   
iris_feature = u'花萼长度', u'花萼宽度', u'花瓣长度', u'花瓣宽度', u'类别'
path = 'iris.data'  # 数据文件路径
data = pd.read_csv(path, header= None)
data.columns=iris_feature
data[ '类别']=pd.Categorical(data[ '类别']).codes
   
   

处理完成后，一共有150组数据，数据长下面这样子

取花萼长度和花瓣长度做为特征，训练决策树模型

   
   
x_train = data[[ '花萼长度', '花瓣长度']]
y_train = data[ '类别']
model = DecisionTreeClassifier(criterion= 'entropy', min_samples_leaf= 3)
model.fit(x_train, y_train)
   
   

训练完模型后，现在需要画出分类边界，首先需要在横纵坐标各取500点，一共组成2500个点，然后把这2500个点送进决策树，来算出所属的种类，代码如下：

   
   
N, M = 500, 500 # 横纵各采样多少个值
x1_min, x2_min = x_train.min(axis= 0)
x1_max, x2_max = x_train.max(axis= 0)
t1 = np.linspace(x1_min, x1_max, N)
t2 = np.linspace(x2_min, x2_max, M)
x1, x2 = np.meshgrid(t1, t2) # 生成网格采样点
x_show = np.stack((x1.flat, x2.flat), axis= 1) # 测试点
y_predict=model.predict(x_show)
   
   

接着就可以绘制出分类图了。由于该数据集中一共有三种鸢尾花，所以绘制图片的时候需要三种颜色

   
   
cm_light = mpl.colors.ListedColormap([ '#A0FFA0', '#FFA0A0', '#A0A0FF'])
cm_dark = mpl.colors.ListedColormap([ 'g', 'r', 'b'])
   
   

接着使用plt.pcolormesh来绘制分类图

   
   
plt.pcolormesh(x1, x2, y_predict.reshape(x1.shape), cmap=cm_light)
plt.show()
   
   

plt.pcolormesh()会根据y_predict的结果自动在cmap里选择颜色

结果如下图

接着再把散点图也画上就大功告成了，结果如下：

完整代码如下

   
   
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier

iris_feature = u'花萼长度', u'花萼宽度', u'花瓣长度', u'花瓣宽度', u'类别'
path = 'iris.data' # 数据文件路径
data = pd.read_csv(path, header= None)
data.columns=iris_feature
data[ '类别']=pd.Categorical(data[ '类别']).codes
x_train = data[[ '花萼长度', '花瓣长度']]
y_train = data[ '类别']
model = DecisionTreeClassifier(criterion= 'entropy', min_samples_leaf= 3)
model.fit(x_train, y_train)

N, M = 500, 500 # 横纵各采样多少个值
x1_min, x2_min = x_train.min(axis= 0)
x1_max, x2_max = x_train.max(axis= 0)
t1 = np.linspace(x1_min, x1_max, N)
t2 = np.linspace(x2_min, x2_max, M)
x1, x2 = np.meshgrid(t1, t2) # 生成网格采样点
x_show = np.stack((x1.flat, x2.flat), axis= 1) # 测试点
y_predict=model.predict(x_show)


mpl.rcParams[ 'font.sans-serif'] = [ 'SimHei']
mpl.rcParams[ 'axes.unicode_minus'] = False
cm_light = mpl.colors.ListedColormap([ '#A0FFA0', '#FFA0A0', '#A0A0FF'])
cm_dark = mpl.colors.ListedColormap([ 'g', 'r', 'b'])
plt.xlim(x1_min, x1_max)
plt.ylim(x2_min, x2_max)
plt.pcolormesh(x1, x2, y_predict.reshape(x1.shape), cmap=cm_light)
plt.scatter(x_train[ '花萼长度'],x_train[ '花瓣长度'],c=y_train,cmap=cm_dark,marker= 'o',edgecolors= 'k')
plt.xlabel( '花萼长度')
plt.ylabel( '花瓣长度')
plt.title( '鸢尾花分类')
plt.grid( True,ls= ':')
plt.show()