Plotting legend for 2D numpy array

Question

I like to create a legend based on certain numbers using Numpy and Matplotlib but to no avail. So I started to play around with a test function to get it right before transferring this to my main script.

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

min_xyz = np.random.randint(5, size=(50,1,50))

fig = plt.figure(figsize=(7,7))
ax = fig.add_subplot(111)
ax.set_title('test')
plt.imshow(min_xyz[:,0,:])
ax.set_aspect('equal')

ax.set_xlabel('Distance')
ax.set_ylabel('Depth')
ax.legend()

So this creates something like this

I would like to create the legend that shows the following:

 Mineral 1 = colour_1
 Mineral 2 = colour_2
 Mineral 3 = colour_3
 Mineral 4 = colour_4
 Mineral 5 = colour_5

I tried working with ax.legend() but I can't seem to get it right. Any ideas?

Edit: Solution with vertical colour bars

My solution with the input from j08lue - Vertical colour bars

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.patches as mpatches

min_xyz = np.random.randint(5, size=(50,1,50))

fig, ax = plt.subplots(figsize = (7,40))
ax.set_title('test')
cax = ax.imshow(min_xyz[:,0,:], cmap = plt.cm.Accent)

"""Handles for min_xyz"""
cbar = plt.colorbar(cax, ticks=[0, 1, 2, 3, 4], 
                    orientation='vertical',
                    fraction=0.045, pad=0.05)
cbar.ax.set_yticklabels(['Mineral 1', 'Mineral 2', 'Mineral 3', 'Mineral 4','Mineral 5'])

ax.set_xlabel('Distance')
ax.set_ylabel('Depth')

Edit: Creating customised legend

I have placed the solution suggested by j08lue and managed to get the first legend correct. However, I believe it got to do with the normalising of the colour bar to get the legend reflect the right colour. I know I am missing something but I am not sure what I should be searching for. Any input is greatly appreciated.

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.patches as mpatches

min_xyz = np.random.randint(5, size=(50,50))

fig2, ax2 = plt.subplots(figsize = (7,40))
ax = fig.add_subplot(111)
ax2.set_title('test')
cax2 = ax2.imshow(min_xyz, cmap = plt.cm.Accent, vmin=0, vmax=4)
ax2.set_aspect('equal')

"""Handles for min_xyz"""
 my_colors = {
    'Mineral 1' : 0.,
    'Mineral 2' : 1., # It is normalised to 0 to 1
    'Mineral 3' : 2.,
    'Mineral 4' : 3.,
    'Mineral 5' : 4.,
}
patches = [mpatches.Patch(color=cmap(v), label=k) for k,v in sorted(my_colors.items(), key=lambda t: t[0])]
plt.legend(handles=patches, loc=2, bbox_to_anchor=(1.01,1))

ax2.set_xlabel('Distance')
ax2.set_ylabel('Depth')

Edited post with solution with color bars. Am working on a possible solution with proxy artists. — quenchlemon
– quenchlemon, Commented Feb 16, 2016 at 4:20
Great you are figuring things out. =) However, your values in my_colors are not between 0 and 1. Make them 0, 0.1, 0.2, 0.3, 0.4 instead and it will look much better. Or, even better, use 0, 0.25, 0.5, 0.75, 1 instead, to exploit the whole spectrum of the colormap. This is what I meant by normalizing your data to (0,1). — j08lue
– j08lue, Commented Feb 29, 2016 at 12:20

j08lue · Accepted Answer · 2016-02-09 08:48:48Z

1

Proxy artists

This can be done via proxy artists. Example from the docs:

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt

red_patch = mpatches.Patch(color='red', label='The red data')
plt.legend(handles=[red_patch])

plt.show()

But you need to figure out which colours correspond to which values. E.g.

cmap = plt.cm.viridis
my_colors = {
    'Mineral 1' : 0.1,
    'Mineral 2' : 0.2,
    }

patches = [mpatches.Patch(color=cmap(v), label=k) for k,v in my_colors.items()]

plt.legend(handles=patches)

The numbers in the dictionary correspond to the data normalized to [0,1] and you need to plot your data with the same cmap, of course.

Alternative: Colorbar

Alternatively, you can add a colorbar (the equivalent to the legend in imshow plots and the like) and place your labels on the ticks.

cbar = plt.colorbar(cax, ticks=list(my_colors.values()), orientation='horizontal')
cbar.ax.set_xticklabels(list(my_colors.keys()))

edited Feb 9, 2016 at 8:48

answered Feb 9, 2016 at 8:28

j08lue

1,7162 gold badges21 silver badges39 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

j08lue Over a year ago

NB: You might need to sort the dictionary when making the patches, if you want them to be in alphabetical order: sorted(my_colors.items(), key=lambda t: t[0]).

quenchlemon Over a year ago

Thanks for your input j08lue, I only managed to have a solution by using the colourbar. However, I don't wish to normalise the data (since the data is in integers (0, 1, 2, 3, 4).

j08lue Over a year ago

You do not need to normalize the data. Only the values for cmap(v) must be between 0 and 1.

Collectives™ on Stack Overflow

Plotting legend for 2D numpy array

Edit: Solution with vertical colour bars

Edit: Creating customised legend

1 Answer 1

Proxy artists

Alternative: Colorbar

3 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

Edit: Solution with vertical colour bars

Edit: Creating customised legend

1 Answer 1

Proxy artists

Alternative: Colorbar

3 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related