Make Your Data Move: Creating Animations in Python for Science and Machine Learning

As a data scientist and professor, I often need to explain the inner workings of learning algorithms and mathematical concepts, whether in technical presentations, classroom lectures, or written articles. In many cases, static plots can show the final outcome, but they often fall short when it comes to illustrating the underlying process.

In this context, animations can make a significant difference. By presenting a sequence of frames, each showing a plot that represents a step in the process, you can better capture your audience’s attention and more effectively explain complex concepts and workflows.

This tutorial will show you how to use Python and Matplotlib to bring scientific ideas to life through animation. Whether you’re a data scientist visualizing a Machine Learning algorithm, a physics teacher demonstrating harmonic motion, or a technical writer aiming to convey math intuitively, this guide is for you.

We’ll explore the following topics:

1. Basic animation setup with Matplotlib
2. Math example
3. Physics Example
4. Animating machine learning algorithms
5. Exporting animations for web and presentations

1. Basic animation setup with Matplotlib

Let’s introduce the FuncAnimation class from Matplotlib’s Animation package by animating the sine function. The following steps can be replicated virtually in every case.

• Import required libraries

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

FuncAnimation from matplotlib.animation is the class that allows you to create animations by repeatedly calling an update function.

• Defining the plotted data (sine function)

x = np.linspace(0, 2 * np.pi, 1000)
y = np.sin(x)

y = np.sin(x) computes the sine of each x-value. This is the initial sine wave that will be plotted.

• Creating the initial plot

Python">fig, ax = plt.subplots()
line, = ax.plot(x, y)
ax.set_ylim(-1.5, 1.5)

line, = ax.plot(x, y) plots the initial sine wave and stores the line object in line.

Note: The comma after line is important: it unpacks the single-element tuple returned by plot.

• Defining the Update Function

def update(frame):
    line.set_ydata(np.sin(x + frame / 10))
    return line,

np.sin(x + frame / 10) shifts the sine wave horizontally, creating the effect of a moving wave.

• Creating and displaying the animation

ani = FuncAnimation(fig, update, frames=100, interval=50, blit=True)
plt.show()

That ties everything together to create the animation:

fig: the figure to animate.
update: the function to call for each frame.
frames=100: the number of frames in the animation.
interval=50: delay between frames in milliseconds (50 ms = 20 frames per second).
blit=True: optimizes performance by only redrawing parts of the plot that change.

• Result

2. Animating Physics: Oblique Launch

We’ll show how to animate a classic example in physics classes: the Oblique Launch. We’ll follow similar steps to the basic example shown before.

• Defining motion parameters and the time vector

g = 9.81  # gravity (m/s^2)
v0 = 20   # initial velocity (m/s)
theta = np.radians(45)  # launch angle in radians

# total time the projectile will be in the air
t_flight = 2 * v0 * np.sin(theta) / g
# time vector with 100 equally spaced values between 0 and t_flight
t = np.linspace(0, t_flight, 100)

• Computing the trajectory

x = v0 * np.cos(theta) * t # horizontal position at time t
y = v0 * np.sin(theta) * t - 0.5 * g * t**2 # vertical position at time t

• Setting up the plot

fig, ax = plt.subplots()
ax.set_xlim(0, max(x)*1.1)
ax.set_ylim(0, max(y)*1.1)
ax.set_title("Oblique Launch")
ax.set_xlabel("Distance")
ax.set_ylabel("Height")
line, = ax.plot([], [], lw=2)
point, = ax.plot([], [], 'ro')  # red dot for projectile

• Initialization Function

In this example, we’ll use an initialization function to set everything to an empty state. It returns the plot elements to be updated during the animation.

def init():
    line.set_data([], [])
    point.set_data([], [])
    return line, point

• Update function and animation

def update(frame):
    line.set_data(x[:frame], y[:frame])     # trajectory up to current frame
    point.set_data(x[frame], y[frame])      # current projectile position
    return line, point

The update function is called at each frame. The frame parameter indexes into the time array, so x[frame] and y[frame] give current coordinates.

ani = FuncAnimation(fig, update, frames=len(t), init_func=init, blit=True, interval=30)

frames=len(t): total number of animation steps.

interval=30: time (in milliseconds) between frames (~33 fps).

blit=True: improves performance by only redrawing parts of the frame that changed.

• Result

3. Animating Math: Fourier Series

In the following example, we’ll show how to build a square wave from sine functions using the Fourier Transform.

