Scatterplot animation with matplotlib in Jupyter

scatter_animation.py

      
import matplotlib.pyplot as plt
import matplotlib
import matplotlib.animation as animation
from matplotlib.animation import PillowWriter
import numpy as np
from IPython.display import HTML

# Example data for two states
state1_x = np.random.rand(10)  # x-coordinates for state 1
state1_y = np.random.rand(10)  # y-coordinates for state 1
state2_x = np.random.rand(10)  # x-coordinates for state 2
state2_y = np.random.rand(10)  # y-coordinates for state 2

labels = [f"Point {i}" for i in range(10)]  # Labels for the points

# Create the figure and axes
fig, ax = plt.subplots()

# Initialize the scatter plot
scatter = ax.scatter(state1_x, state1_y)

# Set axis limits (important for consistent animation)
ax.set_xlim(0, 1)  # Adjust as needed for your data
ax.set_ylim(0, 1)

# Function to update the scatter plot for each frame
def animate(i):
    # Interpolate between state 1 and state 2
    interpolated_x = state1_x * (1 - i / 100) + state2_x * (i / 100)  # 100 frames
    interpolated_y = state1_y * (1 - i / 100) + state2_y * (i / 100)

    # Update the scatter plot data
    scatter.set_offsets(np.c_[interpolated_x, interpolated_y])

    for annotation in ax.get_children(): # Iterate through all children of the axes
        if isinstance(annotation, matplotlib.text.Annotation):  # Check if it's an annotation
            annotation.remove()

    # Add new annotations
    for j, label in enumerate(labels):
        ax.annotate(label, (interpolated_x[j], interpolated_y[j]), textcoords="offset points", xytext=(0,10), ha='center')

    return scatter,

# Create the animation
ani = animation.FuncAnimation(fig, animate, frames=100, interval=20, blit=False)

# Save the animation as a GIF using PillowWriter (optional)
writer = PillowWriter(fps=100)  # Adjust fps as needed

# Display in Jupyter Notebook
HTML(ani.to_jshtml())  # Or HTML(ani.to_html5_video())