如何为不同的圆形标签设置颜色图例

时间:2019-08-04 13:03:44

标签: matplotlib

我正在用Matplotlib绘制不同的圆圈。每个圆圈都有一个标签,每个标签都有一种颜色。我应该怎么做才能为这些不同的标签添加颜色图例?

我已经在网上尝试了很多解决方案,包括仅通过添加plt.colorbar()来解决的最幼稚的解决方案,

RuntimeError('No mappable was found to use for colorbar')

这是我完整的代码。有点长。请注意,关键部分仅从if labels is None:开始。我只是为了完整性而包括所有内容。

def plot_gaussian_circles(loc_list, scale_list, save_path=None, sigma_coe=3, num_to_plot=300, labels=None):
    mu_x_max = -float('inf')
    mu_y_max = -float('inf')

    mu_x_min = float('inf')
    mu_y_min = float('inf')

    color_idx = 0

    rvs = []

    lim_loc_list = loc_list[:num_to_plot]
    lim_scale_list = scale_list[:num_to_plot]
    for a_mu_, a_sigma_ in zip(lim_loc_list, lim_scale_list):
        a_mu = a_mu_.squeeze()
        a_sigma_ = a_sigma_.squeeze()
        if not type(a_sigma_) is np.ndarray:
            a_sigma_ = a_sigma_.numpy()

        radius = sigma_coe * np.max(a_sigma_)
        a_mu_x = a_mu[0]
        a_mu_y = a_mu[1]

        if (a_mu_x + radius) >= mu_x_max:
            mu_x_max = a_mu_x + radius
        if (a_mu_x - radius) <= mu_x_min:
            mu_x_min = a_mu_x - radius

        if (a_mu_y + radius) >= mu_y_max:
            mu_y_max = a_mu_y + radius
        if (a_mu_y - radius) <= mu_y_min:
            mu_y_min = a_mu_y - radius

        if labels is None:
            rv = plt.Circle(a_mu, radius, fill=False, clip_on=False)
        else:
            colors = cm.rainbow(np.linspace(0, 1, len(set(labels))))
            rv = plt.Circle(a_mu, radius, color=colors[labels[color_idx]], fill=False, clip_on=False)

        rvs.append(rv)
        color_idx = (color_idx + 1)

    fig, ax = plt.subplots()
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')

    axes = plt.gca()
    axes.set_xlim([mu_x_min - 1, mu_x_max + 1])
    axes.set_ylim([mu_y_min - 1, mu_y_max + 1])

    for rv in rvs:
        ax.add_artist(rv)

    if not(labels is None):
        # plt.legend(colors, list(range(len(set(labels)))))
        plt.colorbar()
    if save_path is None:
        plt.plot()
        plt.show()
        # plt.savefig('plotcircles_test.png')
    else:
        plt.savefig(save_path, dpi=200)

这里的图像是我目前正在获得的图像,而我希望有一个色彩图的图例。 enter image description here

1 个答案:

答案 0 :(得分:0)

我发现这样做会为圆圈添加颜色。感谢ImportanceOfBeingErnest的评论建议PatchCollection。

p = PatchCollection(rvs, cmap=cm.jet, alpha=0.4)
p.set_array(labels)
ax.add_collection(p)
fig.colorbar(p, ax=ax)