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)

• Reinderien: Have you read about Numpy?
• susthebus: How would I use numpy to optimize this?
• Reinderien: Read about vectorisation. Reinterpret your sand particles as being in a vector upon which operations need to be done on all of them at once.
• susthebus: I will look into it!
• Francesco Pasa: I am not sure if it supports pygame, but a quick way would be to try running the code with pypy instead of cpython.

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.

• susthebus: This did help raise the framerate a little bit! Thanks.