Python implementation of Conway's Game of life with Pygame

So I've started learning programming a few months ago and this is one of my first projects. The program works, but I feel like it's a bit slow compared to the other implementations of Conway's Game of life. Is there any ways I could speed up my code ?

import numpy
import copy
import pygame
from pygame.locals import *
import random

pygame.init()

height = 400
width = 400
block_size = 4
window = pygame.display.set_mode((height,width),DOUBLEBUF)
window.set_alpha(None)
alive = (255,255,255)

def check(y_axis,x_axis,grid):
tmp_grid = copy.copy(grid)
for y in range(1,y_axis-1):
for x in range(1,x_axis-1):
a = 0
for i in range(-1,2):
for j in range(-1,2):
if grid[y+i][x+j] == 1:
a += 1
if grid[y][x] == 1:
a -= 1
if a < 2 or a > 3:
tmp_grid[y][x] = 0
elif grid[y][x] == 0:
if a == 3:
tmp_grid[y][x] = 1
grid = tmp_grid
return grid

def display(y_axis,x_axis,grid):
rects = []
for y in range(1,y_axis-1):
for x in range(1,x_axis-1):
rect = pygame.Rect(x*block_size, y*block_size, block_size, block_size)
if grid[x][y] == 1:
rects.append(pygame.draw.rect(window,alive,rect))
elif grid[x][y] == 0:
pygame.display.update(rects)

def generate(y_axis,x_axis,grid):
for y in range(1,y_axis-1):
for x in range(1,x_axis-1):
grid[y][x] = 1
else:
grid[y][x] = 0

def main():
y_axis = height//block_size
x_axis = width//block_size
grid = numpy.zeros((y_axis,x_axis))
generate(y_axis,x_axis,grid)
while True:
display(y_axis,x_axis,grid)
grid = check(y_axis,x_axis,grid)

if __name__ == '__main__':
main()

1. In your generate function, numpy.random.randint can do everything you need. I'd suggest renaming the function to generate_grid and just returning the generated grid.

2. In your display and check functions, the grid parameter knows its own size - you don't need to pass that in.

3. In display, your alive and dead colors could be in a tuple. Then you could index the tuple by  or  (the value of the cell) to determine the color:

rect = pygame.Rect(x*block_size, y*block_size, block_size, block_size)

rects.append(pygame.draw.rect(window, Colors[grid[x][y]], rect)

4. Pretty much everything with a dot (.) or an expression should be cached inside display.

pgrect = pygame.Rect
rappend = rects.append
drawrect = pygame.draw.rect

5. Your x/y loop can be replaced using np.apply_over_axes

6. The scipy.signal.convolve2d function will do most of the heavy lifting for you in terms of neighbor counting &c

• Thank you for your answer, I don't understand how I can implement the functions np.apply_over_axes and scipy.signal.convolve2d, I've looked at the documentation, but I don't really understand what they do. – Antetokounpo Dec 29 '17 at 23:54