在matplotlib中绘图时循环颜色：按实例跟踪状态

Question

我正在尝试为Matplotlib中的Axes实例构建一个简单的状态跟踪功能。每当我创建一个新的坐标轴对象（无论是直接创建还是通过其他函数，比如subplots()），我希望这个实例有一个绑定的方法a.next_color()，这样我就可以在为坐标轴添加新线条时循环使用颜色。我写了类似这样的代码：

def set_color_sequence(colors = ['r', 'g', 'b', 'c', 'm', 'y']):
    i = [0]
    def cf(self):
        i[0] += 1
        return colors[(i[0]-1) % len(colors)]
    return cf

我还觉得自己很聪明，把它添加到了父类中：

plt.Axes.next_color = set_color_sequence()

问题是状态变量i似乎被所有的Axes实例共享，而不是每个新实例都有自己的状态。有没有什么优雅的方法可以确保所有新实例都有自己的状态跟踪功能？（顺便说一下，我希望在不修改原始matplotlib代码的情况下做到这一点。）

Answer 1

matplotlib 其实已经有类似的功能了。假设你有一个叫 ax1 的对象，它是 AxesSubplot 的一个实例（或者是任何从 Axes 衍生出来的东西），你可以通过 _get_lines 来访问线条的属性。这些属性包括 prop_cycler，你可以通过以下方式来覆盖它：

ax2._get_lines.prop_cycler = ax1._get_lines.prop_cycler

举个例子：

# Generate some data
x = linspace(-10, 10)
y_lin = 2*x
y_quad = x**2

# Create empty axes
_, ax1 = subplots()

# Plot linear
l1 = ax1.plot(y_lin)
ax1.set_ylabel('linear', color=l1[0].get_color())
ax1.tick_params('y', colors=l1[0].get_color())

# Generate secondary y-axis
ax2 = ax1.twinx()
# Make sure that plots on ax2 continue color cycle
ax2._get_lines.prop_cycler = ax1._get_lines.prop_cycler

# Plot quadratic
l2 = ax2.plot(y_quad)
ax2.set_ylabel('quadratic', color=l2[0].get_color())
ax2.tick_params('y', colors=l2[0].get_color())

这样会输出：

如果不执行 ax2._get_lines.prop_cycler = ax1._get_lines.prop_cycler，我们得到的结果是：

注意：如果你在 jupyter（或者 nteract）中使用 %pylab inline 这个魔法命令（这样你就不能直接访问 matplotlib.pyplot 模块），那么像 plt.Axes 这样的用法会有点麻烦。

Answer 2

你现在的函数可以正常工作，只要你把next_color这个属性赋值给Axes的一个实例，而不是直接赋给类本身。

首先，通过set_color_sequence，你实际上是在间接实现一个生成器。为了简化，我们可以用一行代码通过itertools.cycle来实现同样的功能：

from itertools import cycle
...
axes_instance.next_color = cycle(['r', 'g', 'b', 'c', 'm', 'y']).next

实际上，这就是matplotlib用来跟踪颜色循环进度的方法。例如，如果你查看matplotlib.axes._subplots.AxesSubplot的一个实例，你会发现它有一个属性_get_lines.color_cycle，这个属性是一个itertools.cycle（你可以试着调用color_cycle.next()）。

现在看看这两个例子：

class MyClass1(object):

    # next_color is an attribute of the *class itself*
    next_color = cycle(['r', 'g', 'b', 'c', 'm', 'y']).next


class MyClass2(object):

    def __init__(self):

        # next_color is an attribute of *this instance* of the class
        self.next_color = cycle(['r', 'g', 'b', 'c', 'm', 'y']).next

在第一个例子中，赋值操作

next_color = cycle(['r', 'g', 'b', 'c', 'm', 'y]).next

只会在类第一次被导入时执行一次。这意味着每当你创建一个新的MyClass1实例时，它的next_color属性都会指向同一个itertools.cycle实例，因此所有的MyClass1实例都会共享同一个状态：

a = MyClass1()
b = MyClass1()

print a.next_color is b.next_color
# True

print a.next_color(), a.next_color(), b.next_color()
#  r g b

然而，__init__方法在每次创建新类的实例时都会被调用。因此，每个MyClass2的实例都会有自己的itertools.cycle，也就是说每个实例都有自己的状态：

a = MyClass2()
b = MyClass2()

print a.next_color is b.next_color
# False

print a.next_color(), a.next_color(), b.next_color()
# r g r

如果你的目标是继承plt.Axes，你需要把赋值放在一个地方，这样每个新实例都会调用它（可能是在__init__里）。但是，如果你只是想把这个方法添加到一个现有的实例中，那么你只需要做：

axes_instance.next_color = get_color_sequence()

Answer 3

AMacK的评论让我找到了一个简单的解决办法：

def next_color(self):
    return next(self._get_lines.color_cycle)
plt.Axes.next_color = next_color

这个方法虽然没有我之前提到的自定义颜色顺序的功能，但它确实实现了我想要的每个实例都有不同表现的效果，而且代码简洁易懂。（如果我想要生成一个自定义的颜色顺序，应该可以直接重写颜色循环的迭代器。）