Simulation of sand particle collision

Question

I am making a module that allows you to create sand and simulate it. I have successfully done this, but it runs very slowly. I need this to be able to simulate at least 1000 particles, at a fps higher than 5. How do I optimize this code?

My code:

import pygame, random, sys
pygame.init()
FPS = 60 
fpsClock = pygame.time.Clock()
screen = pygame.display.set_mode((400, 300), 0, 32)
pygame.display.set_caption('Sand')

class SandSource:
    def __init__(self, screen, colliders):
        self.screen = screen
        self.tiles = colliders
        self.sand_arr = []
        self.colour = (255, 255, 0, 255)
        self.gradient = True
        self.gradients = []

    def update(self):
        for sand in self.sand_arr:
            sand1 = self.custom_sand_collision(sand)
            if self.gradient:
                pygame.draw.rect(self.screen, sand[2], sand1)
            else:
                pygame.draw.rect(self.screen, self.colour, sand1)

    def check_pos_of(self, x, y):
        if self.screen.get_at((x, y)) in self.gradients:
            return True
        else:
            return False

    def create_sand(self, x, y, sands=1):
        for i in range(sands):
            if self.gradient:
                COLOR = list(self.colour)
                COLOR[1] = COLOR[1] + random.randint(-50, 0)
                COLOR[0] = COLOR[0] + random.randint(-50, 0)
                COLOR = tuple(COLOR)
                if COLOR not in self.gradients:
                    self.gradients.append(COLOR)
                self.sand_arr.append([x, y, COLOR])
            else:
                self.sand_arr.append([x, y])

    def check_collide(self, x, y, rext):
        rext.x = x
        rext.y = y
        if self.collision_test(rext) or self.check_pos_of(x, y):
            return False
        else:
            return True

    def custom_sand_collision(self, sand):
        ox = sand[0]
        oy = sand[1]
        vel = 4
        check = pygame.Rect((ox, oy), (1, 1))
        try:
            if self.collision_test(check):
                pass
            elif self.check_collide(ox, oy + 1, check):
                for i in range(vel):
                    if self.check_collide(ox, oy + i, check):
                        sand[1] = oy + i
            elif self.check_collide(ox - 1, oy + 1, check):
                sand[1] = oy + 1
                sand[0] -= 1
            elif self.check_collide(ox + 1, oy + 1, check):
                sand[1] += 1
                sand[0] += 1
        except:
            pass
        check.x = sand[0]
        check.y = sand[1]
        return check

    def collision_test(self, rect):
        for thing in self.tiles:
            if thing.colliderect(rect):
                return True
        return False

a = SandSource(screen, colliders=[pygame.Rect((0, 250), (300, 300))])
a.create_sand(100, 100, 1000)


while True:
    screen.fill((0, 0, 0))

    a.update()
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

    pygame.display.update()
    fpsClock.tick(FPS)

Answer 1

• Don't write, never commit/publish uncommented/undocumented code.
  Tell, in the code, what "everything" is there for.
  Python got it right specifying docstrings such that it is easy to copy them with the code, and tedious to copy the code without them.

• just return condition (or return bool(condition) where not already boolean) instead of

  if condition:
      return True
  else:
      return False
  

• I find using sand[1] = oy + 1 in one branch and sand[1] += 1 in the next irritating.
  Coming to think of it, the "velocity branch" looks particularly peculiar, starting with checking the original location and continuing with a re-check of ox, oy + 1.

  - checking each abutting coordinate independently looks brute force.

• If self.gradients gets beyond small, using a set may reduce execution time.