Efficiency in Infections - Pygame

Question

I am currently making an infection survival game for my A-level coursework, and I am struggling on how I can make this efficient.
When the cells get infected I need to check more and more cells, my computer science teacher recommended I save infections as a boolean value as I can do a lot more with that later, however it makes it more inefficient due to me having to eventually check the amount of cells squared which causes a lot of framerate issues.
My original idea was to store the uninfected and infected in separate lists but my comp sci teacher said I was over complicating it, however this didn't have any framerate issues.
A lot of my code has taken inspiration from this question Random movement pygame, especially when it comes to the cell movement.

TLDR: I want to make my code more efficient but I can't think of how
My code:

import sys, random, pygame
import matplotlib.pyplot as plt
from pygame.locals import *
import time
pygame.init()

#Window details
windowWidth = 400
windowHeight = 400
pixSize = 2
FPS = 60
screen = pygame.display.set_mode((windowWidth, windowHeight))
pygame.display.set_caption("Infection Game")
class Cell:
    def __init__(self):
        self.xPos = random.randrange(1,windowWidth)
        self.yPos = random.randrange(1,windowHeight)
        self.speed = 2
        self.isInfected = False
        self.infectionRange = 5
        self.move = [None, None]
        self.direction = None
    def cellDraw(self):
        if self.isInfected == False:
            pygame.draw.rect(screen, (255,255,255), (self.xPos,self.yPos,pixSize,pixSize),0)
        else:
            pygame.draw.rect(screen, (0,255,0), (self.xPos,self.yPos,pixSize,pixSize),0)
    def cellMovement(self):
        directions = {"S":((-1,2),(1,self.speed)),"SW":((-self.speed,-1),(1,self.speed)),"W":((-self.speed,-1),(-1,2)),"NW":((-self.speed,-1),(-self.speed,-1)),"N":((-1,2),(-self.speed,-1)),"NE":((1,self.speed),(-self.speed,-1)),"E":((1,self.speed),(-1,2)),"SE":((1,self.speed),(1,self.speed))} #((min x, max x)(min y, max y))
        directionsName = ("S","SW","W","NW","N","NE","E","SE") #possible directions
        if random.randrange(0,5) == 2: #move about once every 5 frames
            if self.direction == None: #if no direction is set, set a random one
                self.direction = random.choice(directionsName)
            else:
                a = directionsName.index(self.direction) #get the index of direction in directions list
                b = random.randrange(a-1,a+2) #set the direction to be the same, or one next to the current direction
                if b > len(directionsName)-1: #if direction index is outside the list, move back to the start
                    b = 0
                self.direction = directionsName[b]

            self.move[0] = random.randrange(directions[self.direction][0][0],directions[self.direction][0][1]) + 0.35
            self.move[1] = random.randrange(directions[self.direction][1][0],directions[self.direction][1][1]) + 0.35
        if self.xPos < 5 or self.xPos > windowWidth - 5 or self.yPos < 5 or self.yPos > windowHeight - 5: #if cell is near the border of the screen, change direction
            if self.xPos < 5:
                self.direction = "E"
            elif self.xPos > windowWidth - 5:
                self.direction = "W"
            elif self.yPos < 5:
                self.direction = "S"
            elif self.yPos > windowHeight - 5:
                self.direction = "N"

            self.move[0] = random.randrange(directions[self.direction][0][0],directions[self.direction][0][1]) + 0.35
            self.move[1] = random.randrange(directions[self.direction][1][0],directions[self.direction][1][1]) + 0.35
        if self.move[0] != None: #add the relative coordinates to the cells coordinates
            self.xPos += self.move[0]
            self.yPos += self.move[1]
    def Infect(self):
        for i in cellList:
            if (self.xPos > i.xPos - self.infectionRange and self.xPos < i.xPos + self.infectionRange) and (self.yPos > i.yPos - self.infectionRange and self.yPos < i.yPos + self.infectionRange):
                i.isInfected = True
        
xgraph = []
ygraph = []
cellList = []
startTime = time.time()
for i in range(1000):
    cell = Cell()
    cellList.append(cell)
