GitHub - ali-ramadhan/matplotloom: Weave your frames into matplotlib animations

matplotloom

🧵🧶🪡Weave your frames into matplotlib animations!

Installation

matplotloom is published on PyPI so you can install matplotloom via pip

or uv

or poetry

or conda

conda install matplotloom

matplotloom is continuously tested with Python 3.10+ on Linux, Windows, and Mac.

You will need to have ffmpeg installed so that animations can be generated.

Why use matplotloom?

To visualize simulation output for computational fluid dynamics I've had to make long animations with complex figures for a long time. The animations consist of thousands of frames and the figures are too complex for FuncAnimation and ArtistAnimation. This package aims to simplify and massively speed up the process of making these kinds of animations.

The main idea behind matplotloom is to describe how to generate each frame of your animation from scratch, instead of generating an animation by modifying one existing plot. This simplifies generating animations. See the examples below and how the code inside the for loops is plain and familiar matplotlib. It also ensures that every feature can be animated and that the generation process can be easily parallelized.
matplotlib has two tools for making animations: FuncAnimation and ArtistAnimation. But to use them you have to write your plotting code differently to modify an existing frame. This makes it difficult to go from plotting still figures to making animations. And some features are non-trivial to animate.
celluloid is a nice package for making matplotlib animations easily, but as it relies on ArtistAnimation under the hood it does come with some limitations such as not being able to animate titles. It also hasn't been maintained since 2018.
Plotting many frames (hundreds to thousands+) can be slow but with matplotloom you can use a parallel Loom to plot each frame in parallel, speeding up the animation process significantly especially if you can dedicate many cores to plotting.

Some notes to users

You can use loom.show() to display animations in Jupyter notebooks.
Anxious about animation progress? Pass verbose=True or use tqdm to monitor progress.
Animations taking too long to make or do you have tons of frames? You can parallelize frame creating by looming in parallel.
You have to call loom.save_frame(fig) for each frame (see the examples). While the Loom object can be made to do this automatically it would have to create and own the Figure instance and I wanted full control over the creation of the Figure for maximum flexibility.
matplotloom is going to be slow. But it's flexible and compatible with all of matplotlib! The real speedup comes from parallelizing frame creation, especially if you have a ton of frames to make.

Examples

Sine wave

import numpy as np
import matplotlib.pyplot as plt
from matplotloom import Loom

with Loom("sine_wave_animation.gif", fps=30) as loom:
    for phase in np.linspace(0, 2*np.pi, 100):
        fig, ax = plt.subplots()

        x = np.linspace(0, 2*np.pi, 200)
        y = np.sin(x + phase)

        ax.plot(x, y)
        ax.set_xlim(0, 2*np.pi)

        loom.save_frame(fig)

Rotating circular sine wave

import numpy as np
import matplotlib.pyplot as plt
from matplotloom import Loom

with Loom("rotating_circular_sine_wave.mp4", fps=10) as loom:
    for i in range(36):
        fig, ax = plt.subplots(figsize=(12, 8), subplot_kw={"projection": "3d"})

        X = np.arange(-5, 5, 0.25)
        Y = np.arange(-5, 5, 0.25)
        X, Y = np.meshgrid(X, Y)
        R = np.sqrt(X**2 + Y**2)
        Z = np.sin(R)

        surf = ax.plot_surface(X, Y, Z, cmap="coolwarm")

        ax.view_init(azim=i*10)
        ax.set_zlim(-1.01, 1.01)
        fig.colorbar(surf, shrink=0.5, aspect=5)

        loom.save_frame(fig)

rotating_circular_sine_wave.mp4

Bessel wave

Compare with animatplot's blocks example. With matplotloom you just use regular matplotlib abstractions.

import numpy as np
import matplotlib.pyplot as plt
from scipy.special import j0
from cmocean import cm
from matplotloom import Loom

def bessel_wave(r, t, k, omega, A):
    return A * j0(k*r - omega*t)

def create_frame(x, y, t):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    r = np.sqrt(x**2 + y**2)
    z = bessel_wave(r, t, k=2, omega=1, A=1)

    pcm = ax1.pcolormesh(x, y, z, cmap=cm.balance, shading='auto', vmin=-1, vmax=1)
    fig.colorbar(pcm, ax=ax1)

    ax1.set_title(f"Bessel wave: t = {t:.3f}")
    ax1.set_xlabel("x")
    ax1.set_ylabel("y")
    ax1.set_xlim(-10, 10)
    ax1.set_ylim(-10, 10)
    ax1.set_aspect("equal", adjustable="box")

    mid = z.shape[0] // 2
    ax2.plot(x[mid], z[mid])

    ax2.set_xlim(x.min(), x.max())
    ax2.set_ylim(-1.1, 1.1)
    ax2.set_title("Cross-section at y = 0")
    ax2.set_xlabel("x")
    ax2.set_ylabel("z")

    return fig

