显示 matplotlib 动态创建的图例

Question

我的 df 有 4 列：x、y、z 和分​​组。我创建了一个 3D 图，每个点的指定颜色由它在该行中所属的分组决定。作为参考，“分组”可以是从 1 到 6 的任意数字。代码如下所示：

fig = plt.figure()
ax = Axes3D(fig)
ax.scatter3D(df.x, df.y, df.z, c=df.grouping)
plt.show()

我想在图上显示一个图例，显示哪种颜色属于哪个分组。以前，我使用 Seaborn 进行 2D 绘图，并且自动绘制了图例。如何在 matplotlib 中添加此功能？

Answer 1

如果要进行颜色映射的值是数字，则解决方案可以很简单：

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

a = np.random.rand(3,40)
c = np.random.randint(1,7, size=a.shape[1])

fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
sc = ax.scatter3D(*a, c=c)
plt.legend(*sc.legend_elements())
plt.show()