heatmap_colorized.py 文件源码

def heatmap(name, data, title, text):
    fig, ax = plt.subplots()
    ticks_font = font_manager.FontProperties(family='Decima Mono')
    plt.style.use([os.path.join(sys.path[0], 'ethplot.mplstyle')])
    LEFT = 0.125
    fig.suptitle(title,
             horizontalalignment='left',
             weight='bold', fontsize=20,
             x=LEFT, y=1)
    t = fig.text(LEFT, 0.92, text,
                 horizontalalignment='left',
                 weight='medium', fontsize=16, color='#555555')

    labels1 = ['PR','HD','SSSP','SCC']
    labels2 = ['PR','HD','SSSP','SCC']
    ax.set_xticklabels(labels1)
    ax.set_yticklabels(labels2)
    ax.set_yticks(np.arange(data.shape[0]) + 0.5)
    ax.set_xticks(np.arange(data.shape[1]) + 0.5)

    ax.tick_params(pad=11)

    plt.setp(ax.get_xticklabels(), fontproperties=ticks_font)
    plt.setp(ax.get_yticklabels(), fontproperties=ticks_font)

    norm = MidpointNormalize(midpoint=1.0)
    c = plt.pcolor(data, cmap = colors, vmin=0.5, vmax=2.5, norm=norm)

    values = data.as_matrix()
    for x in range(data.shape[0]):
        for y in range(data.shape[1]):
            #color = 'white' if values[y][x] > 2.3 else 'black'
            color = 'black'
            plt.text(x + 0.5, y + 0.5, '%.2f' % values[y][x],
                     horizontalalignment='center',
                     verticalalignment='center',
                     color=color,
                     fontproperties=ticks_font)

    colorbar = plt.colorbar(c)
    plt.setp(colorbar.ax.get_yticklabels(), fontproperties=ticks_font)

    plt.savefig(name + ".png", format='png')
    #ppad_inched=0.08 here because otherwise it cuts off the numbers...
    #plt.savefig(name + ".pdf", format='pdf', pad_inches=0.08)