Python Pygame Convay’s Game Of Life

Pygame implementation of Convay’s Game Of Life.

import pygame
import sys
import numpy as np


# ? 2019 TheFlyingKeyboard and released under MIT License
# theflyingkeyboard.net

def life_step(x, y, world):
    neighbours = np.sum(world[x - 1: x + 2, y - 1: y + 2]) - world[x, y]

    if world[x, y] == 1 and not 2 <= neighbours <= 3:
        return 0
    elif neighbours == 3:
        return 1
    return world[x, y]


screenSize = 1200
tileSize = 30

backgroundColor = (255, 255, 255)
cellColor = (0, 0, 0)

pygame.init()
screen = pygame.display.set_mode((screenSize, screenSize))
pygame.display.set_caption('Convay\'s Game Of Life')

clock = pygame.time.Clock()

world = np.zeros((screenSize // tileSize + 1, screenSize // tileSize + 1))

world = np.random.choice(a=[0, 1], size=(screenSize // tileSize + 1, screenSize // tileSize + 1))

fps = 10
epoch = 0

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                pygame.quit()
                sys.exit()
            if event.key == pygame.K_r:  # Restart world
                world = np.random.choice(a=[0, 1], size=(screenSize // tileSize + 1, screenSize // tileSize + 1))
                epoch = 0
            if event.key == pygame.K_q:
                fps += 1
            if event.key == pygame.K_a:
                fps -= 1
                if fps == 0:
                    fps = 1

    pygame.display.set_caption("Convay\'s Game Of Life " + str(fps) + "fps " + str(epoch) + " epoch")
    screen.fill(backgroundColor)

    newWorld = np.copy(world)

    for (x, y), value in np.ndenumerate(world):
        newWorld[x, y] = life_step(x, y, world)

        if newWorld[x, y] == 1:
            pygame.draw.rect(screen, cellColor, (tileSize * (x - 1), tileSize * (y - 1), tileSize, tileSize), 0)

    world = newWorld
    epoch += 1

    pygame.display.update()

    msElapsed = clock.tick(fps)