gravity.py (2723B)
1 import pygame 2 import random 3 import time 4 5 G = (6.674 * 10**-11) 6 7 class Body(): 8 def __init__(self, x, y, mass, radius, vel_x, vel_y): 9 self.radius = radius 10 self.mass = mass 11 self.x = x 12 self.y = y 13 self.vel_x = vel_x 14 self.vel_y = vel_y 15 self.last_timestep = time.time() 16 17 def distance(self, body): 18 return ((((self.x - body.x) ** 2) + ((self.y - body.y) ** 2)) ** .5) 19 20 def render(self, display): 21 pygame.draw.circle(display, white, (self.x, self.y), self.radius) 22 23 def get_delta_time(self): 24 current = time.time() 25 result = current - self.last_timestep 26 self.last_timestep = current 27 return result 28 29 def update_position(self): 30 self.x += self.vel_x * self.delta_time 31 self.y += self.vel_y * self.delta_time 32 33 def update_velocity(self, bodies): 34 self.delta_time = self.get_delta_time() 35 36 m_1 = self.mass 37 38 for body in bodies: 39 if body == self: 40 continue 41 r = self.distance(body) 42 43 if r < self.radius + body.radius: 44 if body.mass > self.mass: 45 return False 46 47 if r < 20: 48 continue 49 50 m_2 = body.mass 51 52 self.vel_x += ((G * m_2 / r ** 2) * (body.x - self.x) / r) * self.delta_time 53 self.vel_y += ((G * m_2 / r ** 2) * (body.y - self.y) / r) * self.delta_time 54 55 if (self.x > 1200 and self.vel_x > 0) or (self.x < 0 and self.vel_x < 0): 56 self.vel_x *= -1 57 if (self.y > 1000 and self.vel_y > 0) or (self.y < 0 and self.vel_y < 0): 58 self.vel_y *= -1 59 60 return True 61 62 63 64 65 66 67 pygame.init() 68 69 display = pygame.display.set_mode((1200, 1000)) 70 71 white = (255, 255, 255) 72 red = (255, 0, 0) 73 black = (0, 0, 0) 74 75 bodies = [] 76 77 for i in range(0,200): 78 # would be better to calculate radius, volume, and then use this to calculate mass 79 size_rnd = random.random() * .25 80 size = (4.2 * 10**15) * size_rnd 81 radius = size_rnd * 10 82 bodies.append(Body(random.randint(0,1200), random.randint(0,1000), size, radius, random.randint(-100,100), random.randint(-100,100))) 83 84 size = (4.2 * 10**15) * 10 85 radius = 100 86 87 bodies.append(Body(600, 500, size, radius, 0,0)) 88 89 while True: 90 for event in pygame.event.get(): 91 if event.type == pygame.QUIT: 92 pygame.quit() 93 quit() 94 95 display.fill(black) 96 97 to_remove = [] 98 99 for body in bodies: 100 if not body.update_velocity(bodies): 101 to_remove.append(body) 102 103 for body in to_remove: 104 bodies.remove(body) 105 106 for body in bodies: 107 body.update_position() 108 body.render(display) 109 110 111 pygame.display.update()