Switched core calculation to numpy

gravity_simulation.py

@@ -4,42 +4,24 @@
 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
 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))
+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()