From 9f45d0a2db12ca31a5a9d7bcab933a9d797c50de Mon Sep 17 00:00:00 2001 From: Jim Caten Date: Thu, 6 Jun 2019 03:22:28 -0500 Subject: [PATCH 1/2] Add features and some code clean up Adds collisions, which adds mass and velocity Adds feature where the "screen" will follow the center of mass for all Bodies Adds feature where a line is drawn to virtual 0,0 to the screen's center Moves Body class to it's own file Checks __name__ is == to __main__, just in case multiprocessing is used later --- gravity_simulation.py | 216 ++++++++++++++++++++++++------------------ 1 file changed, 122 insertions(+), 94 deletions(-) diff --git a/gravity_simulation.py b/gravity_simulation.py index 8f89ca3..9ce5e22 100644 --- a/gravity_simulation.py +++ b/gravity_simulation.py @@ -1,100 +1,128 @@ # Very simple implementation of simulation of gravity on bodies in 2D. Does not handle the case when 2 or more # bodies collide with each other import random -import math import sys +import math import pygame -class Body: - def __init__(self, pos, a, v, m): - self.pos = pos # pos is a list of x and y position of that body in pixels eg : [500,600] - self.a = a # a is a list of x and y components of accelaration of that body in pixel units - self.v = v # b is a list of x and y components of velocity of that body in pixel units - self.m = m # m is the mass of that object - - -def calculate_forces(pos_a, pos_b, m_a, m_b): - x_diff = pos_b[0] - pos_a[0] - y_diff = pos_b[1] - pos_a[1] - hypotenuse = math.sqrt(((x_diff) ** 2 + (y_diff) ** 2)) - sin = x_diff / hypotenuse - cos = y_diff / hypotenuse - f = G * m_a * m_b / hypotenuse ** 2 - fx = f * sin - fy = f * cos - - return fx, fy - - -G = 6.67408e-11 * 100_000_000 # Otherwise the bodies would not move given the small value of gravitational constant -NUM_OF_BODIES = 10 -WIDTH = 900 -HEIGHT = 800 -WHITE = (255, 255, 255) -BLACK = (0, 0, 0) -BLUE = (109, 196, 255) - -bodies = [] -for i in range(NUM_OF_BODIES): - px = random.randint(10, WIDTH - 10) - py = random.randint(10, HEIGHT - 10) - m = random.randint(1, 25) - bodies.append(Body([px, py], [0, 0], [0, 0], m)) - -# Some predefined bodies for the purpose of testing -# bodies.append(Body([500,500],[0,0],[0,0],20)) -# bodies.append(Body([510,503],[0,0],[0,0],7)) -# bodies.append(Body([400,400],[0,0],[0,0],14)) -# bodies.append(Body([10,600],[0,0],[0,0],9)) -# bodies.append(Body([250,198],[0,0],[0,0],18)) -# bodies.append(Body([340,700],[0,0],[0,0],24)) - -pygame.init() -size = WIDTH, HEIGHT -screen = pygame.display.set_mode(size) - -font = pygame.font.SysFont('Arial', 16) -text = font.render('0', True, BLUE) -textRect = text.get_rect() -while True: - screen.fill(BLACK) - for event in pygame.event.get(): - if event.type == pygame.QUIT: sys.exit() - - for body_a in bodies: - pos_a = body_a.pos - m_a = body_a.m - fx_total = 0 - fy_total = 0 - - for body_b in bodies: - if body_b.pos == pos_a: - continue - fx, fy = calculate_forces(pos_a, body_b.pos, m_a, body_b.m) - fx_total += fx - fy_total += fy - - body_a_acceleration = body_a.a - - body_a_acceleration[0] = fx_total / m_a - body_a_acceleration[1] = fy_total / m_a - - body_a.v[0] = body_a.v[0] + body_a_acceleration[0] - body_a.v[1] = body_a.v[1] + body_a_acceleration[1] - - pos_a[0] = pos_a[0] + body_a.v[0] - pos_a[1] = pos_a[1] + body_a.v[1] - - mass_text = 'M={0}'.format(m_a) - # force_text = 'F=({0},{1})'.format(fx_total.__round__(3), fy_total.__round__(3)) - # velocity_text = 'V=({},{})'.format(body_a.v[0].__round__(3),body_a.v[1].__round__(3)) - # text_str = mass_text + ' ' + force_text + ' ' + velocity_text - text_str = mass_text - - text = font.render(text_str, True, BLUE) - textRect.center = (pos_a[0] + m_a + 10, pos_a[1] + m_a + 10) - - screen.blit(text, textRect) - - pygame.draw.rect(screen, (255, 255, 255), pygame.Rect(pos_a[0], pos_a[1], m_a, m_a)) - pygame.display.flip() +from Planet import Body + +if __name__ == "__main__": + NUM_OF_BODIES = 20 + WIDTH = 900 + HEIGHT = 800 + WHITE = (255, 255, 255) + BLACK = (0, 0, 0) + BLUE = (109, 196, 255) + + bodies = [] + for i in range(NUM_OF_BODIES): + px = random.randint(10, WIDTH - 10) + py = random.randint(10, HEIGHT - 10) + m = random.randint(1, 20) + bodies.append(Body([px, py], [0, 0], m, i)) + + pygame.init() + size = WIDTH, HEIGHT + screen = pygame.display.set_mode(size) + + font = pygame.font.SysFont('Arial', 16) + text = font.render('0', True, BLUE) + textRect = text.get_rect() + velocity_diff = [0, 0] + while True: + screen.fill(BLACK) + for event in pygame.event.get(): + if event.type == pygame.QUIT: + sys.exit() + + # Don't process bodies that where hit/removed + bodies = [body for body in bodies if not body.hit] + + # Get Bodies and find the center of mass of all + x = [p.pos[0] for p in bodies] + y = [p.pos[1] for p in bodies] + centroid = (sum(x) / len(bodies), sum(y) / len(bodies)) + lx = bodies[0].pos[0] + ly = bodies[0].pos[1] + for body in bodies: + body.pos[0] = body.pos[0] - centroid[0] + WIDTH / 2 + body.pos[1] = body.pos[1] - centroid[1] + HEIGHT / 2 + + # Draw circles and line for reference point of origin + velocity_diff[0] = lx - centroid[0] + velocity_diff[1] = ly - centroid[1] + # Origin Point + textRect.center = (velocity_diff[0] + 10, velocity_diff[1] + 10) + screen.blit(font.render("{0},{1}".format(0, 0), True, BLUE), textRect) + pygame.draw.circle(screen, (255, 255, 127), [int(velocity_diff[0]), int(velocity_diff[1])], 3, 1) + # Center of all objects + textRect.center = (WIDTH/2 + 10, HEIGHT/2 + 10) + screen.blit(font.render("{0},{1}".format( + int(math.floor(velocity_diff[0])), + int(math.floor(velocity_diff[1]))), True, BLUE), textRect) + # Draw line from origin to center of all objects + pygame.draw.line(screen, (255, 255, 0), (WIDTH/2, HEIGHT/2), velocity_diff) + pygame.draw.circle(screen, (255, 255, 127), [int(WIDTH/2), int(HEIGHT/2)], 3, 1) + + for body_a in bodies: + # Remove invulnerability flag + body_a.invuln = False + + f_totals = body_a.n_body(bodies, 0, 0) + + body_a.velocity[0] = body_a.velocity[0] + f_totals[0] / body_a.mass + body_a.velocity[1] = body_a.velocity[1] + f_totals[1] / body_a.mass + + body_a.pos[0] = body_a.pos[0] + body_a.velocity[0] + body_a.pos[1] = body_a.pos[1] + body_a.velocity[1] + + mass_text = 'M={0}'.format(body_a.mass) + # force_text = 'F=({0},{1})'.format(fx_total.__round__(3), fy_total.__round__(3)) + # velocity_text = 'V=({},{})'.format(body_a.v[0].__round__(3),body_a.v[1].__round__(3)) + # text_str = mass_text + ' ' + force_text + ' ' + velocity_text + text_str = mass_text + + text = font.render(text_str, True, BLUE) + textRect.center = ( + body_a.pos[0] + body_a.size + 10, + body_a.pos[1] + body_a.size + 10) + + screen.blit(text, textRect) + + pygame.draw.circle( + screen, + (255, 255, 255), + [int(body_a.pos[0]), int(body_a.pos[1])], int(body_a.size)) + + # Get a list of bodies, except for body_a + next_bodies = [body for body in bodies if not body_a.id == body.id] + for body in next_bodies: + # if body is invulnerable then continue on to the next body + if body.invuln: + continue + # Find the distance to body_a + distance = math.sqrt( + ((body_a.pos[0] - body.pos[0]) * (body_a.pos[0] - body.pos[0])) + + ((body_a.pos[1] - body.pos[1]) * (body_a.pos[1] - body.pos[1]))) + # If bodied touch then "remove" one by setting flag and making the other invulnerable + if distance < int(body_a.size) + int(body.size): + if body_a.mass >= body.mass: + body_a.mass += body.mass + body_a.update_size() + body_a.velocity[0] = (body_a.mass * body_a.velocity[0] + body.mass * body.velocity[0] + ) / (body_a.mass + body.mass) + body_a.velocity[1] = (body_a.mass * body_a.velocity[1] + body.mass * body.velocity[1] + ) / (body_a.mass + body.mass) + body.hit = True + body_a.invuln = True + else: + body.mass += body_a.mass + body.update_size() + body.velocity[0] = (body_a.mass * body_a.velocity[0] + body.mass * body.velocity[0] + ) / (body_a.mass + body.mass) + body.velocity[1] = (body_a.mass * body_a.velocity[1] + body.mass * body.velocity[1] + ) / (body_a.mass + body.mass) + body_a.hit = True + body.invuln = True + pygame.display.flip() From 35976a8814c9485e22801ab337208b8e63587956 Mon Sep 17 00:00:00 2001 From: Jim Caten Date: Thu, 6 Jun 2019 03:27:13 -0500 Subject: [PATCH 2/2] Create Planet.py Moves Body class to this file Removes acceleration as it isn't needed Adds id for unique identifying bodies Adds hit and invuln for collision detection Moves n-body logic to Body class Moves calculate_forces to Body class Creates update_size Moved G constant to here as it isn't needed in main script file --- Planet.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 Planet.py diff --git a/Planet.py b/Planet.py new file mode 100644 index 0000000..846fda1 --- /dev/null +++ b/Planet.py @@ -0,0 +1,35 @@ +import math + +G = 6.67408e-11 * 100_000_000 + + +class Body: + def __init__(self, pos, velocity, mass, id_num): + self.pos = pos # pos is a list of x and y position of that body in pixels eg : [500,600] + self.velocity = velocity # b is a list of x and y components of velocity of that body in pixel units + self.mass = mass # m is the mass of that object + self.size = math.log2(self.mass) # size is used for only visualisation of the body + self.id = id_num # id is the unique identifier for this body + self.hit = False # tell us if this object was hit by another + self.invuln = False # a temporary flag for showing what is invulnerable + + def update_size(self): + self.size = math.log2(self.mass) + + def n_body(self, other_bodies, n_fx_total, n_fy_total): + for body_b in other_bodies: + if body_b.pos != self.pos: + fx, fy = self.calculate_forces(self.pos, body_b.pos, self.mass, body_b.mass) + n_fx_total += fx + n_fy_total += fy + return [n_fx_total, n_fy_total] + + @staticmethod + def calculate_forces(c_pos_a: list, c_pos_b: list, c_m_a: float, c_m_b: float): + x_diff = c_pos_b[0] - c_pos_a[0] + y_diff = c_pos_b[1] - c_pos_a[1] + f = G * c_m_a * c_m_b / math.sqrt(x_diff ** 2 + y_diff ** 2) + angle = math.atan2(y_diff, x_diff) + cfx = f * math.cos(angle) + cfy = f * math.sin(angle) + return cfx, cfy