Speed-up matplotlib creation of frame images for animation

Question

I am simulating some physical processes with Python and wish to regularly plot the situation of the system in order to make an animation. What I have been doing so far is to draw a snapshot of the system at regular intervals, save it to file, and finally convert the images into a video using convert or ffmpeg. Lately, however, the plot has been growing in complexity and it takes now 2-3 seconds to generate just one image, which is too long for my needs.

In my system, I have some background elements that are fixed and never change, and some other elements that are subject to movements depending on time. As a minimum working example, consider the following code:

import matplotlib.pyplot as plt
import numpy as np
import time


class MWE():
    def __init__(self):
        self.x = 0.0
        self.y = 0.0
        self.dt = 0.01
        self.imageCounter = 0
        self.elapsed_time = 0.0

    def evolve(self):
        theta = 2 * np.pi * self.elapsed_time
        self.x = np.cos(theta)
        self.y = np.sin(theta)
        self.elapsed_time += self.dt

    def draw(self):
        fig, ax = plt.subplots(figsize=(2, 2), layout='constrained')
        ax.set_aspect('equal')
        ax.axis('off')

        ax.set_xlim((-2, 2))
        ax.set_ylim((-2.2, 2))

        for center_x in np.linspace(-2,2, 21)[1::2]:
            for center_y in np.linspace(-2,2, 21)[1::2]:
                blue_circle = plt.Circle((center_x, center_y), 0.1, color='blue', fill=False)
                ax.add_patch(blue_circle)

        red_circle = plt.Circle((self.x, self.y), 0.1, color='red', fill=True)
        ax.add_patch(red_circle)
        self.imageCounter += 1

        ax.text(-1.9, -2.2, f'Time: {self.elapsed_time:.2f} s', color='black', size=8)

        fig.savefig(f'MWE_{str(self.imageCounter).zfill(3)}.png')
        plt.clf()
        plt.cla()
        plt.close()


mwe = MWE()

start_time = time.time()
while mwe.elapsed_time < 1.0:
    mwe.evolve()
    mwe.draw()
print(f"Execution: {time.time() - start_time:.2f} s" % ())

convert -resize 100% -delay 2 -loop 0 MWE_*.png loop.gif

The execution on my computer takes Execution: 51.32 s, therefore I am looking for alternative way to handle the plots, be it keeping in memory the background drawing, or using an additional package for the animation.

Answer 1

Without any fancy blitting (which probably speeds up things even more dramatically), these few simple changes bring execution time down to Execution: 4.10 s on my side:

• initialize figure only once
• draw artists only once and update positions
• switch to agg backend to avoid unnecessary guis

import matplotlib.pyplot as plt
import numpy as np
import time

plt.switch_backend("agg")

class MWE():
    def __init__(self):
        self.x = 0.0
        self.y = 0.0
        self.dt = 0.01
        self.imageCounter = 0
        self.elapsed_time = 0.0
        
        self.init_figure()

    def init_figure(self):
        self.fig, ax = plt.subplots(figsize=(2, 2), layout='constrained')
        ax.set_aspect('equal')
        ax.axis('off')

        ax.set_xlim((-2, 2))
        ax.set_ylim((-2.2, 2))

        for center_x in np.linspace(-2,2, 21)[1::2]:
            for center_y in np.linspace(-2,2, 21)[1::2]:
                blue_circle = plt.Circle((center_x, center_y), 0.1, color='blue', fill=False)
                ax.add_patch(blue_circle)

        self.red_circle = plt.Circle((self.x, self.y), 0.1, color='red', fill=True)
        ax.add_patch(self.red_circle)
        self.info_text = ax.text(-1.9, -2.2, f'Time: {self.elapsed_time:.2f} s', color='black', size=8)

    def evolve(self):
        theta = 2 * np.pi * self.elapsed_time
        self.x = np.cos(theta)
        self.y = np.sin(theta)
        self.elapsed_time += self.dt

    def draw(self):
        self.red_circle.set_center((self.x, self.y))
        self.info_text.set_text(f'Time: {self.elapsed_time:.2f} s')
        self.imageCounter += 1

        self.fig.savefig(f'MWE_{str(self.imageCounter).zfill(3)}.png')


mwe = MWE()
mwe.init_figure()

# %%
start_time = time.time()
while mwe.elapsed_time < 1.0:
    mwe.evolve()
    mwe.draw()
print(f"Execution: {time.time() - start_time:.2f} s" % ())