Run Conway's Game of Life on a GPU using PyTorch

game_of_life_pytorch.py

import torch
import torch.nn.functional as F
from PIL import Image

# Check CUDA availability
device = "cuda" if torch.cuda.is_available() else "cpu"

game_of_life_kernel = torch.tensor([
    [1, 1, 1],
    [1, 0, 1],
    [1, 1, 1],
], dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)

def game_of_life_step(life):
    # Use convolution to sum neighbors
    # we have to use float32 because uint8 is not supported by conv2d
    float_life = life.to(torch.float32)
    neighbors = F.conv2d(float_life.unsqueeze(0).unsqueeze(0), game_of_life_kernel, padding=1).to(torch.uint8)
    new_life = ((neighbors == 3) | (life & (neighbors == 2))).to(torch.uint8)

    return new_life.squeeze(0).squeeze(0)


def animate_game_of_life(size=128, generations=50, filename='game_of_life.gif'):
    life = torch.randint(0, 2, (size, size), dtype=torch.uint8).to(device)
    # Prepare for animation
    frames = []
    for _ in range(generations):
        frame = Image.fromarray(life.cpu().numpy() * 255)
        frames.append(frame)
        life = game_of_life_step(life)
    
    # Save the frames as an animated GIF
    frames[0].save(filename, save_all=True, append_images=frames[1:], duration=100, loop=1)

if __name__ == '__main__':
    animate_game_of_life(256, 256, filename="game_of_life.256.256.gif")