• Creating x values and a plot figure

x = np.linspace(-np.pi, np.pi, 1000)
y = np.zeros_like(x)
fig, ax = plt.subplots()
line, = ax.plot(x, y, lw=2)

# Setting text labels and axis limits
ax.set_title("Fourier Series Approximation of a Square Wave")
ax.set_ylim(-1.5, 1.5)
ax.set_xlim(-np.pi, np.pi)
text = ax.text(-np.pi, 1.3, '', fontsize=12)

x: An array of 1000 evenly spaced points from −π to π, which is the typical domain for periodic Fourier series.

y: Initializes a y-array of zeros, same shape as x.

• Defining the Fourier Series function

The formula used for the Fourier Series is given by:

def fourier_series(n_terms):
    result = np.zeros_like(x)
    for n in range(1, n_terms * 2, 2):  # Only odd terms: 1, 3, 5, ...
        result += (4 / (np.pi * n)) * np.sin(n * x)
    return result

• Setting the Update Function

def update(frame):
    y = fourier_series(frame + 1)
    line.set_ydata(y)
    text.set_text(f'{2*frame+1} terms')
    return line, text

This function updates the plot at each frame of the animation:

• It computes a new approximation with frame + 1 terms.
• Updates the y-values of the line.
• Updates the label to show how many terms are used (e.g., “3 terms”, “5 terms”, etc.).
• Returns the updated plot elements to be redrawn.

• Creating the visualization

ani = FuncAnimation(fig, update, frames=20, interval=200, blit=True)
plt.show()

4. Machine Learning in Action: Gradient Descent

Now, we’ll show how the classical machine learning algorithm finds a minimum on a three-dimensional parabolic function.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Define the function and its gradient
def f(x, y):
    return x**2 + y**2

def grad_f(x, y):
    return 2*x, 2*y

# Initialize parameters
lr = 0.1
steps = 50
x, y = 4.0, 4.0  # start point
history = [(x, y)]

# Perform gradient descent
for _ in range(steps):
    dx, dy = grad_f(x, y)
    x -= lr * dx
    y -= lr * dy
    history.append((x, y))

# Extract coordinates
xs, ys = zip(*history)
zs = [f(xi, yi) for xi, yi in history]

# Prepare 3D plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X, Y = np.meshgrid(np.linspace(-5, 5, 100), np.linspace(-5, 5, 100))
Z = f(X, Y)
ax.plot_surface(X, Y, Z, alpha=0.6, cmap='viridis')
point, = ax.plot([], [], [], 'ro', markersize=6)
line, = ax.plot([], [], [], 'r-', lw=1)

# Animation functions
def init():
    point.set_data([], [])
    point.set_3d_properties([])
    line.set_data([], [])
    line.set_3d_properties([])
    return point, line

def update(i):
    point.set_data(xs[i], ys[i])
    point.set_3d_properties(zs[i])
    line.set_data(xs[:i+1], ys[:i+1])
    line.set_3d_properties(zs[:i+1])
    return point, line

ani = FuncAnimation(fig, update, frames=len(xs), init_func=init, blit=True, interval=200)

• Result

5. Exporting animations for web and presentations

Finally, to export the animated plots to a file, you can use the animation.save() function.

• Example

# Export as GIF (optional)
ani.save("launch.gif", writer='pillow', fps=30)

In the example above, the function takes the FuncAnimation object, renders it frame by frame using the Pillow library, and exports the result as a .gif file called launch.gif at 30 frames per second.

Conclusion

In this article, we saw how the animation class from matplotlib can be handy for demonstrating the inner workings of algorithms, mathematical, and physical processes. The examples explored in this article can be expanded to create impactful visuals for blog posts, lectures, and reports.

In order to make the present article even more profitable, I suggest using the examples shown to create your own animations and simulate processes related to your field.

Check out my GitHub repository, which has complete code examples and animations available here.

References

[1] GeeksforGeeks. Using Matplotlib for animations. https://www.geeksforgeeks.org/using-matplotlib-for-animations/

[2] TutorialsPoint. Matplotlib – Animations. https://www.tutorialspoint.com/matplotlib/matplotlib_animations.htm

[3] The Matplotlib Development Team. Animations using Matplotlib. https://matplotlib.org/stable/users/explain/animations/animations.html

The post Make Your Data Move: Creating Animations in Python for Science and Machine Learning appeared first on Towards Data Science.