Making a gravity simulator using Python

Photo by NASA on Unsplash

“Hello world!” Today I am going to show you how to make a gravity simulator using Python. You should be able to define bodies with a specific mass and initial velocity. Then let the programmed force of gravity take over!

Modules

We will be using pygame as a canvas to show the movement of bodies and numpy to calculate all the vectors in the simulation. You can install both of them easily enough.

pip install numpy
pip install pygame

Use conda instead of pip if you use anaconda. I just setup a virtual environment and installed them there using pip. Also I’ll be using PyCharm as my IDE, although, it shouldn’t matter much.

Let’s begin

This is the file structure I used:

Gravity_sim
|
|-- src
    |-- classes
        |-- body.py
        |-- physics_engine.py
        |-- simulation.py
    |-- main.py
    |-- README.md
|-- venv

Let’s start with thebody.py file. Here we will define the characteristics of a body in our simulation. We will make it a class:

import numpy as np
import pygame
# constants
TIME_DELAY = 0.0005
class Body:
    def __init__(self, position_array, mass, color, radius=10):
        self.velocity = np.array([[0, 0, 0]])
        self.force = np.array([[0, 0, 0]])
        self.mass = mass
        self.position = position_array
        self.radius = radius
        self.thickness = self.radius * 2
        self.color = color

The time delay is set to update the sim in such a way so that it is visible to us. Next we can define several functions. The first thing we can define is a draw function that will draw the circle representing our body:

def draw(self, surface):
    pygame.draw.circle(surface, self.color, (self.position[0][0],
                       self.position[0][1]), self.radius,
                       self.thickness)

The surface argument is just the surface we have defined in pygame. Since we might have multiple simulations, I decided to add a surface argument. As you see, we have defined a velocity array, so we need to add a function that can update the velocity at the start of the simulation. This is very trivial:

def add_velocity(self, velocity_array):
    self.velocity = self.velocity + velocity_array

The same thing goes for a force. However, this will be used every time we update the simulation. This means every 0.0005 (TIME_DELAY) seconds:

def add_force(self, force_array):
    self.force = self.force + force_array

Now that we have a velocity and a force, we need to update the position of the body. Since we are going to use newton’s laws here, we can just use the laws of motion. You all should remember this from high school, but these are the laws I have used:

v = u + a*t

d = s * t

where v is the final velocity, u is the initial velocity, a is the acceleration, t is the time, and d is the displacement.

In code:

def move(self):
    self.velocity = self.velocity + ((self.force / self.mass) *
                                     TIME_DELAY)
    self.position = self.position + self.velocity * TIME_DELAY

Now comes the hard part. Defining the “physics engine” that would run our sim. In the physics_engine.py file, first import numpy. Then we will define 2 helper functions: one for calculating the magnitude of a vector and another for calculating the unit vector.

# shape of all arrays will be 1, 3.
# example: np.array([[x, y, z]])
def _magnitude(arr):
    return np.sqrt(pow(arr[0], 2) + pow(arr[1], 2) + pow(arr[2], 2))
def _unit_vector(arr):
    mag = _magnitude(arr)
    return np.array([
        [
            arr[0] / mag,
            arr[1] / mag,
            arr[2] / mag,
        ]
    ])

Now we can define a physics engine class:

class PhysicsEngine:
    def __init__(self):
        self.body_pos_array = np.array([]).reshape((0, 3))
        self.body_list = None

We need a way for the physics engine to know the bodies in the simulation, so we can make a function that does just that:

def define_bodies(self, body_list):
    self.body_list = [np.array([i, body]) for i, body in
                      enumerate(body_list)]

I decided to use the enumerate function and assign an index to every body so that we can keep track of the body being analysed. This will come in handy later. Now comes the most difficult and gruesome part of the project, computing the force vectors:

def compute_force_vectors(self):
    distance_list = []
    force_list = []
    net_force = []

The distance list will hold the distance of every other body from a body in the simulation. The force list will do the same, just the force with respect to every other body in the simulation. The net force list will hold the final force vector for every body.

Now the difficult part is updating the force list and displacement list. I’ll post the code and then explain:

First we consider a single body. Then we find the displacement vector from that body to every other body in the simulation. We append the array to the distance_list. We multiplied by np.array([[-1, -1, 1]]) as pygame’s coordinate system is slightly different from the cartesian coordinate system. The origin is at the top-left corner and positive y is in the downward direction.

Now let’s calculate the force vector:

If you see the formula, I have multiplied the magnitude by 0.25 so that the effects of the attractive force is magnified so that the effects are visible to the observer.

The for loop is relatively simple. First we take a primary body and calculate the force it feels by every other body. However, we do not want a np.inf in our calculations when the secondary body and the primary body are the same or if two bodies are overlapping, hence we put a if statement that ignores such circumstances and just puts the force as zero. Then I append the array to a list.

for obj in force_list:
    net_force.append(obj.sum(axis=0))


return net_force

After that I just sum the forces in each direction for every body and return the list.

Let’s start making the Simulation class. Let’s import pygame, time and the physics engine class.

import pygame
import time
from src.classes.physics_engine import PhysicsEngine
class Simulation:
    physics_engine = PhysicsEngine()
    
    def __init__(self):
        self.run, self.space, self.bodies = None, None, None

self.run will be used to govern the loop for pygame. “self.bodies” will hold the bodies in the simulation and self.space is just the pygame display. We will define all of that later so I have just set it to None for now. Let’s make a function to initialise the pygame display and to define all the bodies in the simulation class:

def initialise_environment(self, body_list):
    self.bodies = body_list
    space_plane_size = (1000, 800)  # width, height of canvas
    self.run = True

    # setting up pygame
    pygame.init()
    pygame.display.set_caption("Orbit simulator")
    self.space = pygame.display.set_mode(space_plane_size)

    # setting up physics engine
    self.physics_engine.define_bodies(body_list)

This code is self explanatory. I just set the size of the display and pass the body list to the physics engine. Now we define the function to run the sim:

def show_environment(self):
    while self.run:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.run = False

        self.space.fill((0, 0, 0))

        net_force = self.physics_engine.compute_force_vectors()
        for i, body in enumerate(self.bodies):
            body.force = net_force[i]

        for body in self.bodies:
            body.draw(self.space)

        for body in self.bodies:
            body.move()
        time.sleep(0.0005)
        pygame.display.update()

Even this code is very self explanatory. We setup the condition to quit. Then we fill the display with black. Then we use the draw function from all the bodies. Then we use the move function from all the bodies. Then we wait for 0.0005 seconds, then we update the display. That’s all. Now we just run the code in main.py:

body1 = Body(np.array([[500, 300, 0]]), 6 * pow(10, 15), (255, 255,
                                                          255))
body1.add_velocity(np.array([[40, 0, 0]]))

body2 = Body(np.array([[600, 200, 0]]), 6 * pow(10, 15), (0, 0,
                                                          255))
body2.add_velocity(np.array([[-40, 0, 0]]))

body3 = Body(np.array([[300, 500, 0]]), 6 * pow(10, 15), (0, 255,
                                                          0))
body3.add_velocity(np.array([[50, 0, 0]]))

sim = Simulation()
sim.initialise_environment([body1, body2, body3])
sim.show_environment()

Conclusion

Well that’s it. If you want to see the full code, you can access it here. I hope you enjoyed reading this blog. If you want more blogs related to programming and space news, follow me here.

Thanks for reading!