Animation based on only updating colours in a plot

Question

I have a plot which consists of great number of lines. At each step the colours of lines should get updated in the animation, but doing a for loop on lines seems to be really costly. Is there any better way to do that?

Here is my code:

import numpy as np
lines=[]
from matplotlib import pyplot as plt
import matplotlib.animation as animation

#initial plot
fig=plt.figure()
ax=plt.subplot(1,1,1)
for i in range(10):
    lines.append([])
    for j in range(10):
        lines[i].append(ax.plot([i,j],color='0.8'))
lines=np.asarray(lines)


##Updating the colors 10 times 
im=[]
for steps in range(10):
    colors=np.random.random(size=(10,10))
    for i in range(10):
        for j in range(10):
            lines[i,j][0].set_color(str(colors[i,j])) 
    plt.draw()
#    im.append(ax)
    plt.pause(.1)
#ani = animation.ArtistAnimation(fig, im, interval=1000, blit=True,repeat_delay=1000)
plt.show()

Plus I couldn't make it to work with animation artist! I used draw. What is wrong with the animation lines

Now increasing those 10s to 100 makes the program terribly slow:

import numpy as np
lines=[]
from matplotlib import pyplot as plt
import matplotlib.animation as animation

#initial plot
fig=plt.figure()
ax=plt.subplot(1,1,1)
for i in range(100):
    lines.append([])
    for j in range(100):
        lines[i].append(ax.plot([i,j],color='0.8'))
lines=np.asarray(lines)


##Updating the colors 10 times 
im=[]
for steps in range(10):
    colors=np.random.random(size=(100,100))
    for i in range(100):
        for j in range(100):
            lines[i,j][0].set_color(str(colors[i,j])) 
    plt.draw()
#    im.append(ax)
    plt.pause(.1)
#ani = animation.ArtistAnimation(fig, im, interval=1000, blit=True,repeat_delay=1000)
plt.show()

As I said I want to run it side by side with an animation. Therefore I prefer to make it an animation. I think that would solve the lagging problem at least after the animation starts but right now the way I defined it, it doesn't work.

Answer 1

It's easiest to use a LineCollection for this. That way you can set all of the colors as a single array and generally get much better drawing performance.

The better performance is mostly because collections are an optimized way to draw lots of similar objects in matplotlib. Avoiding the nested loops to set the colors is actually secondary in this case.

With that in mind, try something more along these lines:

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.collections import LineCollection
import matplotlib.animation as animation

lines=[]
for i in range(10):
    for j in range(10):
        lines.append([(0, i), (1, j)])

fig, ax = plt.subplots()
colors = np.random.random(len(lines))
col = LineCollection(lines, array=colors, cmap=plt.cm.gray, norm=plt.Normalize(0,1))
ax.add_collection(col)
ax.autoscale()

def update(i):
    colors = np.random.random(len(lines))
    col.set_array(colors)
    return col,

# Setting this to a very short update interval to show rapid drawing.
# 25ms would be more reasonable than 1ms.
ani = animation.FuncAnimation(fig, update, interval=1, blit=True, 
                              init_func=lambda: [col])
# Some matplotlib versions explictly need an `init_func` to display properly...
# Ideally we'd fully initialize the plot inside it. For simplicitly, we'll just
# return the artist so that `FuncAnimation` knows what to draw.
plt.show()

Answer 2

If you want to speed up a for loop, there are several good ways to do that. The best one for what you are trying to do, generator expressions, is probably like this:

iterator = (<variable>.upper() for <samevariable> in <list or other iterable object>)

(for more specific information on these there is documentation at http://www.python.org/dev/peps/pep-0289/ and https://wiki.python.org/moin/Generators)

There are also other, non-for loop ways to update color, but they are unlikely to be any faster than a generator. You could create some form of group for the lines, and call something like:

lines.update()

on all of them.