Matplotlib的subplot画图, 共享colorbar

Python画图,利用Matplotlib中subplot画3*3的heatmap图,所有热力图共享一个colorbar。

import numpy as np
import matplotlib
matplotlib.use('AGG')
import matplotlib.pyplot as plt
import csv
import pandas as pd  

fig, axs = plt.subplots(3, 3, sharex=True, sharey=True, constrained_layout=True, figsize=(6, 6))
fig.subplots_adjust(hspace=0.1)
datasets = ['envi', 'enfr','enge']
types = ['enc', 'enc-dec' ,'dec']
Datasets = ['En-Vi', 'En-Fr','En-Ge']

def draw_map(ax,xlabels, ylabels, tt, attns, index):   #画热力图
    im = ax.imshow(attns,interpolation='none',  cmap='Blues', vmin=0, vmax=1)
    ax.set_title(tt, fontsize=10)
style='italic' )
    return im
xlabels = '0 1 2 3 4 5'.split()
ylabels = '5 4 3 2 1 0'.split()

for i, data in enumerate(datasets):
    datafile = '{}.csv'.format(data) #读取数据
    data = pd.read_csv(datafile, delimiter='	', header=None)
    for j, tt in enumerate(types):
        A = data.values[:,j*6:(j+1)*6] #取出所需热力图的数据
        ax = axs[i, j]
        im = draw_map(ax, xlabels,ylabels, '%s (%s)'%(Datasets[i],tt), A, j)

                
plt.xticks(range(len(xlabels)),xlabels)    #设置横坐标label
plt.yticks(range(len(ylabels)),ylabels)    #设置纵坐标label

cb_ax = fig.add_axes([1.0, 0.1, 0.02, 0.8]) #设置colarbar位置
cbar = fig.colorbar(im, cax=cb_ax)     #共享colorbar
plt.tight_layout()
plt.show()
fig.savefig('weights_visualization.pdf', bbox_inches='tight', dpi=500)

画图结果:

原文地址:https://www.cnblogs.com/huadongw/p/14315944.html