Animation with matplotlib - "4 bugs in square"

In this famous problem 4 bugs start at the corners of a square and move chasing each other. Bug nr 1 moves towards nr 2, nr 2 to 3, nr 3 to 4 and the last one towards the 1st.

Solving the motion equation and find the analytical expression for their trajectories is not trivial. You can easily Google for the mathemathical solutions (for instance here).

Our goal is just to visualize each bug's motion.

Do the animation

Let's use the FuncAnimation() class of the matplotlib.animate package.

Firstly, setup a general animation able to handle the 4 bugs (points on a 2D plane)

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

from matplotlib import animation
from IPython.display import HTML

fig, ax = plt.subplots(figsize=(5,5))
ax.set_xlim((0, 1))
ax.set_ylim((0, 1))

# point marker/color
fmts=['or','ob','og','ok']

# prepare a list to contain the points and fill it
points=[]

for i in range(4):
    p, = ax.plot([], [], fmts[i])
    points.append(p)

Create the init and animate functions to feed FuncAnimation()

Remember the they have to return a iterable of Artists (lines, points... whatever is in the plot is an Artist)

## the "velocity" (equal for the 4 points)
vx, vy = 0.03, 0.03

# starting points
start_positions = [ [0,0], [1,0], [1,1], [0,1] ]

def init():
    for i in range(4):
        points[i].set_data(start_positions[i])
    return points

def animate(t):
    # Let the points move towards the center 
    # just to see if the animations is ok up to this point)
    for i in range(4):
        x, y = points[i].get_data()
        newx = x - vx*(x-0.5)
        newy = y - vy*(y-0.5)
        points[i].set_data(newx,newy)
    return points

anim = animation.FuncAnimation(fig, animate, init_func=init, frames=100, interval=50,blit=True)

Display the animation

HTML(anim.to_jshtml())

Once Loop Reflect

Hint: going towards one point (the center in the previous case) is simply a shift along a line connecting that point and the target. This point can be computed as the weighted average between the 1st bug position and the target's one. The weigth represent the "speed" with which the bug will reach the target.

Let's see with a simple skectch

### this plot is just to illustrate the "weighted average trick"

plt.subplots()
plt.axis([0,1,0,1])
plt.axis(False)
x, y = [0.1, 0.8], [0.2,0.9]
plt.plot(x,y,'--k',x,y,'ro',markersize=10)
a = plt.annotate('P1 (x1, y1)', xy=(x[0],y[0]), xytext=(x[0],y[0]-0.1), fontsize=14)
plt.annotate('P2 (x2, y2)', xy=(x[1],y[1]), xytext=(x[1],y[1]-0.1), fontsize=14)

stepw = 0.3  # if 0.3, P2 willarrive at ~1/3 of the path towards P1

p2newpos = x[0]*stepw + x[1]*(1-stepw) , y[0]*stepw + y[1]*(1-stepw)
plt.plot(p2newpos[0],p2newpos[1], 'ob', markersize=10)
plt.annotate('P2\' ( x1*(1-{0})+x2*{0}, y1*(1-{0})+y2*{0} )'.format(stepw), 
             xy=p2newpos, xytext=(p2newpos[0], p2newpos[1]-0.1),
             fontsize=14)
plt.annotate('',xy=(p2newpos[0]+0.03, p2newpos[1]+0.03), xytext=(x[1]-0.01,y[1]-0.01), 
             arrowprops=dict(facecolor='blue',edgecolor='blue', shrink=0.02) )
plt.show()

Now set up the correct movement rule exploiting the above illustrated method.

In order to do it, let's define a function that take 2 points as arguments and returns new x and y (computed as the weighted average of the 2 points) to be used in other parts of the program. Make it more flexible taking "velocity" (that will be used as weigths normalized to 1) as arguments as well

def move_to(p1,p2, vx, vy):
    pos1 = p1.get_data()
    pos2 = p2.get_data()
    newx = pos1[0]*(1-vx) + pos2[0]*vx
    newy = pos1[1]*(1-vy) + pos2[1]*vy
    return newx, newy

Now we can us this function to move the points according to the problem rules

def animate(t):
    for i in range(4):
        newx , newy = move_to(points[i%4], points[(i+1)%4], vx, vy)
        points[i].set_data(newx,newy)
    return points

