# Gravity simulation using Numpy and Pygame

This is a simple gravity simulator coded in Python 3.7 using Numpy and Pygame. I was wondering if it can be optimized further. Initially I had coded it using pure Python lists, using nested loops to calculate forces between each body, but then the good people of Reddit suggested me to use Numpy, Cython or Numba to improve the runtime of the code. Some even suggested a few highly optimized algorithms which are used to improve the time complexity of N-Body simulations like the Barnes Hut algorithm. For v2, I decided to use Numpy and have kept Barnes-Hut for the scope of future. Here is the implementation in numpy :

Can it be optimized further?

import sys
import time

import numpy as np
import pygame

G = 6.67408e-11 * 100_000_000  # Otherwise the bodies would not move given the small value of gravitational constant
NUM_OF_BODIES = 4000
WIDTH = 1200
HEIGHT = 800
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (109, 196, 255)

vx = np.zeros((NUM_OF_BODIES,),dtype=np.float)
vy = np.zeros((NUM_OF_BODIES,),dtype=np.float)

px = np.random.uniform(low=10, high=WIDTH-10,size=NUM_OF_BODIES)
py = np.random.uniform(low=10, high=HEIGHT-10,size=NUM_OF_BODIES)

m = np.random.randint(1,25,size=NUM_OF_BODIES)

fx = np.zeros((NUM_OF_BODIES,),dtype=float)
fy = np.zeros((NUM_OF_BODIES,),dtype=float)

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

in_t = time.time()
for i in range(0,NUM_OF_BODIES):
xdiff = (px - px[i])
ydiff = (py - py[i])

distance = np.sqrt(xdiff ** 2 + ydiff ** 2)

f = G * m[i] * np.divide(m,distance ** 2)

sin = np.divide(ydiff,distance)
cos = np.divide(xdiff,distance)

fx_total = np.nansum(np.multiply(f, cos))
fy_total = np.nansum(np.multiply(f,sin))

vx[i] = vx[i] + fx_total / m[i]
vy[i] = vy[i] + fy_total / m[i]

px[i] = px[i] + vx[i]
py[i] = py[i] + vy[i]

pygame.draw.rect(screen, (255, 255, 255), pygame.Rect(px[i], py[i], m[i],m[i]))
print(time.time() - in_t)
pygame.display.flip()


• After f = m[i] * ... and f_total / m[i] the mass cancels. Multiplication and division could be safely omitted.
• Along the same line, scale the masses by G once, before the simulation begins.