CODEWITHSUNDEEP
pythonarraysnumpymatplotlib
DJTS
DJTS

Reputation: 43

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

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: enter image description here

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

Upvotes: 4

Views: 1906

Answers (1)

JohanC
JohanC

Reputation: 80279

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()

resulting plot

Upvotes: 3

Related Questions