重置Matplotlib中的颜色循环

Question

假设我有关于三种交易策略的数据，每种策略都有考虑和不考虑交易成本的情况。我想在同一张图上绘制这六种不同的情况（3种策略 * 2种交易成本）。我希望“考虑交易成本”的线条用 alpha=1 和 linewidth=1 来绘制，而“没有交易成本”的线条则用 alpha=0.25 和 linewidth=5 来绘制。不过，我希望每种策略的两种版本使用相同的颜色。

我想要的代码大概是这样的：

fig, ax = plt.subplots(1, 1, figsize=(10, 10))

for c in with_transaction_frame.columns:
    ax.plot(with_transaction_frame[c], label=c, alpha=1, linewidth=1)

****SOME MAGIC GOES HERE TO RESET THE COLOR CYCLE

for c in no_transaction_frame.columns:
    ax.plot(no_transaction_frame[c], label=c, alpha=0.25, linewidth=5)

ax.legend()

我想知道在指定的那一行上，应该写什么代码来重置颜色循环，这样在第二次循环时颜色就能“回到起点”。

Answer 1

你可以这样从seaborn获取颜色：colors = sns.color_palette()。这样的话，Ffisegydd的回答就能很好地使用了。你也可以通过取余运算符（%）来获取绘图的颜色：mycolor = colors[icolumn % len(colors)]。我自己也经常用这种方法。所以你可以这样做：

for icol, column in enumerate(with_transaction_frame.columns): mycolor = colors[icol % len(colors)] ax.plot(with_transaction_frame[col], label=col, alpha=1.0, color=mycolor)

不过，Ffisegydd的回答可能更符合“python风格”。

Answer 2

简单来说，你只需要选择好颜色，然后把它们放到一个列表里。接着，在绘制数据的时候，可以使用一个叫做zip的工具，把你的数据列和你想要的颜色配对起来。

colors = ['red', 'blue', 'green']

for col, color in zip(colors, with_transaction_frame.columns):
    ax.plot(with_transaction_frame[col], label=col, alpha=1.0, linewidth=1.0, color=color)

for col, color in zip(no_transaction_frame.columns):
    ax.plot(no_transaction_frame[col], label=col, alpha=0.25, linewidth=5, color=color)

zip这个工具会把你所有列表里的元素组合成一个新的列表。这样，你就可以很方便地同时处理这两个列表里的内容了。

Answer 3

既然你提到你在使用seaborn，我建议你可以这样做：

with sns.color_palette(n_colors=3):

    ax.plot(...)
    ax.plot(...)

这样做会把颜色调色板设置为当前活动的颜色循环，但只会使用其中的前三种颜色。这也是一个通用的解决方案，适合任何时候你想设置临时颜色循环的情况。

需要注意的是，实际上只有创建Axes对象的部分需要放在with块里面（比如plt.subplots、fig.add_subplot()等）。这是因为matplotlib的颜色循环是这样工作的。

如果你想要的具体操作是“重置”颜色循环，这也是可以做到的，但这是一种小技巧，我不建议在任何正式的代码中使用。不过，这里有一种实现方法：

f, ax = plt.subplots()
ax.plot(np.random.randn(10, 3))
ax._get_lines.color_cycle = itertools.cycle(sns.color_palette())
ax.plot(np.random.randn(10, 3), lw=5, alpha=.25)

Answer 4

由于@pelson给出的答案使用了set_color_cycle，而这个在Matplotlib 1.5中已经不再推荐使用，所以我觉得更新一下他的解决方案，改用set_prop_cycle会更有帮助：

import matplotlib.pyplot as plt
import numpy as np

for i in range(3):
    plt.plot(np.arange(10) + i)

plt.gca().set_prop_cycle(None)

for i in range(3):
    plt.plot(np.arange(10, 0, -1) + i)

plt.show()

另外要注意，我把np.arange(10,1,-1)改成了np.arange(10,0,-1)。前者只生成了9个元素的数组。这可能是因为使用了不同版本的Numpy。我这边用的是1.10.2。

编辑：去掉了使用rcParams的必要。感谢@divenex在评论中指出这一点。

Answer 5

你可以通过 Axes.set_color_cycle 来把颜色循环重置回最初的状态。看一下这个代码，它里面有一个函数专门用来完成这个工作：

def set_color_cycle(self, clist=None):
    if clist is None:
        clist = rcParams['axes.color_cycle']
    self.color_cycle = itertools.cycle(clist

还有一个在 Axes 上的方法会使用这个函数：

def set_color_cycle(self, clist):
    """
    Set the color cycle for any future plot commands on this Axes.

    *clist* is a list of mpl color specifiers.
    """
    self._get_lines.set_color_cycle(clist)
    self._get_patches_for_fill.set_color_cycle(clist)

简单来说，你只需要用 None 作为唯一的参数来调用 set_color_cycle，它就会用默认的颜色循环来替换掉当前的颜色循环，这个默认值在 rcParams['axes.color_cycle'] 中。

我用下面的代码试了一下，得到了预期的结果：

import matplotlib.pyplot as plt
import numpy as np

for i in range(3):
    plt.plot(np.arange(10) + i)

# for Matplotlib version < 1.5
plt.gca().set_color_cycle(None)
# for Matplotlib version >= 1.5
plt.gca().set_prop_cycle(None)

for i in range(3):
    plt.plot(np.arange(10, 1, -1) + i)

plt.show()