本文主要介绍如何利用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()