matplotlib: colorbars and its text labels

前端 未结 2 1282
夕颜
夕颜 2020-11-28 20:44

I\'d like to create a colorbar legend for a heatmap, such that the labels are in the center of each discrete color. Example borrowed from here:

2条回答
  •  盖世英雄少女心
    2020-11-28 21:04

    import matplotlib.pyplot as plt
    import numpy as np
    from matplotlib.colors import ListedColormap
    
    #discrete color scheme
    cMap = ListedColormap(['white', 'green', 'blue','red'])
    
    #data
    np.random.seed(42)
    data = np.random.rand(4, 4)
    fig, ax = plt.subplots()
    heatmap = ax.pcolor(data, cmap=cMap)
    
    #legend
    cbar = plt.colorbar(heatmap)
    
    cbar.ax.get_yaxis().set_ticks([])
    for j, lab in enumerate(['$0$','$1$','$2$','$>3$']):
        cbar.ax.text(.5, (2 * j + 1) / 8.0, lab, ha='center', va='center')
    cbar.ax.get_yaxis().labelpad = 15
    cbar.ax.set_ylabel('# of contacts', rotation=270)
    
    
    # put the major ticks at the middle of each cell
    ax.set_xticks(np.arange(data.shape[1]) + 0.5, minor=False)
    ax.set_yticks(np.arange(data.shape[0]) + 0.5, minor=False)
    ax.invert_yaxis()
    
    #labels
    column_labels = list('ABCD')
    row_labels = list('WXYZ')
    ax.set_xticklabels(column_labels, minor=False)
    ax.set_yticklabels(row_labels, minor=False)
    
    plt.show()
    

    You were very close. Once you have a reference to the color bar axis, you can do what ever you want to it, including putting text labels in the middle. You might want to play with the formatting to make it more visible.

    demo

提交回复
热议问题