loom = Loom(
    "bessel_wave.mp4",
    fps = 30,
    overwrite = True,
    verbose = True,
    savefig_kwargs = {
        "dpi": 100,
        "bbox_inches": "tight"
    }
)

with loom:
    x = np.linspace(-10, 10, 500)
    y = np.linspace(-10, 10, 500)
    x, y = np.meshgrid(x, y)

    for t in np.linspace(0, 50, 300):
        fig = create_frame(x, y, t)
        loom.save_frame(fig)

bessel_wave.mp4

Double pendulum

Compare with matplotlib's double pendulum example.

import numpy as np
import matplotlib.pyplot as plt

from scipy.integrate import solve_ivp
from tqdm import tqdm
from matplotloom import Loom

g = 9.80665 # standard acceleration of gravity [m/s²]
l1, l2 = 1, 1  # pendulum arms lengths [m]
m1, m2 = 1, 1  # pendulum masses [kg]

# Calculate dy/dt where y = [θ1, ω1, θ2, ω2].
def derivatives(t, state):
    θ1, ω1, θ2, ω2 = state
    dydt = np.zeros_like(state)

    dydt[0] = ω1
    dydt[2] = ω2

    Δθ = θ2 - θ1

    denominator1 = (m1 + m2) * l1 - m2 * l1 * np.cos(Δθ)**2
    dydt[1] = (m2 * l1 * ω1**2 * np.sin(Δθ) * np.cos(Δθ)
               + m2 * g * np.sin(θ2) * np.cos(Δθ)
               + m2 * l2 * ω2**2 * np.sin(Δθ)
               - (m1 + m2) * g * np.sin(θ1)) / denominator1

    denominator2 = (l2 / l1) * denominator1
    dydt[3] = (-m2 * l2 * ω2**2 * np.sin(Δθ) * np.cos(Δθ)
               + (m1 + m2) * g * np.sin(θ1) * np.cos(Δθ)
               - (m1 + m2) * l1 * ω1**2 * np.sin(Δθ)
               - (m1 + m2) * g * np.sin(θ2)) / denominator2

    return dydt

t_span = (0, 20)
y0 = [np.pi/2, 0, np.pi/2, 0]
sol = solve_ivp(derivatives, t_span, y0, dense_output=True)

times = np.linspace(t_span[0], t_span[1], 1000)

θ1, ω1, θ2, ω2 = sol.sol(times)
x1 =  l1 * np.sin(θ1)
y1 = -l1 * np.cos(θ1)
x2 = x1 + l2 * np.sin(θ2)
y2 = y1 - l2 * np.cos(θ2)

loom = Loom(
    "double_pendulum.mp4",
    fps = 60,
    overwrite = True,
    savefig_kwargs = {"bbox_inches": "tight"}
)

with loom:
    for i, t in tqdm(enumerate(times), total=len(times)):
        fig, ax = plt.subplots(figsize=(8, 8))

        ax.plot(
            [0, x1[i], x2[i]],
            [0, y1[i], y2[i]],
            linestyle = "solid",
            marker = "o",
            color = "black",
            linewidth = 3
        )

        ax.plot(
            x2[:i+1],
            y2[:i+1],
            linestyle = "solid",
            linewidth = 2,
            color = "red",
            alpha = 0.5
        )

        ax.set_title(f"Double Pendulum: t = {t:.3f}s")

        ax.set_xlim(-2.2, 2.2)
        ax.set_ylim(-2.2, 2.2)
        ax.set_aspect("equal", adjustable="box")

        loom.save_frame(fig)

double_pendulum.mp4

Night time shading

matplotloom works out of the box with anything that is built on top of matplotlib. Here we're extending a Cartopy example.

import datetime

import matplotlib.pyplot as plt
import cartopy.crs as ccrs

from cartopy.feature.nightshade import Nightshade
from joblib import Parallel, delayed
from matplotloom import Loom

def plot_frame(day_of_year, loom, frame_number):
    date = datetime.datetime(2024, 1, 1, 12) + datetime.timedelta(days=day_of_year-1)

    fig = plt.figure(figsize=(15, 5))

    proj1 = ccrs.Orthographic(central_longitude=0, central_latitude=30)
    proj2 = ccrs.Orthographic(central_longitude=120, central_latitude=0)
    proj3 = ccrs.Orthographic(central_longitude=240, central_latitude=-30)

    ax1 = fig.add_subplot(1, 3, 1, projection=proj1)
    ax2 = fig.add_subplot(1, 3, 2, projection=proj2)
    ax3 = fig.add_subplot(1, 3, 3, projection=proj3)

    fig.suptitle(f"Night time shading for {date} UTC")

    ax1.stock_img()
    ax1.add_feature(Nightshade(date, alpha=0.2))

    ax2.stock_img()
    ax2.add_feature(Nightshade(date, alpha=0.2))

    ax3.stock_img()
    ax3.add_feature(Nightshade(date, alpha=0.2))

    loom.save_frame(fig, frame_number)