cellList[0].isInfected = True

def gameLoop():
    while True:
        infectList = []
        for event in pygame.event.get():
            if event.type == QUIT:
                pygame.quit()
                sys.exit()
        screen.fill((0,0,0))
        for i in cellList:
            i.cellDraw()
            i.cellMovement()
        for i in cellList:
            if i.isInfected == True:
                i.Infect()
                infectList.append(i)
        xgraph.append(time.time()-startTime)
        ygraph.append(len(infectList))
        plt.plot(xgraph,ygraph)
        plt.xlabel('time (s)')
        plt.ylabel('infected')
        if len(infectList) == 1000:
            plt.show()
        pygame.display.update() #update display
        pygame.time.Clock().tick(FPS) #limit FPS

gameLoop()

user_6766457 · Answer

First off, I've changed some of your code:

if self.isInfected == False:
if self.direction == None:

To

if not self.isInfected:
if self.direction is None:

Just so it reads a little a nicer.

Secondly, I've vectorized the Infect function:

uninfected = [i for i in cellList if not i.isInfected]
uninfected_array = np.array([[i.xPos, i.yPos] for i in uninfected])
indices = np.greater(uninfected_array[:, 0], self.xPos - self.infectionRange) * \
          np.greater(self.xPos + self.infectionRange, uninfected_array[:, 0]) * \
          np.greater(uninfected_array[:, 1], self.yPos - self.infectionRange) * \
          np.greater(self.yPos + self.infectionRange, uninfected_array[:, 1])

    
for i in np.where(indices)[0]:
    uninfected[i].isInfected = True

It takes the same time for this number of cells, but should scale better.

It turns out creating the array takes almost all the time. So you can create it once, pull it out of the loop and shave off a bunch of time:

def Infect(self, uninfected, uninfected_array):
    indices = np.greater(uninfected_array[:, 0], self.xPos - self.infectionRange) * \
              np.greater(self.xPos + self.infectionRange, uninfected_array[:, 0]) * \
              np.greater(uninfected_array[:, 1], self.yPos - self.infectionRange) * \
              np.greater(self.yPos + self.infectionRange, uninfected_array[:, 1])

    for i in np.where(indices)[0]:
        uninfected[i].isInfected = True

uninfected = [i for i in cellList if not i.isInfected]
uninfected_array = np.array([[i.xPos, i.yPos] for i in uninfected])
# To prevent errors with empty arrays
if len(uninfected) > 0:
    for i in cellList:
        if i.isInfected:
            i.Infect(uninfected, uninfected_array)
# To prevent errors when everyone is infected
if infected == 0:
    infected = len(cellList) - len(uninfected)

Lastly, you don't really seem to be using the infectList, so I replaced it with a counter:

infected = 0
if len(uninfected) > 0:
    for i in cellList:
        if i.isInfected:
            infected += 1

As a side note, I'd change the UI controls a bit so it's easier to graph, instead of quitting using sys.exit it's nicer just to break out of the while loop. You also only plot the results once:

running = True
while running:
    infectList = []
    for event in pygame.event.get():
        if event.type == QUIT:
            running = False
    ...
pygame.quit()
plt.plot(xgraph, ygraph)
plt.xlabel('time (s)')
plt.ylabel('infected')
plt.show()

Implementing all this results in:

import random
import pygame
import matplotlib.pyplot as plt
from pygame.locals import *
import time
import numpy as np

pygame.init()

# Window details
windowWidth = 400
windowHeight = 400
pixSize = 2
FPS = 60
screen = pygame.display.set_mode((windowWidth, windowHeight))
pygame.display.set_caption("Infection Game")


