Planetary simulation in python

Question

I have created a program in python that calculates forces between bodies (i.e earth, moon and a hypothetical moon) and make them move according to the changes in velocity and forces. This is the code for the program:

from math import sin,cos,sqrt,atan2,pi
import pygame
pygame.init()

class Planet:
    dt = 1/100
    G = 6.67428e-11 #G constant
    scale = 1/(1409466.667) #1 m = 1/1409466.667 pixlar
    def __init__(self,x=0,y=0,radius=0,color=(0,0,0),mass=0,vx=0,vy=0):
        self.x = x #x-coordinate pygame-window
        self.y = y #y-coordinate pygame-window
        self.radius = radius
        self.color = color
        self.mass = mass
        self.vx = vx #velocity in the x axis
        self.vy = vy #velocity in the y axis
        
    def draw(self,screen):
        pygame.draw.circle(screen, self.color, (self.x, self.y), self.radius)
    
    def orbit(self,trace):
        pygame.draw.rect(trace, self.color, (self.x, self.y, 2, 2))
        
    def update_vel(self,Fnx,Fny):
        ax = Fnx/self.mass #Calculates acceleration in x- and y-axis for body 1.
        ay = Fny/self.mass
        self.vx -= ((ax * Planet.dt)/Planet.scale)
        self.vy -= ((ay * Planet.dt)/Planet.scale)
        self.update_pos()
     
    def update_pos(self):
        self.x += ((self.vx * Planet.dt)) #changes position considering each body's velocity.
        self.y += ((self.vy * Planet.dt))
        
    def move(self,body):
        dx = (self.x - body.x) #Calculates difference in x- and y-axis between the bodies
        dy = (self.y - body.y)
        r = (sqrt((dy**2)+(dx**2))) #Calculates the distance between the bodies
        angle = atan2(dy, dx) #Calculates the angle between the bodies with atan2!
        if r < self.radius: #Checks if the distance between the bodies is less than the radius of the bodies. Uses then Gauss gravitational law to calculate force.
            F = 4/3 * pi * r
            Fx = cos(angle) * F
            Fy = sin(angle) * F
        else:  
            F = (Planet.G*self.mass*body.mass)/((r/Planet.scale)**2) #Newtons gravitational formula.
            Fx = cos(angle) * F
            Fy = sin(angle) * F
        return Fx,Fy

def motion():
    for i in range(0,len(bodies)):
        Fnx = 0 #net force
        Fny = 0
        for j in range(0,len(bodies)):
            if bodies[i] != bodies[j]:
                Fnx += (bodies[i].move(bodies[j]))[0]
                Fny += (bodies[i].move(bodies[j]))[1]
            elif bodies[i] == bodies[j]:
                continue
        bodies[i].update_vel(Fnx,Fny)
        bodies[i].draw(screen)
        bodies[i].orbit(trace)
        Fnx,Fny=0,0 

screen = pygame.display.set_mode([900,650]) #width - height
trace = pygame.Surface((900, 650))
pygame.display.set_caption("Moon simulation")
FPS = 60 #how quickly/frames per second our game should update. Change?

earth = Planet(450,325,30,(0,0,255),5.97219*10**(24),-24.947719394204714/2) #450= xpos,325=ypos,30=radius
luna = Planet(450,(575/11),10,(128,128,128),7.349*10**(22),1023)
moon = Planet() #the second moon
bodies = [earth,luna]

running = True
clock = pygame.time.Clock()

while running: #if user clicks close window
    clock.tick(FPS)    
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
            
    screen.fill((0,0,0))
    pygame.Surface.blit(screen, trace, (0, 0))
    motion()

    pygame.display.flip() #update? flip? 

pygame.quit()

