将图例添加到具有二进制数据的 numpy 数组的 matplotlib plot

Question

我正在编写一个程序来模拟放射性衰变。 我正在使用 class 和一个decay_sim方法，该方法生成一个二维numpy数组，其中一个（未衰减）和零（衰减）称为self.nuclei 。 这是一个伪随机过程，所以每次的数组都不一样。 然后我使用visual方法将 plot 这个数组作为一个颜色网格使用matplotlib 。

这是self.nuclei数组的示例：

[[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]]

这是visual方法：

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

这是它的 output：

我正在寻找一种方法来为这个 plot 添加一个图例，它可以 label 衰减和未衰减的颜色。 我目前正在使用plt.colorbar() ，但这只会给我价值 - 而不是 label。 理想情况下，它应该看起来像：“红色 = 未衰减”和“蓝色 = 衰减”（尽管我希望稍后可以更改配色方案）。 我还想在图例中包括腐烂/未腐烂细胞的数量。 我想我可以使用以下方法获取号码：

import numpy as np

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

Answer 1

您可以从矩形补丁创建自定义图例。 网格的 colors 可以通过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()