Pythonで千個くらいのボールの衝突判定をハッシュを使って軽くする（新型コロナウィルス関連）
ワシントンポスト-オンラインの記事
この↓ワシントンポストの新型コロナウィルスの感染シミュレーションの記事（現象の説明）
に刺激されて、沢山のボールが衝突するプログラムをPythonで作ってみた。
（ちなみにこの記事はワシントンポストのオンラインで、それまでで最も見られた記事になったようです。）
overshoot.py
スペースキーをプレスすると、最初の感染が始まる。
Pythonista バージョン
Pythonisitaバージョンはspawnで感染スタート
ボールの衝突そのものは、 Simpson College Computer Scienceの以下の記事などを
参考にした。
http://programarcadegames.com/python_examples/f.php?file=bouncing_balls.py
問題は、ボールが1000個くらいになったら、とても動きがスローになる。常に一つのボールが他の
すべてのボールと衝突判定をしているためである。
そこで、近傍のボールとだけの衝突判定をすればいいわけだが、この近傍のボール同士を決めるの
にPythonのハッシュ関数を使う。
レイトレーシングなどでレイとオブジェクトの交差判定の計算をする場合も、Oct Treeなどで、
空間分割をするのに似ている。
ハッシュとは？
ハッシュとは、乱暴に言うと、ある変数（文字列、数値など）を与えると、それが入る箱の名前を決める関数（関数と言っていいのか？）で変数は複数でもよい。今回の場合、ボールの二次元座標（x,y)を与えるとその座標値からその位置が入る箱の番号を与える。
たとえば、x,yが、1.24と3.24とすると、それぞれの小数点以下を切り捨てて、(1,3)という名前の箱に入れることができる。この箱には(1.55, 1.99)の位置のボールも同じ名前の箱に入る。箱の名前を決定する計算式は用途に合わせて適宜考える。Pythonにはこのハッシュが用意されているので便利！
こうして、交差判定を同じ箱に入っているもの同士のボールに限れば計算がぐっと軽くなる。しかし、実際はボールに大きさがあるので、隣あう箱どうしにボールの中心が存在して、それと衝突する可能性も考慮する。
(ただし下記のコードがすべてボールの衝突を完全に正しく反映しているかは保証できない。）
一応動いたのが以下のコードであるが、Pygameを用いているので、そのモジュールもインストールする必要がある。このコードではただボールが衝突して衝突したら必ず感染するのですぐに爆発的に感染するが、潜伏期、発症期、死亡、抗体獲得、再生産数、などを設定しなければならない。
コード
（コメントに@shiracamusが、classを使ってコードを整理してくださったバージョンがありすので、、そちらをご覧下さい）
overshoot.py
import pygame
import random
import math
#import numpy as np
#import copy
# Define colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BLUE = (90, 90, 180)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
GREY = (50, 50, 255)
SCREEN_WIDTH = 700
SCREEN_HEIGHT = 650
start_time = pygame.time.get_ticks()
class Ball:
"""
Class to keep track of a ball's location and vector.
"""
def __init__(self):
# radius
self.r = 4
# diameter
self.d = self.r * 2
# position
self.x = 0
self.y = 0
# velocity
self.vx = 0
self.vy = 0
# color
self.color = (0, 255, 0)
# neighbour
self.px = False
self.mx = False
self.py = False
self.my = False
def make_ball():
"""
Function to make a new, random ball.
"""
ball = Ball()
# Starting position of the ball.
# Take into account the ball size so we don't spawn on the edge.
ball.x = random.randrange(ball.d, SCREEN_WIDTH - ball.d)
ball.y = random.randrange(ball.d, SCREEN_HEIGHT - ball.d - 200)
# Speed and direction of rectangle
ball.vx = float(random.randrange(-2, 2)) * 0.2
ball.vy = float(random.randrange(-2, 2)) * 0.2
return ball
def main():
"""
main program.
"""
pygame.init()
xval = 0
bcount = 0
#redcount = 0
# Set the height and width of the screen
size = [SCREEN_WIDTH, SCREEN_HEIGHT]
screen = pygame.display.set_mode(size)
pygame.display.set_caption("OverShoot")
# Loop until the user clicks the close button.
done = False
# Used to manage how fast the screen updates
clock = pygame.time.Clock()
# dictionary
results = {}
cell_size = 40
ball_amount = 800
# Spawn many balls
for i in range(ball_amount):
ball = make_ball()
results.setdefault((int(ball.x/cell_size), int(ball.y/cell_size)), []).append(ball)
print(len(results.keys()))
#print(results)
screen.fill((20,20,20), rect=(0,SCREEN_HEIGHT - 250, SCREEN_WIDTH,SCREEN_HEIGHT))
# -------- Main Program Loop -----------
while not done:
# --- Event Processing
for event in pygame.event.get():
if event.type == pygame.QUIT:
done = True
elif event.type == pygame.KEYDOWN:
# Space bar! Spawn a new ball.
if event.key == pygame.K_SPACE:
ball = make_ball()
ball.color = RED
results.setdefault((int(ball.x/cell_size), int(ball.y/cell_size)), []).append(ball)
#
if event.key == pygame.K_g:
ball = make_ball()
ball.color = (255,0,255)
results.setdefault((int(ball.x/cell_size), int(ball.y/cell_size)), []).append(ball)
# --- Drawing
# Set the screen background
screen.fill(GREY, rect=(0,0,SCREEN_WIDTH,SCREEN_HEIGHT - 200 + ball.d))
# Draw the balls
cresults = {}
for ky in results:
blist = results[ky]
blist2 = blist.copy()
#
for bl in blist:
blist2.remove(bl)
if bl.px:
k = (ky[0]+1,ky[1])
if k in results:
ek = results[k]
blist2.extend(ek)
if bl.mx:
k = (ky[0]-1,ky[1])
if k in results:
ek = results[k]
blist2.extend(ek)
if bl.py:
k = (ky[0],ky[1]+1)
if k in results:
ek = results[k]
blist2.extend(ek)
if bl.my:
k = (ky[0],ky[1]-1)
if k in results:
ek = results[k]
blist2.extend(ek)
if bl.px and bl.py:
k = (ky[0]+1,ky[1]+1)
if k in results:
ek = results[k]
blist2.extend(ek)
if bl.mx and bl.py:
k = (ky[0]-1,ky[1]+1)
if k in results:
ek = results[k]
blist2.extend(ek)
if bl.px and bl.my:
k = (ky[0]+1,ky[1]-1)
if k in results:
ek = results[k]
blist2.extend(ek)
if bl.mx and bl.my:
k = (ky[0]-1,ky[1]-1)
if k in results:
ek = results[k]
blist2.extend(ek)
#
for bl2 in blist2:
#if bl == bl2:
# continue
# Circle Intersect
sqDist = (bl2.x-bl.x)*(bl2.x-bl.x)+(bl2.y-bl.y)*(bl2.y-bl.y)
#if sqDist > 0 and sqDist <= ((BALL_SIZE)+(BALL_SIZE)):
#print(sqDist, math.sqrt(sqDist))
if sqDist > 0 and sqDist <= (bl.d*bl.d):
# detect collision
dist = math.sqrt(sqDist)
# vCollison
vc = ([bl2.x-bl.x, bl2.y-bl.y])
# vCollisonNorm
vcn = ([vc[0]/dist, vc[1]/dist])
# vRelativeVelocity
vrv = ([bl.vx-bl2.vx, bl.vy-bl2.vy])
speed = vrv[0]*vcn[0] + vrv[1]*vcn[1]
#print(speed)
if speed > 0:
bl.vx -= speed * vcn[0]
bl.vy -= speed * vcn[1]
bl2.vx += speed * vcn[0]
bl2.vy += speed * vcn[1]
#
#infection
if bl.color == RED and bl2.color != RED:
bl2.color = RED
bcount += 1
if bl2.color == RED and bl.color != RED:
bl.color = RED
bcount += 1
#
col = bl.color
pygame.draw.circle(screen, col, [round(bl.x), round(bl.y)], bl.r)
bl.x += bl.vx
bl.y += bl.vy
# avoid line y
if bl.y > SCREEN_HEIGHT - 200 + bl.r:
bl.y = SCREEN_HEIGHT - 200
bl.vy = -bl.vy
if bl.y < bl.r :
bl.y = bl.r
bl.vy = -bl.vy
if bl.x < bl.r:
bl.x = bl.r
bl.vx = -bl.vx
if bl.x <= 0 :
bl.vx = 0.1
if bl.x > SCREEN_WIDTH - bl.r:
bl.x = SCREEN_WIDTH - bl.r
bl.vx = -bl.vx
#
#Bounce the ball if needed
if bl.y > SCREEN_HEIGHT - 200 or bl.y < bl.r:
bl.vy *= -1
if bl.x > SCREEN_WIDTH - bl.r or bl.x < bl.r:
bl.vx *= -1
# set key and get hash and append
cresults.setdefault((int(bl.x/cell_size), int(bl.y/cell_size)), []).append(bl)
#
# check neighbour with new key
bl.px = int((bl.x+bl.r)/cell_size) > int(bl.x/cell_size)
bl.mx = int((bl.x-bl.r)/cell_size) < int(bl.x/cell_size)
bl.py = int((bl.y+bl.r)/cell_size) > int(bl.y/cell_size)
bl.my = int((bl.y-bl.r)/cell_size) < int(bl.y/cell_size)
results.clear()
results.update(cresults)
# show stat
timepassed = pygame.time.get_ticks()-start_time
if timepassed % 12 == 0:
if bcount > 0 and bcount < ball_amount:
pygame.draw.line(screen,(200,0,0),(xval*2,SCREEN_HEIGHT-10),(xval*2,SCREEN_HEIGHT-int(bcount/5)-10),2)
xval += 1
#for res in results:
# a = results[res]
# for bl in a:
# pygame.draw.circle(screen, WHITE, [round(bl.x), round(bl.y)], BALL_SIZE)
# break
# Go ahead and update the screen with what we've drawn.
clock.tick(60)
pygame.display.flip()
# Close everything down
pygame.quit()
if __name__ == "__main__":
main()
Pythonistaバージョン
overshoot.py
import ui
from scene import *
import random
import math
import time
import sys
import os
# Define colors
BLACK = (0, 0, 0)
WHITE = (1, 1, 1)
BLUE = (0.5, 0.5, 0.8)
GREEN = (0, 1, 0)
RED = (1, 0, 0)
GREY = (0.5, 0.5, 0.5)
CELL_SIZE = 20
BALL_AMOUNT = 800
class Ball:
"""
Class to keep track of a ball's location and vector.
"""
def __init__(self):
# radius
self.r = 3
# diameter
self.d = self.r * 2
# position
self.x = 0
self.y = 0
# velocity
self.vx = 0
self.vy = 0
# color
self.color = (0, 1, 0)
# neighbour
self.nb = (False,False,False,False)
def make_ball(self):
"""
Function to make a new, random ball.
"""
ball = Ball()
# Starting position of the ball.
# Take into account the ball size so we don't spawn on the edge.
ball.x = random.randrange(ball.d, self.size.width - ball.d)
ball.y = random.randrange(270, self.size.height- ball.d-200)
# Speed and direction of rectangle
ball.vx = float(random.randrange(-2, 2)) * 0.2
ball.vy = float(random.randrange(-2, 2)) * 0.2
return ball
def button_tapped(sender):
#print('hhh')
os.abort()
class MyScene2(Scene):
def setup(self):
self.val = (0,0)
self.bar = []
self.last = 0
def add(self):
self.bar.append(self.val[1])
print(self.val[1])
def draw(self):
#print(self.val)
background(0.5,0.5,0.5)
fill(0.9,0.0,0.5)
#
for i in range(len(self.bar)):
rect(40+i*2, 4, 2, 4+self.bar[i]/4)
class MyScene(Scene):
def setup(self):
self.startTime = time.time()
self.xval = 0
self.bcount = 0
self.last = 0
#self.MySprite = SpriteNode(color = 'red', size = (32,32), parent = self)
#self.spot=ShapeNode(path=ui.Path.rect (100,0,50 ,30),parent=self)
#self.spot.position = (200,500)
self.results = {}
self.cell_size = CELL_SIZE
self.ball_amount = BALL_AMOUNT
# Spawn many balls
for i in range(self.ball_amount):
ball = make_ball(self)
self.results.setdefault((int(ball.x/self.cell_size), int(ball.y/self.cell_size)), []).append(ball)
# red ball
#aball = make_ball(self)
#aball.color = RED
#self.results.setdefault((int(ball.x/self.cell_size), int(ball.y/self.cell_size)), []).append(aball)
print("cell number = ",len(self.results.keys()))
# disp hashmap data
cnt = 0
ad = 0
ocl = 0
for aky in self.results.keys():
ad += len(aky)
cnt += 1
if len(aky) == 1:
ocl += 1
print("balls / cell = ", ad/cnt)
print("1 ball cell = ", ocl)
# scene view for stat
self.sv = SceneView(frame=(0, self.view.height-250,self.view.width,250),flex='WH')
self.s2 = MyScene2()
self.sv.scene = self.s2
self.view.add_subview(self.sv)
self.sv.border_color = (0.9,0.6,0.7)
self.sv.border_width = 3
#
button = ui.Button()
button.frame = (110,100,120,30)
button.action = self.button_tapped
button.title = 'spawn'
button.background_color = (0.6,0.3,0.5)
button.tint_color = ('white')
self.view.add_subview(button)
def button_tapped(self,sender):
#print('hhhgggg')
aball = make_ball(self)
aball.r= 3
aball.d = 60
aball.color = RED
self.results.setdefault((int(aball.x/self.cell_size), int(aball.y/self.cell_size)), []).append(aball)
def draw(self):
cell_size = CELL_SIZE
#bcount = self.bcount
background(0, 0, 150)
fill(110,10,10)
tint(200,5,5)
# Draw the balls
#self.cresults = {}
cresults = {}
for ky in self.results:
blist = self.results[ky]
blist2 = blist.copy()
#
for bl in blist:
blist2.remove(bl)
skip = True
if bl.nb[0] and bl.nb[1]:
k = (ky[0]+1,ky[1]+1)
if k in self.results:
blist2.extend(self.results[k])
skip = False
if skip and bl.nb[2] and bl.nb[1]:
k = (ky[0]-1,ky[1]+1)
if k in self.results:
blist2.extend(self.results[k])
skip = False
if skip and bl.nb[0] and bl.nb[2]:
k = (ky[0]+1,ky[1]-1)
if k in self.results:
blist2.extend(self.results[k])
skip = False
if skip and bl.nb[2] and bl.nb[2]:
k = (ky[0]-1,ky[1]-1)
if k in self.results:
blist2.extend(self.results[k])
#
if bl.nb[0]:
k = (ky[0]+1,ky[1])
if k in self.results:
blist2.extend(self.results[k])
if bl.nb[2]:
k = (ky[0]-1,ky[1])
if k in self.results:
blist2.extend(self.results[k])
if bl.nb[1]:
k = (ky[0],ky[1]+1)
if k in self.results:
blist2.extend(self.results[k])
if bl.nb[3]:
k = (ky[0],ky[1]-1)
if k in self.results:
blist2.extend(self.results[k])
#
for bl2 in blist2:
#if bl == bl2:
# continue
# Circle Intersect
sqDist = (bl2.x-bl.x)*(bl2.x-bl.x)+(bl2.y-bl.y)*(bl2.y-bl.y)
#if sqDist > 0 and sqDist <= ((BALL_SIZE)+(BALL_SIZE)):
#print(sqDist, math.sqrt(sqDist))
if sqDist > 0 and sqDist <= ((bl.r+bl2.r)*(bl.r+bl2.r)):
# detect collision
dist = math.sqrt(sqDist)
# vCollison
vc = ([bl2.x-bl.x, bl2.y-bl.y])
# vCollisonNorm
vcn = ([vc[0]/dist, vc[1]/dist])
# vRelativeVelocity
vrv = ([bl.vx-bl2.vx, bl.vy-bl2.vy])
speed = vrv[0]*vcn[0] + vrv[1]*vcn[1]
#print(speed)
if speed > 0:
bl.vx -= speed * vcn[0]
bl.vy -= speed * vcn[1]
bl2.vx += speed * vcn[0]
bl2.vy += speed * vcn[1]
#
#infection
if bl.color == RED and bl2.color != RED:
bl2.color = RED
self.bcount += 1
if bl2.color == RED and bl.color != RED:
bl.color = RED
self.bcount += 1
#
# draw shape
fill(bl.color)
ellipse(int(bl.x)-bl.r, int(bl.y)-bl.r, bl.r*2, bl.r*2)
bl.x += bl.vx
bl.y += bl.vy
# avoid line y
if bl.y > self.size.height - 200 + bl.r:
bl.y = self.size.height - 200
bl.vy = -bl.vy
if bl.y < bl.r :
bl.y = bl.r
bl.vy = -bl.vy
if bl.x < bl.r:
bl.x = bl.r
bl.vx = -bl.vx
if bl.x <= 0 :
bl.vx = 0.1
if bl.x > self.size.width - bl.r:
bl.x = self.size.width - bl.r
bl.vx = -bl.vx
#
#Bounce the ball if needed
if bl.y > self.size.height - 200 or bl.y < bl.r + 260:
bl.vy *= -1
if bl.x > self.size.width - bl.r or bl.x < bl.r:
bl.vx *= -1
# set key and get hash and append
cresults.setdefault((int(bl.x/cell_size), int(bl.y/cell_size)), []).append(bl)
#
# check neighbour with new key
px = int((bl.x+bl.r)/cell_size) > int(bl.x/cell_size)
mx = int((bl.x-bl.r)/cell_size) < int(bl.x/cell_size)
py = int((bl.y+bl.r)/cell_size) > int(bl.y/cell_size)
my = int((bl.y-bl.r)/cell_size) < int(bl.y/cell_size)
bl.nb =(px,py,mx,my)
self.results.clear()
self.results.update(cresults)
#print((selfself.spot.position = (100,800).now-dt.now).seconds))
timepassed = time.time()-self.startTime
#print(int(timepassed))
if int(timepassed) != self.last and self.ball_amount > self.bcount:
self.s2.add()
self.s2.val = (self.xval, self.bcount)
self.last = int(timepassed)
self.xval += 1
s=MyScene()
#Scene.run(s, show_fps=True)
main_view = ui.View()
scene_view = SceneView(frame=main_view.bounds, flex='WH')
main_view.title = 'OverShoot'
main_view.tint_color = (0,0,0)
main_view.add_subview(scene_view)
scene_view.scene = s
button = ui.Button()
button.frame = (10,100,main_view.width-20,30)
button.action = button_tapped
button.title = 'quit'
button.background_color = (0.6,0.3,0.5)
button.tint_color = ('white')
main_view.add_subview(button)
#scene_view.add_subview(sv)
#sv.background_color = (1,0,0)
main_view.present(hide_title_bar=False, animated=False)