loom = Loom(
    "night_time_shading.mp4",
    fps = 10,
    overwrite = True,
    parallel = True,
    verbose = True,
    savefig_kwargs = {
        "bbox_inches": "tight"
    }
)

with loom:
    n_days_2024 = 366
    days_of_year = range(1, n_days_2024 + 1)

    Parallel(n_jobs=-1)(
        delayed(plot_frame)(day_of_year, loom, i)
        for i, day_of_year in enumerate(days_of_year)
    )

night_time_shading.mp4

Lorenz Attractor

This example shows how you can use matplotloom to create an animation of a Lorenz attractor (inspired by an example from the Makie library).

from dataclasses import dataclass, field

import matplotlib.pyplot as plt
import numpy as np
from joblib import Parallel, delayed
from matplotlib.colors import Normalize
from matplotloom import Loom
from mpl_toolkits.mplot3d.art3d import Line3DCollection


@dataclass
class Lorenz:
    dt: float = 0.01
    sigma: float = 10.0
    rho: float = 28.0
    beta: float = 8.0 / 3.0
    x: float = 1.0
    y: float = 1.0
    z: float = 1.0

    def step(self):
        dx = self.sigma * (self.y - self.x)
        dy = self.x * (self.rho - self.z) - self.y
        dz = self.x * self.y - self.beta * self.z
        self.x += dx * self.dt
        self.y += dy * self.dt
        self.z += dz * self.dt

    @property
    def position(self) -> tuple[float, float, float]:
        return self.x, self.y, self.z


@dataclass
class LorenzPlotter:
    steps_per_frame: int = 20
    attractor = Lorenz()
    points: list[tuple[float, float, float]] = field(default_factory=list)

    def initialize(self, steps: int):
        self.points = [self.attractor.position]
        for _ in range(steps):
            self.attractor.step()
            self.points.append(self.attractor.position)

    @property
    def frames(self) -> list[int]:
        return list(range(1, len(self.points) // self.steps_per_frame))

    def get_frame(self, i: int, loom: Loom):
        fig, ax = plt.subplots(figsize=(12, 8), subplot_kw={'projection': '3d'})
        points = np.array(self.points[: i * self.steps_per_frame])
        xs, ys, zs = points.T
        segments = np.array([points[:-1], points[1:]]).transpose(1, 0, 2)
        norm = Normalize(vmin=0, vmax=len(xs))
        colors = plt.get_cmap('inferno')(norm(np.arange(len(xs) - 1)))
        lc = Line3DCollection(segments, colors=colors, linewidth=0.5)
        ax.add_collection3d(lc)
        ax.set_xlim(-30, 30)
        ax.set_ylim(-30, 30)
        ax.set_zlim(0, 50)
        ax.view_init(
            azim=(np.pi * 1.7 + 0.8 * np.sin(2.0 * np.pi * i * self.steps_per_frame / len(self.frames) / 10))
            * 180.0
            / np.pi
        )
        ax.set_axis_off()
        ax.grid(visible=False)
        loom.save_frame(fig, i - 1)


with Loom('lorenz.mp4', fps=60, parallel=True) as loom:
    attractor = LorenzPlotter()
    attractor.initialize(10000)
    Parallel(n_jobs=-1)(delayed(attractor.get_frame)(i, loom) for i in attractor.frames)

lorenz.mp4

Looming in parallel

By passing parallel=True when creating a Loom, you can save frames using loom.save_frame(fig, frame_number) which allows you to plot and save all your frames in parallel. One easy way to leverage this is by using joblib to parallelize the for loop. For example, here's how you can parallelize the simple sine wave example:

import numpy as np
import matplotlib.pyplot as plt
from matplotloom import Loom
from joblib import Parallel, delayed

def plot_frame(phase, frame_number, loom):
    fig, ax = plt.subplots()

    x = np.linspace(0, 2*np.pi, 200)
    y = np.sin(x + phase)

    ax.plot(x, y)
    ax.set_xlim(0, 2*np.pi)

    loom.save_frame(fig, frame_number)

with Loom("parallel_sine_wave.gif", fps=30, parallel=True) as loom:
    phases = np.linspace(0, 2*np.pi, 100)

    Parallel(n_jobs=-1)(
        delayed(plot_frame)(phase, i, loom)
        for i, phase in enumerate(phases)
    )