Search code examples
pythonmatplotlibscatter-plotmatplotlib-animation

How to create a video that plots the scatter plots one by one


I'm plotting some data as scatter plots which is overlaid on an image. I would like to make an animation of this by plotting one scatter point at a time. This is the code I have right now using and editing the answer from here:

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

x = random.sample(range(0, 287), 20)
y = random.sample(range(0, 380), 20)
size =  [20 for x in range(20)]
colors = ["r" for x in range(20)]


cm = plt.get_cmap('jet')

fig = plt.figure(figsize=(18,9))

graph = plt.scatter([], [],marker='+')
url = 'https://raw.githubusercontent.com/kornelski/pngquant/master/test/img/test.png'
im = plt.imread(url)
def animate(i):
        
    implot = plt.imshow(im)
    graph.set_offsets(np.vstack((x[:i+1], y[:i+1])).T)
    graph.set_sizes(size[:i])
    graph.set_facecolors(colors[:i+1])
    return graph

ani = FuncAnimation(fig, animate, repeat=False, interval=0.1)
plt.show()

There are two things I would like help with.

  1. I would like the color of my scatterplot to change based on a third variable, i.e use a cmap. However, the set_facecolors does not accept such an argument.
  2. When I try to save my animation using ani.save('files/animation.gif',writer='imagemagick', fps=60) my jupyter notebook crashes.

Solution

  • The background image of the graph is drawn by adding ax. The color map is also created according to the number of data, 20, and a list is created so that each color can be displayed. Since the coordinates of the image and the coordinate basis of the graph are different, the y-axis is set in the opposite direction.

    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    import random
    from PIL import Image
    import urllib.request
    
    random.seed(20210702)
    N = 20
    x = random.sample(range(0, 380), N)
    y = random.sample(range(0, 287), N)
    size =  [20 for x in range(N)]
    colors = []
    
    cm = plt.get_cmap('jet', N)
    
    fig,ax = plt.subplots(figsize=(9, 4.5))
    plt.xlim(0, 380)
    plt.ylim(287, 0)
    graph = ax.scatter([], [], marker='+')# 
    
    url = 'https://raw.githubusercontent.com/kornelski/pngquant/master/test/img/test.png'
    im = Image.open(urllib.request.urlopen(url))
    print(im.size)
    
    def animate(i):
        ax.imshow(im)
        graph.set_offsets(np.vstack((x[:i+1], y[:i+1])).T)
        graph.set_sizes(size[:i+1])
        colors.append(cm(i))
        graph.set_facecolors(colors) 
        return graph
    
    ani = FuncAnimation(fig, animate, frames=20, repeat=False, interval=200)
    plt.show()
    

    enter image description here