How does one animate a 2d grid using matplotlib

user_11147253 0
Viewed 133

I'm trying to code langton's ant in python and while the code does work when i test it as an ASCII "simulation" i'm having problems making it work in matplotlib as a 2d animation. I've never really messes with matplotlib but i was trying to learn by doing this.

import datetime,pprint,keyboard,time,os
from matplotlib import pyplot as plt
from matplotlib import animation

antLocation = (15,15) # (y,x)
antDirection = 270
grid = [[0 for i in range(30)].copy() for _ in range(30)]

def update(data):
    global grid
    global antLocation
    global antDirection
    if(grid[antLocation[0]][antLocation[1]] == 0):
        antDirection +=90
        grid[antLocation[0]][antLocation[1]] = 1
    else:
        antDirection -=90
        grid[antLocation[0]][antLocation[1]] = 0
    if (antDirection == 0 or antDirection == 360): #UP
        antLocation = (antLocation[0]-1, antLocation[1])
        antDirection = 0
    elif (antDirection == 90): #RIGHT
        antLocation = (antLocation[0], antLocation[1] + 1)
    elif (antDirection == 180): # DOWN
        antLocation = (antLocation[0]+1, antLocation[1])
    else: #(antDirection == 270 or antDirection == -90): #LEFT
        antLocation = (antLocation[0], antLocation[1] - 1)
        antDirection = 270
    mat.set_data(grid)
    return [mat]
    
fig,ax = plt.subplots()
mat = ax.matshow(grid)
ani = animation.FuncAnimation(fig, update)
plt.show()
matplotlib python
1 Answers

You need a function that creates the individual frames (for each time point t). Then pass this function to animation.FuncAnimation(). Since you're working with global variables, update() doesn't need to return anything.

from matplotlib import pyplot as plt
from matplotlib import animation

antLocation = (15,15) # (y,x)
antDirection = 270
grid = [[0 for i in range(30)].copy() for _ in range(30)]

def update(data):
    global grid
    global antLocation
    global antDirection
    if(grid[antLocation[0]][antLocation[1]] == 0):
        antDirection +=90
        grid[antLocation[0]][antLocation[1]] = 1
    else:
        antDirection -=90
        grid[antLocation[0]][antLocation[1]] = 0
    if (antDirection == 0 or antDirection == 360): #UP
        antLocation = (antLocation[0]-1, antLocation[1])
        antDirection = 0
    elif (antDirection == 90): #RIGHT
        antLocation = (antLocation[0], antLocation[1] + 1)
    elif (antDirection == 180): # DOWN
        antLocation = (antLocation[0]+1, antLocation[1])
    else: #(antDirection == 270 or antDirection == -90): #LEFT
        antLocation = (antLocation[0], antLocation[1] - 1)
        antDirection = 270
    
def frame(t):
    """
    Update the grid and return a plot of the earlier grid.
    t (time) is not used, but necessary to be able to pass
    this function to matplotlib.animation.FuncAnimation.
    """
    plt.cla()
    image = plt.imshow(grid, cmap='Greys')
    update(grid)
    return image

fig, ax = plt.subplots()
ani = animation.FuncAnimation(fig, frame)
plt.show()
user_130141720
Related