The code works and I get the moon to orbit around earth. However since I am an amateur in such an area I wonder if it can be improved on. Do the calculations look alright? Can I make some improvements? I am trying to make my simulation realistic and thus I have a scaling factor that I use to convert pixels to meters and vice versa. Here for example: Is it correct to divide by the scaling factor or should I multiply it (so I don't mix pixels with meters):

    def update_vel(self,Fnx,Fny):
        ax = Fnx/self.mass #Calculates acceleration in x- and y-axis for body 1.
        ay = Fny/self.mass
        self.vx -= ((ax * Planet.dt)/Planet.scale)
        self.vy -= ((ay * Planet.dt)/Planet.scale)
        self.update_pos()
     

Thankful for any suggestions or thoughts on my code and/or the physics involved in it!

Answer 1

Use a vector type to store positions, velocities and forces

Instead of writing out all the calculations for the x and y direction separately, consider using a library that provides you with vector types, like for example PyGLM. This allows you to write simpler expressions, and there will be less chance of accidental mistakes. Using glm.vec2, you can write code like:

def __init__(self, pos=glm.vec2(0, 0), …, vel=glm.vec2(0, 0)):
    self.pos = pos
    …
    self.vel = vel
        
def update_vel(self, Fn):
    a = Fnx / self.mass
    self.vel -= (a * Planet.dt) / Planet.scale
    self.update_pos()

def update_pos(self):
    self.pos += self.vel * Planet.dt

Use the velocity Verlet method

You are using a rather straightforward but naive way to calculate the next position given the current velocity and acceleration, which is also known as the semi-implicit Euler method. While this is already better than the forward Euler method, it's just a first-order approximation, and if the time step is too large it might result in an unstable system.

There are various ways to improve on this. A very commonly used method to discretely integrate the equations of motion is the velocity Verlet method; it is a second order approximation, and results in a much more stable system, but it's still very simple to implement and very efficient.

Avoid using trigonometric functions

You can calculate the force vectors without using atan(), cos() and sin(). Consider that cos(angle) is just dx / r and sin(angle) is dy / r.

Scaling

It would indeed be very nice if you put in all the values in SI units like kilograms, meters, seconds and so on. You can find distances between the Sun, planets and moons on Wikipedia. Try to fill in the values for our solar system. You'll notice that the time step should then be quite large, for example if you set it to 1 hour, then it will take about 24 * 365 = 8760 steps for the Earth to revolve around the Sun once. That's about 2 minutes if your simulation runs at 60 frames per second.

You also have to convert from meters to pixels in such a way that the distance to the outermost planet is slightly less than half the width and height of the window, so everything fits in it. This scaling factor should be global, the same value should be used for all objects you want to show. It should be applied inside draw(), not while updating the positions and velocities: keep all the physics in SI units, and only convert to pixels when it's time to draw.

Answer 2

TL;DR:

I created a Github project with everything I suggested and more (like taking all PyGame-related logic away from the Planet class, leaving it only with planet-related logic). Take a look at the commit history:

https://github.com/kirisakow/planetary_simulation/commits/main

And here's what I previously said in the comments:

Syntax

• for loops: You may want to read about the so-called ”Pythonic way“ to write a loop.

• Too many brackets (don't be afraid to lose them), not enough spaces:

  • 1/(1409466.667) should be 1 / 1409466.667,
  • 5.97219*10**(24) should be 5.97219 * 10**24 or 5.97219E+24,
  • 7.349*10**(22) should be 7.349 * 10**22 or 7.349E+22,
  • (575/11) should be 575 / 11

• Planet() constructor takes 7 arguments, all positional; that's somewhat messy. I can see you added comments for earth to help remember which argument is which. It would be simpler for the arguments to go together with their parameter names, just like in the constructor's definition.

  • Bonus tip: you can enforce keyword-only arguments by putting an asterisk (*) argument at the beginning of the constructor's definition, right after self: def __init__(self, *, x=...).
  • Learn more about parameters from official documentation, or this paper about why everyone should use keyword-only arguments more often.

Object-oriented logic

• Function motion() operates on data which does not belong to it. Using so-called "global variables" is a bad practice, especially as the code grows. Therefore,
  • either bodies should become an argument of that function: motion(bodies)
  • or the function itself should become a method in a separate class, with bodies as an attribute:

motion = Motion(bodies)
motion.update()

Answer 3

Separation of Concerns

Your Planet class does a lot. It concerns itself with the physics of the planet (computing gravitational forces, updating velocity and position), the specific details of the simulation (dt), and details of how a planet is rendered (draw(), orbit(), scale). Consider separating those concerns. Make your Planet class concern itself with just the physics, and create a new class to handle displaying things.

This also goes to your question about scaling. Scaling is important only when you are trying to display something on the screen, so should only come into play there. But since your positions and velocities are in terms of pixels, while your forces and accelerations are in terms of meters, you have to convert between pixels and meters in your calculations. Separate it out, and the physics works in meters, and is converted to pixels when drawing -- and that's the only place they need to be converted.

class Renderer:
    scale = 1/(1409466.667) # 1/1409466.667 pixels/meter
    screen = pygame.display.set_mode([900,650]) #width - height
    trace = pygame.Surface((900, 650))
    pygame.display.set_caption("Moon simulation")

    def drawBody(body):
        x_in_pixels = body.x * scale
        y_in_pixels = body.y * scale
        pygame.draw.circle(...)
        pygame.draw.rect(...)

Rename Misnamed Methods

I don't know if you noticed, but your move() method calculates forces, but doesn't move anything, and your update_vel() method updates the velocities and moves the body. Renaming move(self,body) to be gravity_from(self,body) (or whatever makes sense to you) can make that clearer.

I would also move the update_pos() call out of update_vel(), and explicitly call the former in the loop.

What integrator is that, again?

Because you are putting the code for updating the velocities (and positions) in the outer loop of your motion function, you are calculating the forces on the Earth using the current position of the Earth and Luna, but you are calculating the forces on Luna using the new position of the Earth and the current position of Luna. While the two forces should be equal and opposite, this calculation won't do that.

Consider revisiting this integrator to do each step for all bodies: Compute all the forces, compute all the new velocities, compute all the new positions, rather than compute the force, new velocity, and new position for each body in turn.

Also, you can simplify your force calculations by changing them to acceleration calculations. The force from a body is \$F = -G\frac{Mm}{r^2}\$, while the acceleration is \$a = F/m\$. This gets you \$a = \mu/r^2\$ where \$\mu = -GM\$ is a constant per body. This means that your gravity_from() method would look like:

def acc_from(self,body):
    dx = (self.x - body.x) #Calculates difference in x- and y-axis between the bodies
    dy = (self.y - body.y)
    r2 = dy**2+dx**2       #Calculates the distance between the bodies
    angle = atan2(dy, dx)  #Calculates the angle between the bodies with atan2!
    if r < self.radius: #Checks if the distance between the bodies is less than the radius of the bodies. Uses then Gauss gravitational law to calculate force.
        A = 4/3 * pi * r * G * self.density
    else:  
        A = (self.mass*body.mu)/(r**2) #Newtons gravitational formula.
    Ax = cos(angle) * A
    Ay = sin(angle) * A
    return Ax,Ay

I also corrected your formula for gravity inside a solid sphere (of constant density).