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

Question

我正在编写一个程序来模拟放射性衰变。我正在使用一个带有 decay_sim 方法的类，它生成一个二维 numpy 数组，其中包含一个(未衰减)和零(衰减)，称为 self.nuclei。这是一个伪随机过程，所以每次的数组都不一样。然后，我使用 visual 方法使用 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]]

这是视觉方法:

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

这是它的输出:

我正在寻找一种方法来为该图添加图例，以标记衰减和未衰减的颜色。我目前正在使用 plt.colorbar()，但这只给我值 - 而不是标签。理想情况下，它应该看起来像:“red = not decayed”和“blue = decayed”(尽管我希望以后可以选择更改配色方案)。我还想在图例中包括腐烂/未腐烂细胞的数量。我想我可以使用以下方法获取号码:

import numpy as np

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

Answer 1

您可以创建一个 custom legend从矩形补丁。网格的颜色可以通过 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()