Plinko implemented with numpy and threads

plinko-numpy.py

# Created by Anthony Cook
# fb.com/groups/pygame

import pygame
from pygame import gfxdraw
import random
import numpy as np
import time
import os
from itertools import product, count
from threading import Thread


class SpatialHash:
    def __init__(self):
        self.grid = {}

    def insert_point(self, point):
        x, y = point
        cell = x // 64, y // 64
        items = self.grid.get(cell)
        if items is None:
            self.grid[cell] = [point]
        else:
            items.append(point)

    def insert(self, rect, value):
        for cell in self._rect_cells(rect):
            items = self.grid.get(cell)
            if items is None:
                self.grid[cell] = [value]
            else:
                items.append(value)

    def _rect_cells(self, rect):
        x1, y1 = rect.topleft
        x1 //= 64
        y1 //= 64
        x2, y2 = rect.bottomright
        x2 = x2 // 64 + 1
        y2 = y2 // 64 + 1
        return product(range(x1, x2), range(y1, y2))

    def query(self, rect):
        items = []
        for cell in self._rect_cells(rect):
            items.extend(self.grid.get(cell, ()))
        return items

    def query_point(self, pos):
        x, y = pos
        return self.grid.get((x // 64, y // 64), ())



# pygame display settings
DR = pygame.Rect((0, 0, 1280, 720)) # Display Rectangle
HDW, HDH = DR.center # H = half
FPS = 60

# set up pygame
pygame.init()
PD = pygame.display.set_mode(DR.size) # primary display based of the size of Display Rect (800, 600)
CLOCK = pygame.time.Clock()

# set strength of gravity
GRAVITY = pygame.math.Vector2(0, 9.8)

# set up stage
SCALE = 10
STAGE = pygame.Rect((0, 0, DR.w * SCALE, DR.h * SCALE))

# create a quadtree instance and initialise it with the display rect (size of the display)
spatial = SpatialHash()

# add 1000 randomly placed Plinko pins to quadtree
PIN_COUNT = 1000
pins = []
for index in range(PIN_COUNT):
    pos = pygame.math.Vector2(
        random.randint(0, STAGE.w),
        random.randint(0, STAGE.h)
    )
    bounds = pygame.Rect(pos.x - 30, pos.y - 30, 60, 60)
    spatial.insert(bounds, pos)
    pins.append(pos)

# create 1000 marbles
MARBLE_COUNT = 1000
MARBLE_SPACING = STAGE.w / MARBLE_COUNT

marble_pos = np.array([
    np.linspace(0, STAGE.w, MARBLE_COUNT),
    np.zeros(MARBLE_COUNT)
]).T.copy()
marble_vel = np.array([
    np.random.uniform(-10, 10, MARBLE_COUNT),
    np.zeros(MARBLE_COUNT)
]).T.copy()


# the viewport is like a window that looks onto the stage
# this sets the location and size of the viewport which has a minimum size set to
# that of the primary display surface and starts in the top left of the stage
VIEWPORT = pygame.Rect(DR)

# exit the demo?
exit = False
MARBLE_COLOR = (0, 255, 0)


def update_marbles():
    global marble_pos, marble_vel
    # update marbles
    marble_pos += marble_vel
    marble_pos %= STAGE.bottomright
    marble_vel += GRAVITY / FPS

    for pos, vel in zip(marble_pos, marble_vel):
        # query the quadtree for all the Plinko pins nearest the marble
        # and store the result in closestPins
        closestPins = spatial.query_point(pos)

        # if there are pins then determine if the marble has collided with them
        for cp in closestPins:
            # is the distance between the marble and the pin less than
            # their combined radius's?
            if cp.distance_squared_to(pos) <= (26 * 26):
                # if yes then a collision has occurred and we need to calculate a new
                # trajectory for the marble. This basically the opposite direction to which
				# it was going
                angle = pygame.math.Vector2().angle_to(cp - pos)
                vel[:] = pygame.math.Vector2(*pos).rotate(angle - 180).normalize() * 10


start = time.time()
dur = 0
marble_current_pos = marble_pos.copy()
update = Thread(target=update_marbles)
update.start()
for framenum in count():
    # process events
    for e in pygame.event.get():
        if e.type == pygame.QUIT: # window close gadget
            exit = True
        elif e.type == pygame.MOUSEBUTTONDOWN:
            if e.button == 5: # mouse wheel down
                # increase the size of the viewport
                VIEWPORT.w += 18
                VIEWPORT.h += 10
            elif e.button == 4: # mouse wheel up
                # decrease the size of the viewport
                VIEWPORT.w -= 18
                VIEWPORT.h -= 10
            # limit the mimium size of the viewport
            # to that of the primary display resolution
            if VIEWPORT.w < DR.w:
                VIEWPORT.w = DR.w
                VIEWPORT.h = DR.h

    # exit the demo if ESC pressed or exit is True (set by pressing window x gadget)
    if pygame.key.get_pressed()[pygame.K_ESCAPE] or exit: break

    # get the distance the mouse has travelled since last get_rel() call
    rx, ry = pygame.mouse.get_rel()
    # if left mouse button has been pressed then you can drag the viewport about
    if pygame.mouse.get_pressed()[0]:
        # move the viewport
        VIEWPORT.x -= rx
        VIEWPORT.y -= ry
        # limit the viewport to stage's boundry
        if VIEWPORT.right > STAGE.w:
            VIEWPORT.x = STAGE.w - VIEWPORT.w
        if VIEWPORT.x < 0:
            VIEWPORT.x = 0
        if VIEWPORT.bottom > STAGE.h:
            VIEWPORT.y = STAGE.h - VIEWPORT.h
        if VIEWPORT.y < 0:
            VIEWPORT.y = 0

    # clear the primary display (fill it with black)
    PD.fill((0, 0, 0))


    # calculate the scale of the viewport against the size of the primary display
    scale = VIEWPORT.w / DR.w

    # draw all of the Plinko board's pins if they're within the viewport
    cull = VIEWPORT.inflate(20, 20)
    radius = int(20 / scale)
    for pin in pins:
        if cull.collidepoint(pin):
            pos = (pin - VIEWPORT.topleft) / scale
            pygame.draw.circle(PD, (255, 0, 0), pos, radius)

    update.join()
    marble_current_pos[:] = marble_pos
    update = Thread(target=update_marbles)
    update.start()

    # draw all of the marbles
    marble_screen_pos = (marble_current_pos - VIEWPORT.topleft) / scale
    radius = 6 / scale
    color = MARBLE_COLOR
    if radius < 1:
        color = (0, round(radius * 255), 0)
        radius = 1
    else:
        radius = round(radius)
    for m in marble_screen_pos:
        if cull.collidepoint(m):
            pygame.draw.circle(PD, color, m, radius)

    # update the primary display
    if framenum % 100 == 99:
        os.write(1, f"FPS: {100 / dur:0.2f}\n".encode('ascii'))
        dur = 0
    else:
        pygame.display.update()

    end = time.time()
    dur += end - start
    pygame.display.update()
    dt = CLOCK.tick(FPS) # limit frames
    start = time.time()