cookertron

Quadtree Plinko Board - Python Pygame

May 22nd, 2020
1,355
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. # Created by Anthony Cook
  2. # fb.com/groups/pygame
  3.  
  4. import pygame
  5. from pygame import gfxdraw
  6. import random
  7.  
  8. # quadTree constants
  9. CAPACITY = 4
  10.  
  11. # This breaks down the display area into populated rectangles. The population is based on points
  12. class quadTree:
  13.     def __init__(s, rect): # boundry is a rectangle
  14.         s.rect = rect
  15.         s.points = []
  16.         s.divided = False
  17.        
  18.     def insert(s, point):
  19.         global CAPACITY # maximum capacity on each quadrant
  20.  
  21.         # if the point inserted doesn't fit into this quadrant then don't add it to
  22.         # the points list. Exit insert()!
  23.         if not s.rect.collidepoint(int(point.x), int(point.y)): return False
  24.  
  25.         # check if the quadrant has room for more points
  26.         if len(s.points) < CAPACITY:
  27.             s.points.append(point) # if there's room add it to the points list
  28.             # return True to inform parent quadtree the point has been added
  29.             return True
  30.         else:
  31.             # if the quadrant is full then divide the quadrant into
  32.             # four parts north west, north east, south west, south east...
  33.             if not s.divided:
  34.                 s.subDivide()
  35.        
  36.             # ...and insert the point into one of the four new quadrants
  37.             if s.nw.insert(point):
  38.                 return True
  39.             elif s.ne.insert(point):
  40.                 return True
  41.             elif s.sw.insert(point):
  42.                 return True
  43.             elif s.se.insert(point):
  44.                 return True
  45.  
  46.     def subDivide(s):
  47.         # create four new quadtrees which belong to the parent quadtree
  48.         size = (s.rect.width / 2, s.rect.height / 2)
  49.         s.nw = quadTree(pygame.Rect(s.rect.topleft, size))
  50.         s.ne = quadTree(pygame.Rect(s.rect.midtop, size))
  51.         s.sw = quadTree(pygame.Rect(s.rect.midleft, size))
  52.         s.se = quadTree(pygame.Rect(s.rect.center, size))
  53.         s.divided = True # the parent quadtree can no longer be divided
  54.  
  55.     def query(s, rect, pointsArray = []):
  56.         # if the queried rect doesn't intersect with the quadrant
  57.         # exit query
  58.         if not s.rect.colliderect(rect): return
  59.  
  60.         # if it does then append points to the pointsArray
  61.         for p in s.points:
  62.             if rect.collidepoint(p): pointsArray.append(p)
  63.  
  64.         # because the quadtree is recursive we need to check its children
  65.         # to see if they're within the query parameters (rect)
  66.         if s.divided:
  67.             s.nw.query(rect, pointsArray)
  68.             s.ne.query(rect, pointsArray)
  69.             s.sw.query(rect, pointsArray)
  70.             s.se.query(rect, pointsArray)
  71.  
  72. # define a marble class
  73. class marble:
  74.     def __init__(s, x, y):
  75.         # pos is a vector (x, y)
  76.         s.pos = pygame.math.Vector2(x, y)
  77.         # velocity is also a vector with a starting x velocity of random (-1, 1) * 10
  78.         s.velocity = pygame.math.Vector2(random.uniform(-1, 1), 0) * 10
  79.  
  80.     def update(s):
  81.         global GRAVITY, FPS, STAGE
  82.        
  83.         # apply velocity to the marble's position
  84.         s.pos += s.velocity
  85.  
  86.         # apply gravity to the marble's velocity
  87.         s.velocity += GRAVITY / FPS
  88.  
  89.         # if a marble goes outside of the stage then it reappears
  90.         # on the opposite side. For instance if the marble drops out of the
  91.         # bottom of the stage then it will reappear at the top
  92.         if s.pos.y >= STAGE.bottom:
  93.             s.pos.y = 0
  94.         if s.pos.x > STAGE.right:
  95.             s.pos.x = 0
  96.         if s.pos.x < 0:
  97.             s.pos.x = STAGE.right
  98.  
  99. # pygame display settings
  100. DR = pygame.Rect((0, 0, 1280, 720)) # Display Rectangle
  101. HDW, HDH = DR.center # H = half
  102. FPS = 60
  103.  
  104. # set up pygame
  105. pygame.init()
  106. PD = pygame.display.set_mode(DR.size) # primary display based of the size of Display Rect (800, 600)
  107. CLOCK = pygame.time.Clock()
  108.  
  109. # set strength of gravity
  110. GRAVITY = pygame.math.Vector2(0, 9.8)
  111.  
  112. # set up stage
  113. SCALE = 10
  114. STAGE = pygame.Rect((0, 0, DR.w * SCALE, DR.h * SCALE))
  115.  
  116. # create a quadtree instance and initialise it with the display rect (size of the display)
  117. QT = quadTree(STAGE)
  118.  
  119. # add 1000 randomly placed Plinko pins to quadtree
  120. PIN_COUNT = 1000
  121. for index in range(PIN_COUNT):
  122.     QT.insert(pygame.math.Vector2(random.randint(0, STAGE.w), random.randint(0, STAGE.h)))
  123.  
  124. # create 1000 marbles
  125. MARBLE_COUNT = 1000
  126. MARBLE_SPACING = STAGE.w / MARBLE_COUNT
  127. MARBLES = [marble(index * MARBLE_SPACING, 0) for index in range(MARBLE_COUNT)]
  128.  
  129. # the viewport is like a window that looks onto the stage
  130. # this sets the location and size of the viewport which has a minimum size set to
  131. # that of the primary display surface and starts in the top left of the stage
  132. VIEWPORT = pygame.Rect(DR)
  133.  
  134. # exit the demo?
  135. exit = False
  136.  
  137. while True:
  138.     # process events
  139.     for e in pygame.event.get():
  140.         if e.type == pygame.QUIT: # window close gadget
  141.             exit = True
  142.         elif e.type == pygame.MOUSEBUTTONDOWN:
  143.             if e.button == 5: # mouse wheel down
  144.                 # increase the size of the viewport
  145.                 VIEWPORT.w += 18
  146.                 VIEWPORT.h += 10
  147.             elif e.button == 4: # mouse wheel up
  148.                 # decrease the size of the viewport
  149.                 VIEWPORT.w -= 18
  150.                 VIEWPORT.h -= 10
  151.             # limit the mimium size of the viewport
  152.             # to that of the primary display resolution
  153.             if VIEWPORT.w < DR.w:
  154.                 VIEWPORT.w = DR.w
  155.                 VIEWPORT.h = DR.h
  156.  
  157.     # exit the demo if ESC pressed or exit is True (set by pressing window x gadget)
  158.     if pygame.key.get_pressed()[pygame.K_ESCAPE] or exit: break
  159.  
  160.     # get the distance the mouse has travelled since last get_rel() call
  161.     rx, ry = pygame.mouse.get_rel()
  162.     # if left mouse button has been pressed then you can drag the viewport about
  163.     if pygame.mouse.get_pressed()[0]:
  164.         # move the viewport
  165.         VIEWPORT.x -= rx
  166.         VIEWPORT.y -= ry
  167.         # limit the viewport to stage's boundry
  168.         if VIEWPORT.right > STAGE.w:
  169.             VIEWPORT.x = STAGE.w - VIEWPORT.w
  170.         if VIEWPORT.x < 0:
  171.             VIEWPORT.x = 0
  172.         if VIEWPORT.bottom > STAGE.h:
  173.             VIEWPORT.y = STAGE.h - VIEWPORT.h
  174.         if VIEWPORT.y < 0:
  175.             VIEWPORT.y = 0
  176.  
  177.     # clear the primary display (fill it with black)
  178.     PD.fill((0, 0, 0))
  179.  
  180.     # query the quadtree for the location of all the pins within the viewport
  181.     pinPoints = []
  182.  
  183.     # after calling query() pinPoints contains a list of pygame.math.Vector2() points
  184.     # representing the Plinko boards pins
  185.     QT.query(VIEWPORT, pinPoints)
  186.    
  187.     # calculate the scale of the viewport against the size of the primary display
  188.     scale = VIEWPORT.w / DR.w
  189.  
  190.     # draw all of the Plinko board's pins if they're within the viewport
  191.     for pin in pinPoints:
  192.         pygame.draw.circle(PD, (255, 0, 0), (int((pin.x - VIEWPORT.x) / scale), int((pin.y - VIEWPORT.y) / scale)), int(20 / scale))
  193.  
  194.     # draw all of the marbles and update marble position
  195.     for m in MARBLES:
  196.         pygame.draw.circle(PD, (0, 255, 0), (int((m.pos.x - VIEWPORT.x) / scale), int((m.pos.y - VIEWPORT.y) / scale)), int(6 / scale))
  197.        
  198.         # update position of the marble (apply velocity)
  199.         m.update()
  200.        
  201.         closestPins = []
  202.        
  203.         # query the quadtree for all the Plinko pins nearest the marble
  204.         # and store the result in closestPins
  205.         QT.query(pygame.Rect(m.pos.x - 30, m.pos.y - 30, 60, 60), closestPins)
  206.        
  207.         # if there are pins then determine if the marble has collided with them
  208.         for cp in closestPins:
  209.             # is the distance between the marble and the pin less than
  210.             # their combined radius's?
  211.             if cp.distance_to(m.pos) <= 26:
  212.                 # if yes then a collision has occurred and we need to calculate a new
  213.                 # trajectory for the marble. This basically the opposite direction to which
  214.                 # it was going
  215.                 angle = pygame.math.Vector2().angle_to(cp - m.pos)
  216.                 m.velocity = m.pos.rotate(angle - 180).normalize() * 10
  217.  
  218.     # update the primary display
  219.     pygame.display.update()
  220.     CLOCK.tick(FPS) # limit frames
RAW Paste Data