anim = animation.FuncAnimation(fig, animate, init_func=init, frames=100, interval=50,blit=True)

HTML(anim.to_jshtml())

Once Loop Reflect

Let's track each point position and draw it.

We have to add another Artist to the animation (the line containing each point coordinates)

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

from matplotlib import animation
from IPython.display import HTML

fig, ax = plt.subplots(figsize=(5,5))
ax.set_xlim((0, 1))
ax.set_ylim((0, 1))

# 4 points marker/color + 4 lines  
fmts=['or','ob','og','ok', '-r', '-b', '-g', '-k']

# prepare a list to contain the points+lines and fill it
points=[]

for i in range(8):
    p, = ax.plot([], [], fmts[i], linewidth=1)
    points.append(p)

## the "velocity" (the same for all the 4 points)
vx, vy = 0.03, 0.03

# starting points and track list
start_positions = [ [0,0], [1,0], [1,1], [0,1] ]
# list of list of lists.. 
# To easily keep track of the x and the y of each points
track = [[[],[]], [[],[]], [[],[]], [[],[]] ]

# same function as before
def move_to(p1,p2, vx, vy):
    pos1 = p1.get_data()
    pos2 = p2.get_data()
    newx = pos1[0]*(1-vx) + pos2[0]*vx
    newy = pos1[1]*(1-vy) + pos2[1]*vy
    return newx, newy

def init():
    for i in range(4):
        points[i].set_data(start_positions[i])  # points
        # set the line starting point as equal to the point
        points[i+4].set_data([start_positions[i][0]],[start_positions[i][1]])
        # add each point (this is the first) to the track
        track[i][0].append(start_positions[i][0]) # x
        track[i][1].append(start_positions[i][1]) # y
    return points

def animate_points_and_tracks(t):
    for i in range(4):
        newx , newy = move_to(points[i%4], points[(i+1)%4], vx, vy)
        points[i].set_data(newx,newy)
        # add x and y of each point to the right track. 
        track[i][0].append(newx)
        track[i][1].append(newy)
        # then store the track as new "data" for that line
        points[i+4].set_data(track[i][0], track[i][1])
    return points

anim = animation.FuncAnimation(fig, animate_points_and_tracks, 
                               init_func=init, 
                               frames=50, 
                               interval=50,
                               blit=True)

HTML(anim.to_jshtml())

Once Loop Reflect

What happens if we set different velocities?

We just need to make speed a list of 4 elements.

Note that the vx[] and vy[] should only be replaced in the code where the move_to() function is called, not where it is defined! That means vx and vy has a local scope there.

## the "velocity" (the same for all the 4 points)
vx, vy = [0.03, 0.02, 0.01, 0.015], [0.01, 0.025, 0.015, 0.04]

# starting points and track list
start_positions = [ [0,0], [1,0], [1,1], [0,1] ]
# list of list of lists.. 
# To easily keep track of the x and the y of each points
track = [[[],[]] for i in range(4) ]

# same function as before
def move_to(p1,p2, vx, vy):
    pos1 = p1.get_data()
    pos2 = p2.get_data()
    newx = pos1[0]*(1-vx) + pos2[0]*vx
    newy = pos1[1]*(1-vy) + pos2[1]*vy
    return newx, newy

def init():
    for i in range(4):
        points[i].set_data(start_positions[i])  # points
        # set the line starting point as equal to the point
        points[i+4].set_data([start_positions[i][0]],[start_positions[i][1]])
        # add each point (this is the first) to the track
        track[i][0].append(start_positions[i][0]) # x
        track[i][1].append(start_positions[i][1]) # y
    return points

def animate_points_and_tracks(t):
    for i in range(4):
        newx , newy = move_to(points[i%4], points[(i+1)%4], vx[i], vy[i])
        points[i].set_data(newx,newy)
        # add x and y of each point to the right track. 
        track[i][0].append(newx)
        track[i][1].append(newy)
        # then store the track as new "data" for that line
        points[i+4].set_data(track[i][0], track[i][1])
    return points

anim = animation.FuncAnimation(fig, animate_points_and_tracks, 
                               init_func=init, 
                               frames=200, 
                               interval=30,
                               blit=True)

HTML(anim.to_jshtml())

Once Loop Reflect

What if we want more points?

Let's make the previous code more general: put N points along a circumference of radius 1

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

