Brute Force Voronoi Animation with Moving Sites in Matplotlib

VoronoiBrutAnim.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import KDTree

# Generate some number of sites in 
# the square [0, 1] x [0, 1]
np.random.seed(0)
n_sites = 20
sites = np.random.rand(n_sites, 2)
# Create random velocities
vel = 2*np.random.rand(n_sites, 2)-1
vel *= 0.001

# Create a grid of points in [0, 1] x [0, 1]
res = 1000 # The resolution of the grid
pix = np.linspace(0, 1, res)
x, y = np.meshgrid(pix, pix)
X = np.array([x.flatten(), y.flatten()]).T

plt.figure(figsize=(5, 5))
n_frames = 3000
for i in range(n_frames):
    sites += vel
    vel[sites > 1] *= -1
    vel[sites < 0] *= -1
    sites[sites > 1] = 1
    sites[sites < 0] = 0
    # Find the nearest site to each grid point using
    # a KD Tree
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.html
    
    tree = KDTree(sites)
    neighbs = tree.query(X, 1)[1].flatten()
    neighbs = np.reshape(neighbs, x.shape)
    # Plot the results
    plt.clf()
    plt.imshow(neighbs, extent=(0, 1, 1, 0), cmap='tab20b')
    plt.scatter(sites[:, 0], sites[:, 1], c='w')
    plt.gca().invert_yaxis()
    plt.axis("equal")
    plt.axis("off")
    plt.savefig("{}.png".format(i), bbox_inches='tight')