class Cell:
    def __init__(self):
        self.xPos = random.randrange(1, windowWidth)
        self.yPos = random.randrange(1, windowHeight)
        self.speed = 2
        self.isInfected = False
        self.infectionRange = 5
        self.move = [None, None]
        self.direction = None

    def cellDraw(self):
        if not self.isInfected:
            pygame.draw.rect(screen, (255, 255, 255), (self.xPos, self.yPos, pixSize, pixSize), 0)
        else:
            pygame.draw.rect(screen, (0, 255, 0), (self.xPos, self.yPos, pixSize, pixSize), 0)

    def cellMovement(self):
        directions = {"S": ((-1, 2), (1, self.speed)), "SW": ((-self.speed, -1), (1, self.speed)),
                      "W": ((-self.speed, -1), (-1, 2)), "NW": ((-self.speed, -1), (-self.speed, -1)),
                      "N": ((-1, 2), (-self.speed, -1)), "NE": ((1, self.speed), (-self.speed, -1)),
                      "E": ((1, self.speed), (-1, 2)),
                      "SE": ((1, self.speed), (1, self.speed))}  # ((min x, max x)(min y, max y))
        directionsName = ("S", "SW", "W", "NW", "N", "NE", "E", "SE")  # possible directions
        if random.randrange(0, 5) == 2:  # move about once every 5 frames
            if self.direction is None:  # if no direction is set, set a random one
                self.direction = random.choice(directionsName)
            else:
                a = directionsName.index(self.direction)  # get the index of direction in directions list
                b = random.randrange(a - 1,
                                     a + 2)  # set the direction to be the same, or one next to the current direction
                if b > len(directionsName) - 1:  # if direction index is outside the list, move back to the start
                    b = 0
                self.direction = directionsName[b]

            self.move[0] = random.randrange(directions[self.direction][0][0], directions[self.direction][0][1]) + 0.35
            self.move[1] = random.randrange(directions[self.direction][1][0], directions[self.direction][1][1]) + 0.35
        if self.xPos < 5 or self.xPos > windowWidth - 5 or self.yPos < 5 or self.yPos > windowHeight - 5:  # if cell is near the border of the screen, change direction
            if self.xPos < 5:
                self.direction = "E"
            elif self.xPos > windowWidth - 5:
                self.direction = "W"
            elif self.yPos < 5:
                self.direction = "S"
            elif self.yPos > windowHeight - 5:
                self.direction = "N"

            self.move[0] = random.randrange(directions[self.direction][0][0], directions[self.direction][0][1]) + 0.35
            self.move[1] = random.randrange(directions[self.direction][1][0], directions[self.direction][1][1]) + 0.35
        if self.move[0] is not None:  # add the relative coordinates to the cells coordinates
            self.xPos += self.move[0]
            self.yPos += self.move[1]

    def Infect(self, uninfected, uninfected_array):
        indices = np.greater(uninfected_array[:, 0], self.xPos - self.infectionRange) * \
                  np.greater(self.xPos + self.infectionRange, uninfected_array[:, 0]) * \
                  np.greater(uninfected_array[:, 1], self.yPos - self.infectionRange) * \
                  np.greater(self.yPos + self.infectionRange, uninfected_array[:, 1])

        for i in np.where(indices)[0]:
            uninfected[i].isInfected = True


xgraph = []
ygraph = []
cellList = []
startTime = time.time()
for i in range(1000):
    cell = Cell()
    cellList.append(cell)
cellList[0].isInfected = True


def gameLoop():
    running = True
    while running:
        infectList = []
        for event in pygame.event.get():
            if event.type == QUIT:
                running = False
        screen.fill((0, 0, 0))
        for i in cellList:
            i.cellDraw()
            i.cellMovement()
        infected = 0

        uninfected = [i for i in cellList if not i.isInfected]
        uninfected_array = np.array([[i.xPos, i.yPos] for i in uninfected])
        if len(uninfected) > 0:
            for i in cellList:
                if i.isInfected:
                    i.Infect(uninfected, uninfected_array)
                    infected += 1
        if infected == 0:
            infected = len(cellList) - len(uninfected)
        xgraph.append(time.time() - startTime)
        ygraph.append(infected)
        pygame.display.update()  # update display
        pygame.time.Clock().tick(FPS)  # limit FPS
    pygame.quit()

    # figured this is what you wanted to do ;)
    plt.plot(xgraph, ygraph)
    plt.xlabel('time (s)')
    plt.ylabel('infected')
    plt.show()


gameLoop()

And it runs smooth

Efficiency in Infections - Pygame

1 Answers