﻿ 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.
9. CAPACITY = 4
10.
11. # This breaks down the display area into populated rectangles. The population is based on points
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)
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)
118.
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():
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