Matplotlib如何从已绘图的轴获取Colorbar的Mappable对象？

先给你泼个小冷水：Matplotlib里的轴（Axes）对象并没有内置的get_mappable()方法——毕竟一个轴上可能同时有好几个带颜色映射的绘图元素，程序没法自动判定你要哪一个。不过别慌，我们有几种实用的办法解决这个问题，而且得先纠正你示例里的一个小误区：你用的ax.plot()生成的Line2D对象，哪怕设置了c为颜色数组，它也不是一个“mappable”对象（mappable需要具备get_cmap()、get_norm()这类颜色映射相关的方法）。如果要实现颜色映射+colorbar的效果，你得用scatter()、imshow()、pcolormesh()这类会返回mappable对象的绘图函数。

下面是具体的解决方案：

方案1：修改绘图函数，返回mappable对象 #

这是最直接也最稳妥的方法，只要给你的plot_something函数加个返回值就行：

import numpy as np
import matplotlib.pyplot as plt

def plot_something(ax):
    # 用scatter代替plot，因为scatter会返回mappable的PathCollection对象
    sc = ax.scatter(np.random.random(10), np.random.random(10), c=np.random.random(10))
    return sc

fig, axs = plt.subplots(2)
# 接收第一个轴的mappable对象
mappable0 = plot_something(axs[0])
plot_something(axs[1])
# 用返回的mappable创建colorbar
fig.colorbar(mappable0, ax=axs[0])
plt.show()

方案2：遍历轴上的艺术家，手动查找mappable #

如果没法修改plot_something函数，你可以遍历轴上的所有绘图元素（艺术家），筛选出带有颜色映射属性的对象：

import numpy as np
import matplotlib.pyplot as plt

def plot_something(ax):
    ax.scatter(np.random.random(10), np.random.random(10), c=np.random.random(10))

def get_mappable_from_ax(ax):
    # 遍历轴上的所有子元素
    for artist in ax.get_children():
        # 判断是否是mappable（具备cmap和norm属性）
        if hasattr(artist, 'cmap') and hasattr(artist, 'norm'):
            return artist
    return None

fig, axs = plt.subplots(2)
plot_something(axs[0])
plot_something(axs[1])

# 从轴中提取mappable
mappable = get_mappable_from_ax(axs[0])
if mappable:
    fig.colorbar(mappable, ax=axs[0])
plt.show()

⚠️ 注意：如果轴上有多个mappable对象，这个方法会返回第一个找到的，可能不符合你的预期；另外如果绘图用的是plot()，这个方法找不到有效对象，因为Line2D不是mappable。

方案3：手动创建ScalarMappable（适配plot的情况） #

如果你坚持要用ax.plot()，可以在绘图函数里把颜色范围和映射规则存到轴的自定义属性中，然后手动创建mappable：

import numpy as np
import matplotlib.pyplot as plt

def plot_something(ax):
    x = np.random.random(10)
    y = np.random.random(10)
    colors = np.random.random(10)
    # 用plot绘制，同时把颜色信息存到轴的自定义属性里
    ax.plot(x, y, c=plt.cm.viridis(plt.Normalize(0,1)(colors)))
    ax._color_info = {
        'cmap': plt.get_cmap('viridis'),
        'norm': plt.Normalize(vmin=colors.min(), vmax=colors.max())
    }

fig, axs = plt.subplots(2)
plot_something(axs[0])
plot_something(axs[1])

# 从轴的自定义属性中提取颜色信息，创建ScalarMappable
color_info = axs[0]._color_info
mappable = plt.cm.ScalarMappable(norm=color_info['norm'], cmap=color_info['cmap'])
fig.colorbar(mappable, ax=axs[0])
plt.show()

这个方法需要你在绘图时手动处理颜色的映射，适合必须用plot()且不能修改函数返回值的场景。

内容的提问来源于stack exchange，提问作者Hoseung Choi