Matplotlib: change colormap after the fact

Question

I'm charting the progress of a differential equation solver (boundary value problem). Each iteration yields a complete set of function evaluations f(x), which can then be plotted against x. Each graph is (supposedly) closer to the correct solution than the last until convergence is reached. A sequential colormap is used to make earlier graphs faded and later ones saturated.

This works fine when the number of iterations is predetermined:

import matplotlib.pyplot as plt

ax = plt.subplot(111)
cm = plt.get_cmap('OrRd')
ax.set_color_cycle([cm(1.*i/(iter+1)) for i in range(1,iter+2)]) 

ax.plot(x,y)
for k in range(iter):
    # iterative solve 
    ax.plot(x,y)

However, if I use a convergence criterion instead of a predetermined number of iterations, I won't be able to set_color_cycle beforehand. And putting that line after the loop doesn't work.

I know that I can store my intermediate results and plot only after convergence is reached, but this strikes me as heavy-handed because I really have no use for all the intermediate results other than to see them on the plot.

So here are my questions: 1. How do I change the colormap of the existing graphs after plotting? (This is easy in MATLAB.) 2. How do I do the same thing with another collection of graphs on the same plot (e.g. from a different initial guess, converging to a different solution) without disturbing the first collection, so that two colormaps distinguish the collections from one another. (This should be obvious with the answer to Question 1, but just in case.)

Many thanks.

Answer 1

You can also use plt.set_cmap, see here or (more elaborately, scroll down) here:

import numpy as np
import matplotlib.pyplot as plt

plt.imshow(np.random.random((10,10)), cmap='magma')
plt.colorbar()
plt.set_cmap('viridis')

Answer 2

Use the update_colors() to update the colors of all lines:

import pylab as pl
import numpy as np
cm = pl.get_cmap('OrRd')

x = np.linspace(0, 1, 100)

def update_colors(ax):
    lines = ax.lines
    colors = cm(np.linspace(0, 1, len(lines)))
    for line, c in zip(lines, colors):
        line.set_color(c)

fig, ax = pl.subplots()
for i in range(10):
    ax.plot(x, x**(1+i*0.1))
    update_colors(ax)