Animating Langdon's Ant in Python with Matplotlib

Question

doing some problems on Project Euler, I came across Langdon's Ant and thought it would be a nice idea to try to code an animation of it in Python. As a basis, I went and used the matplotlib function animation, a very nice example can be seen on this post here

I have managed to get a version working, the code is shown below. The idea is to simulate the black and white grid the ant moves on with a large matrix and plot that matrix with 'imshow'.

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

# Initialize the Simulation

dim=50
a=np.matrix(np.zeros((dim,dim)))
pos=np.matrix([[dim//2],[dim//2]]) # current position of ant
direction=np.matrix([[1],[0]])  # direction ant is currently moving

#Rotation Matrices 
clock=np.matrix([[0,1],[-1,0]])
counter=np.matrix([[0,-1],[1,0]])

def takestep(a,pos,direction):
    pos[:]=pos+direction       
    if a[pos[0,0],pos[1,0]]==0: #landed on white
        a[pos[0,0],pos[1,0]]=1
        direction[:]=clock*direction
    else:        
        a[pos[0,0],pos[1,0]]=0
        direction[:]=counter*direction

#Plotting    
fig = plt.figure()        
im=plt.imshow(a,interpolation='none')  

def init():
    im=plt.imshow(a,interpolation='none')
    return [im]

def animate(i):
    takestep(a,pos,direction)
    im=plt.imshow(a,interpolation='none')
    #im.set_data(np.array(a))        
    return [im]


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

This code is functional, but I'm unclear about somethings:

• The animation runs continuously (until the ant is out of bounds) and doesn't reset after 200 frames

• Even if the interval is set to 0, the animation runs quite slowly, updateing about twice a second

• I think I can speed it up by using the im.set_data() function (commented out in the code), in my implementation this doesn't work though (unchanging visualization)

It would be great if someone could give me some pointers on how to improve this animation, thanks in advance!

Best regards, Raphael

Answer 1

Here is how you could use im.set_data to improve frames/second.

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

# Initialize the Simulation

dim = 50
a = np.matrix(np.zeros((dim, dim)))
pos = np.matrix([[dim // 2], [dim // 2]])  # current position of ant
direction = np.matrix([[1], [0]])  # direction ant is currently moving

# Rotation Matrices
clock = np.matrix([[0, 1], [-1, 0]])
counter = np.matrix([[0, -1], [1, 0]])


def takestep(a, pos, direction):
    pos[:] = pos + direction
    if a[pos[0, 0], pos[1, 0]] == 0:  # landed on white
        a[pos[0, 0], pos[1, 0]] = 1
        direction[:] = clock * direction
    else:
        a[pos[0, 0], pos[1, 0]] = 0
        direction[:] = counter * direction

fig = plt.figure()
im = plt.imshow(a, interpolation='none', vmin=0, vmax=1)

def animate(i):
    takestep(a, pos, direction)
    im.set_data(a)
    return [im]


anim = animation.FuncAnimation(fig, animate,
                               frames=200, interval=0, blit=True,
                               repeat=False)
plt.show()

---

• Use repeat=False to stop the animation after 200 frames.
• Set vmin=0, vmax=1 in the initial call to imshow. The initial call to plt.imshow sets vmin and vmax, which governs the correspondence between values and colors. That setting is used for all subsequent renderings of the imshow AxesImage. Since the initial plt.imshow uses all-zero data, vmin and vmax both get set to 0 (the min and max of the initial data). That makes all values show up as the same color. So to stop that from happening, supply the vmin and vmax yourself.