from matplotlib import animation
from IPython.display import HTML

fig, ax = plt.subplots(figsize=(5,5))
ax.set_xlim((-1, 1))
ax.set_ylim((-1, 1))
ax.axis(False)
N = 15

# prepare a list to contain the points+lines and fill it
points=[]

for i in range(N):
    p, = ax.plot([], [], marker='o')
    points.append(p)
for i in range(N):
    l, = ax.plot([], [], linewidth=1)
    points.append(l)

## the "velocity" (the same for all the 4 points)
vx, vy = 0.1, 0.1

# starting points and track list

x = np.cos(2*np.pi/N ), np.cos(2*np.pi/N )

start_positions = [ [np.cos(2*np.pi/N * i),np.sin(2*np.pi/N * i)] for i in range(N) ]

# list of list of lists.. 
# To easily keep track of the x and the y of each points
track = [[[],[]] for i in range(N) ]

# same function as before
def move_to(p1,p2, vx, vy):
    pos1 = p1.get_data()
    pos2 = p2.get_data()
    newx = pos1[0]*(1-vx) + pos2[0]*vx
    newy = pos1[1]*(1-vy) + pos2[1]*vy
    return newx, newy

def init():
    for i in range(N):
        points[i].set_data(start_positions[i])  # points
        # set the line starting point as equal to the point
        points[i+N].set_data([start_positions[i][0]],[start_positions[i][1]])
        # add each point (this is the first) to the track
        track[i][0].append(start_positions[i][0]) # x
        track[i][1].append(start_positions[i][1]) # y
    return points

def animate_points_and_tracks(t):
    for i in range(N):
        newx , newy = move_to(points[i%N], points[(i+1)%N], vx, vy)
        points[i].set_data(newx,newy)
        # add x and y of each point to the right track. 
        track[i][0].append(newx)
        track[i][1].append(newy)
        # then store the track as new "data" for that line
        points[i+N].set_data(track[i][0], track[i][1])
    return points

anim = animation.FuncAnimation(fig, animate_points_and_tracks, 
                               init_func=init, 
                               frames=150, 
                               interval=50,
                               blit=True)

HTML(anim.to_jshtml())

Once Loop Reflect

Variants

Change the move_to() function to make it funnier:

• randomly swap a bug with its target
• assign a random target from time to time

start_positions = [ [np.cos(2*np.pi/N * i),np.sin(2*np.pi/N * i)] for i in range(N) ]

# list of list of lists.. 
# To easily keep track of the x and the y of each points
track = [[[],[]] for i in range(N) ]

def animate_swap(t):
    for i in range(N):
        # every 20 time steps swap
        sign = 1 if t%20>9 else -1
        newx , newy = move_to(points[i%N], points[(i+sign*1)%N], vx, vy)
        points[i].set_data(newx,newy)
        # add x and y of each point to the right track. 
        track[i][0].append(newx)
        track[i][1].append(newy)
        # then store the track as new "data" for that line
        points[i+N].set_data(track[i][0], track[i][1])
    return points

anim = animation.FuncAnimation(fig, animate_swap, 
                               init_func=init, 
                               frames=200, 
                               interval=50,
                               blit=True)

HTML(anim.to_jshtml())

Once Loop Reflect

import random

## redefine "velocity" (the same for all the 4 points)
vx, vy = 0.01, 0.01

start_positions = [ [np.cos(2*np.pi/N * i),np.sin(2*np.pi/N * i)] for i in range(N) ]

# list of list of lists.. 
# To easily keep track of the x and the y of each points
track = [[[],[]] for i in range(N) ]

targets = [(i+1)%N for i in range(N)]

def random_follow(t):
    for i in range(N):
        if t%20==0:
            targets[i] = random.choice( [j for j in range(N) if j!=i] )
        newx , newy = move_to(points[i%N], points[targets[i]], vx, vy)
        points[i].set_data(newx,newy)
        # add x and y of each point to the right track. 
        track[i][0].append(newx)
        track[i][1].append(newy)
        # then store the track as new "data" for that line
        points[i+N].set_data(track[i][0], track[i][1])
    return points

anim = animation.FuncAnimation(fig, random_follow, 
                               init_func=init, 
                               frames=200, 
                               interval=50,
                               blit=True)

HTML(anim.to_jshtml())

Once Loop Reflect