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 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()