life game in python

life.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import sys

SIZE = int(sys.argv[1])
E = np.eye(SIZE, k=1) + np.eye(SIZE) + np.eye(SIZE, k=-1)
world = np.random.randint(2, size=(SIZE, SIZE))


def update(world):
    neighbor = E @ world @ E - world
    return (neighbor == 3) + world * (neighbor == 2)


def animate(i):
    global world
    plt.cla()
    world = update(world)
    im = plt.imshow(world, interpolation='none')

fig = plt.figure()
ani = animation.FuncAnimation(fig, animate, interval=100)
plt.show()