Gravity simulation using Numpy and Pygame

Question

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

Answer 1

• The positions of the bodies are modified in-place, while the loop is running. It means that in your universe the action is not equal to reaction. Your simulation will exhibit some nonrealistic behaviors, like drift of the center of the mass.

• Euler integration method is not the most accurate. Consider Runge-Kutta. At least, monitor the motion invariants (total energy, momentum, and angular momentum of the ensemble), and do corrective actions when they start diverging.

• 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.