如何自定义Matplotlib图例的句柄与标签？Surpyval绘图适配问题

方法1：捕获绘图句柄，手动指定图例配对 #

Surpyval的plot()方法会返回本次绘图生成的所有Matplotlib Artist对象（如线条、散点等），我们可以捕获这些返回值，按需挑选句柄与对应标签匹配：

import surpyval as surv 
import matplotlib.pyplot as plt 
import numpy as np  # 补充缺失的numpy导入

y_a = np.array([181, 183, 190,190, 195, 195, 198, 198, 198, 201,202, 202, 202, 
       204, 205, 205, 206,206, 206, 206,207, 209 , 213, 214, 218, 219])

y_s = np.array([161, 179, 196,196, 197, 198, 204, 205, 209, 211,215, 218, 227, 
       230, 231, 232, 232 ,236, 237, 237,240, 243, 244, 246, 252, 255])

model_1 = surv.Weibull.fit(y_a)
model_2 = surv.Weibull.fit(y_s)

ax = plt.gca()

# 捕获每个模型绘图返回的句柄列表
handles_a = model_1.plot(ax=ax)
handles_s = model_2.plot(ax=ax)

# 定义对应标签，按绘图返回的句柄顺序匹配
labels_a = ['A', 'a_u_CI','a_l_CI', 'a_fit']
labels_s = ['S', 's_u_CI','s_l_CI', 's_fit']

# 合并句柄和标签生成完整图例
all_handles = handles_a + handles_s
all_labels = labels_a + labels_s
ax.legend(all_handles, all_labels)

plt.show()

如果仅需显示两个模型的主拟合线，可筛选每个模型返回句柄的第一个元素：

# 仅保留主拟合线的句柄与标签
selected_handles = [handles_a[0], handles_s[0]]
selected_labels = ['A', 'S']
ax.legend(selected_handles, selected_labels)

方法2：修改绘图元素的label属性 #

若不想捕获返回值，可在绘图后遍历Axes内的线条，手动设置label后调用图例：

ax = plt.gca()
model_1.plot(ax=ax)
model_2.plot(ax=ax)

# 获取Axes中所有线条对象
lines = ax.get_lines()

# 按绘图顺序对应设置标签
labels = ['A', 'a_u_CI','a_l_CI', 'a_fit', 'S', 's_u_CI','s_l_CI', 's_fit']
for line, label in zip(lines, labels):
    line.set_label(label)

ax.legend()
plt.show()

注意：此方法依赖绘图元素的生成顺序，需确保线条顺序与预期的模型绘图顺序一致。

内容的提问来源于stack exchange，提问作者A.E