From 14beabde9f5adc6a5b28b6e43d5c4381cf81baa0 Mon Sep 17 00:00:00 2001 From: Filip Zastavnik Date: Tue, 4 Jun 2019 16:41:19 +0200 Subject: [PATCH 1/2] Switched core calculation to numpy --- gravity_simulation.py | 91 ++++++++++++++++++------------------------- 1 file changed, 38 insertions(+), 53 deletions(-) diff --git a/gravity_simulation.py b/gravity_simulation.py index 8f89ca3..bddfb7d 100644 --- a/gravity_simulation.py +++ b/gravity_simulation.py @@ -4,29 +4,12 @@ import math import sys import pygame +import numpy as np -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 +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 @@ -34,12 +17,11 @@ def calculate_forces(pos_a, pos_b, m_a, m_b): 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)) +x = np.random.randint(10, WIDTH - 10, NUM_OF_BODIES).astype(np.float32) +y = np.random.randint(10, HEIGHT - 10, NUM_OF_BODIES).astype(np.float32) +v_x = np.zeros(NUM_OF_BODIES, dtype=np.float32) +v_y = np.zeros(NUM_OF_BODIES, dtype=np.float32) +m = np.random.randint(1, 25, NUM_OF_BODIES).astype(np.float32) # Some predefined bodies for the purpose of testing # bodies.append(Body([500,500],[0,0],[0,0],20)) @@ -53,48 +35,51 @@ def calculate_forces(pos_a, pos_b, m_a, m_b): size = WIDTH, HEIGHT screen = pygame.display.set_mode(size) -font = pygame.font.SysFont('Arial', 16) -text = font.render('0', True, BLUE) +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 + if event.type == pygame.QUIT: + sys.exit() - 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 + fx_total = np.zeros(NUM_OF_BODIES) + fy_total = np.zeros(NUM_OF_BODIES) - body_a_acceleration = body_a.a + x_diff = np.repeat(x[:, np.newaxis], NUM_OF_BODIES, axis=1) - np.repeat( + x[np.newaxis, :], NUM_OF_BODIES, axis=0 + ) + y_diff = np.repeat(y[:, np.newaxis], NUM_OF_BODIES, axis=1) - np.repeat( + y[np.newaxis, :], NUM_OF_BODIES, axis=0 + ) + hypotenuse = np.sqrt(x_diff ** 2 + y_diff ** 2) + hypotenuse[np.diag_indices(NUM_OF_BODIES)] = np.inf - body_a_acceleration[0] = fx_total / m_a - body_a_acceleration[1] = fy_total / m_a + sin = x_diff / hypotenuse + cos = y_diff / hypotenuse + f = G * m[np.newaxis, :] * m[:, np.newaxis] / hypotenuse ** 2 + fx = f * sin + fy = f * cos - body_a.v[0] = body_a.v[0] + body_a_acceleration[0] - body_a.v[1] = body_a.v[1] + body_a_acceleration[1] + v_x += np.sum(fx, axis=1) / m + v_y += np.sum(fy, axis=1) / m - pos_a[0] = pos_a[0] + body_a.v[0] - pos_a[1] = pos_a[1] + body_a.v[1] + x += v_x + y += v_y - mass_text = 'M={0}'.format(m_a) + for idx in range(NUM_OF_BODIES): + mass_text = "M={0}".format(m[idx]) # 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) + textRect.center = (x[idx] + m[idx] + 10, y[idx] + m[idx] + 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.draw.rect( + screen, (255, 255, 255), pygame.Rect(x[idx], y[idx], m[idx], m[idx]) + ) pygame.display.flip() From 37dfcfd408d7ff5e735d1c21b300be18e7599870 Mon Sep 17 00:00:00 2001 From: Filip Zastavnik Date: Tue, 4 Jun 2019 16:52:44 +0200 Subject: [PATCH 2/2] Ups, wrong gravity direction... --- gravity_simulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gravity_simulation.py b/gravity_simulation.py index bddfb7d..170d9f8 100644 --- a/gravity_simulation.py +++ b/gravity_simulation.py @@ -58,7 +58,7 @@ sin = x_diff / hypotenuse cos = y_diff / hypotenuse - f = G * m[np.newaxis, :] * m[:, np.newaxis] / hypotenuse ** 2 + f = -G * m[np.newaxis, :] * m[:, np.newaxis] / hypotenuse ** 2 fx = f * sin fy = f * cos