Adding a legend to a matplotlib plot of a numpy array with binary data

Question

I am writing a program to simulate radioactive decay. I am using a class with a decay_sim method which generates a 2D numpy array with ones (not decayed) and zeros (decayed) called self.nuclei. It is a pseudo-random process, so the array is different each time. I then use a visual method to plot this array as a colour grid using matplotlib.

This is an example of the self.nuclei array:

[[1 0 0 0 1 1 1 1 1 1]
 [1 0 1 0 1 0 1 1 0 0]
 [1 1 1 1 0 0 1 1 1 0]
 [1 1 0 0 1 1 0 0 0 0]
 [0 0 1 1 0 0 0 1 1 1]
 [1 0 0 1 1 1 0 0 1 1]
 [0 0 0 0 1 1 0 1 0 0]
 [1 0 0 0 0 0 1 0 0 1]
 [0 0 1 1 1 0 0 0 0 1]
 [0 1 0 1 1 0 1 0 1 1]]

This is the visual method:

def visual(self):                                                     
    """Creates a colour map of the nuclei, showing which have decayed.
                                                                      
    """                                                               
    plt.style.use('classic')                                             
                                                                      
    plt.title('Radioactive Nuclei')                                   
                                                                      
    plt.pcolormesh(self.nuclei, edgecolors='w', linewidth=1)          
    ax = plt.gca()                                                    
    ax.set_aspect('equal')                                            
    plt.colorbar()                                                    
                                                                      
    plt.show()

And this is its output:

I am looking for a way to add a legend to this plot that can label the decayed and not decayed colours. I am currently using plt.colorbar(), but this only gives me values - not a label. Ideally, it should look something like: "red = not decayed" and "blue = decayed" (although I would like the option to change the colour scheme later). I would also like to include the number of decayed/undecayed cells in the legend. I think I could get the number using:

import numpy as np

undecayed_no = np.count_nonzero(self.nuclei)
decayed_no = self.nuclei.size - undecayed_no

JohanC · Accepted Answer · 2020-08-17 18:58:15Z

3

You can create a custom legend from rectangular patches. The colors for the mesh can be set via LinearSegmentedColormap.from_list().

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.patches import Patch
import numpy as np

nuclei = np.array([[1,0,0,0,1,1,1,1,1,1],[1,0,1,0,1,0,1,1,0,0],[1,1,1,1,0,0,1,1,1,0],[1,1,0,0,1,1,0,0,0,0],[0,0,1,1,0,0,0,1,1,1],[1,0,0,1,1,1,0,0,1,1],[0,0,0,0,1,1,0,1,0,0],[1,0,0,0,0,0,1,0,0,1],[0,0,1,1,1,0,0,0,0,1],[0,1,0,1,1,0,1,0,1,1]])

plt.style.use('classic')

colors = ['dodgerblue', 'crimson']
ax = plt.gca()
ax.pcolormesh(nuclei, edgecolors='w', linewidth=1, cmap=LinearSegmentedColormap.from_list('', colors))
ax.set_aspect('equal')
legend_elements = [Patch(facecolor=color, edgecolor='w') for color in colors]
undecayed_no = np.count_nonzero(nuclei)
decayed_no = nuclei.size - undecayed_no
ax.legend(handles=legend_elements,
          labels=[f"decayed ({decayed_no})", f"non-decayed ({undecayed_no})"],
          loc="upper left", bbox_to_anchor=[1.02, 1])
ax.set_title('Radioactive Nuclei')
plt.tight_layout(pad=4)
plt.show()

edited Aug 17, 2020 at 18:58

answered Aug 17, 2020 at 18:47

JohanC

81.4k8 gold badges54 silver badges90 bronze badges

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

2 Comments

DJTS Over a year ago

This work perfectly, thanks. I just wanted to ask about the loc option in ax.legend. It says "upper left" but the legend appears to be in the upper right. When I change to upper right, the legend is placed inside the grid. I'm a bit confused about this.

JohanC Over a year ago

loc goes together with bbox_to_anchor(x, y) where x is measured from 0 at the left of the axes area to 1 at the right. So 1.02 is just outside. Similarly, y goes from 0 at the bottom to 1 at the top. The box is a bit lower than the top because of some padding. "upper left" means the upper left point of the legend is at the anchor point. If you don't provide an anchor, the legend is put inside the axes (overlapping with the squares in this case).

Collectives™ on Stack Overflow

Adding a legend to a matplotlib plot of a numpy array with binary data

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related