Matplotlib Heatmap with Discrete Custom Colorbar

36 Views Asked by At

I've been working on a plot where the data is composed of integers between 1 and 7 (inclusive). The plot is a heatmap and ideally I could associate whichever colour I want to any number. Accompanying the plot would be a colorbar showing the distinct colours along with their labels.

The issue is that I am unsure how the colour sequence works. I would expect data comprised of 7, to be associated with grey, since it is the last color in my colours list, but it gets associated to white (third element in list).

Below is the code used to produce the plot (I've removed a few things that are unrelated to the question).

def plot(data):
    fig, ax = plt.subplots(figsize=(16, 4), dpi = 300)
    
    colours = ["cyan", "red", "white", "lime", "yellow", "k", 'grey']
    bounds = np.array([1, 2, 3, 4, 5, 6, 7, 8]) -0.5 #0.5 to center it which does not work
    cmap = colors.ListedColormap(colours)
    norm = colors.BoundaryNorm(bounds, cmap.N)
    h = ax.pcolor(data, norm= norm, cmap = cmap)
   
    cbar = plt.colorbar(h)
    cbar.set_ticks(ticks = bounds, labels = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'])
    
    plt.title("Assessment of Different Turbulent Mixing Sources", fontsize = 20)
    plt.xlabel("Time EDT (hh:mm)", fontsize = 20)
    plt.ylabel("Height (m)", fontsize=20)
    plt.show()

Which produced the following plot: Plot

1

There are 1 best solutions below

0
Ratislaus On

The data-colors mapping seems to work perfectly, I tried it out with the values proposed by @JohanC, i.e. [[1, 2, 3, 4], [5, 6, 7, 1]] and I saw "7" showing up as a grey rectangle. In order to center ticks on the color bar segments you could add 0.5 to the ticks argument and use all but the last value of bounds (and of their labels too):

cbar.set_ticks(ticks = bounds[:-1] + 0.5, labels = ['a', 'b', 'c', 'd', 'e', 'f', 'g'])

The output:

enter image description here