matplotlib 用一个colorbar画两个heatmap

代码

import matplotlib.pyplot as plt
import numpy as np
def draw_bin_heatmap_together(data1, data2, xsName, ysName, outfig=None, colorscale=False):
    fig, (ax1, ax2, cax) = plt.subplots(ncols=3, figsize=(5.5, 5), gridspec_kw={"width_ratios":[1,1, 0.05]})
    cnorm = matplotlib.colors.LogNorm() if colorscale else matplotlib.colors.Normalize()
    fig.subplots_adjust(wspace=0.3)
    
    vmax = max(np.max(data1), np.max(data2))
    im1 = ax1.imshow(data1, norm=cnorm, vmin=1, vmax=vmax)
    im2 = ax2.imshow(data2, norm=cnorm, vmin=1, vmax=vmax)
    
    # Create colorbar
    fig.colorbar(im1, cax=cax)

    # We want to show all ticks...
    ax1.set_xticks(np.arange(len(xsName)))
    ax2.set_xticks(np.arange(len(xsName)))
    ax1.set_yticks(np.arange(len(ysName)))
    ax2.set_yticks(np.arange(len(ysName)))
    # ... and label them with the respective list entries
    ax1.set_xticklabels(xsName)
    ax2.set_xticklabels(xsName)
    ax1.set_yticklabels(ysName)
    ax2.set_yticklabels(ysName)

    # Rotate the tick labels and set their alignment.
    plt.setp(ax1.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
    plt.setp(ax2.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")

    ax1.set_title("heapmap 1")
    ax2.set_title("heatmap 2")
    fig.tight_layout()
    plt.show()
    
    if outfig is not None:
        fig.savefig(outfig)

效果

draw_bin_heatmap_together(count_data, countIn_data, [], [], colorscale=True)

参考链接：

ImportanceOfBeingErnest的回答
https://stackoverflow.com/questions/13784201/matplotlib-2-subplots-1-colorbar