2D simulation of Solar system
In this article we'll try to create 2D simulation of Solar system with python and pygame lib. I consider people that are reading it are familiar with pygame lib, just basics. Let's get started.
That's how looks framework of pygame app.
import pygame
import math
# colors in rgb type for drawing planets and Sun
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
YELLOW = (255, 255, 0)
BLUE = (26, 209, 255)
RED = (255, 92, 51)
GREY = (80, 71, 81)
ORANGE = (204, 163, 0)
pygame.init()
WIDTH, HEIGHT = 800, 800
WIN = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Solar system simulation")
def main():
run = True
clock = pygame.time.Clock()
while run:
# set up 60 fps
clock.tick(60)
WIN.fill(BLACK)
for event in pygame.event.get():
if event.type == pygame.QUIT:
run = False
pygame.quit()
main()If you start it, you'll see black sreen and working button for closing window.
Now let's define Planet class. It's also easy to change SCALE in simulation, just increase first number and check it out.
class Planet:
# astronomical unit, distance from the Sun to the Earth in meters
AU = 149.6e6 * 1000
G = 6.67428e-11
SCALE = 80 / AU # 1 AU = 100 pixels
TIMESTEP = 3600 * 24 # 1 day (3600 seconds = 1 hour)
Here is initializing function of Planet class, nothing complicated but I'd like to note I use sun variable to identify Sun. Technically sun isn't a planet, it's star but who cares? :) So, all planets have False value and only Sun has True.
def __init__(self, x, y, radius, color, mass):
self.x = x
self.y = y
self.radius = radius
self.color = color
self.mass = mass
self.sun = False
self.distance_to_sun = 0
self.orbit = []
self.x_velocity = 0
self.y_velocity = 0For now let's create draw function basement.
def draw(self, win):
# WIDTH / 2 and HEIGHT / 2 give a chance to draw in the middle of the scren
# in Pygame center of the screen is at the top left corner
x = self.x * self.SCALE + WIDTH / 2
y = self.y * self.SCALE + HEIGHT / 2
# draw circle with set color and radius in the middle of the screen
pygame.draw.circle(win, self.color, (x, y), self.radius)Now we'll add Sun and other planets to the simulation and check if draw function works right. Code below should be added in the main function before while cycle. The reason why I put - in first value for Earth and Mars it's because I want to see planets in different sides(left or right).
By the way, why this project is calling simulation because I use real numbers to simulate movements of planets with gravity. All numbers easy to google and check.
sun = Planet(0, 0, 15, YELLOW, 1.98892 * 10**30)
sun.sun = True
# further it's easy to add more planets
mercury = Planet(0.387 * Planet.AU, 0, 4, GREY, 0.330 * 10**23)
venus = Planet(0.723 * Planet.AU, 0, 7, WHITE, 4.8685 * 10**24)
earth = Planet(-1 * Planet.AU, 0, 8, BLUE, 5.9742 * 10**24)
mars = Planet(-1.524 * Planet.AU, 0, 6, RED, 6.39 * 10**23)
planets = [sun, mercury, venus, earth, mars]
def main():
run = True
clock = pygame.time.Clock()
while run:
# set up 60 fps
clock.tick(60)
WIN.fill(BLACK)
for event in pygame.event.get():
if event.type == pygame.QUIT:
run = False
for planet in planets:
planet.draw(WIN)
pygame.display.update()
pygame.quit()
main()That's what we got. Planets with actual distances reduced in many times
Now I want to disscus the hardest part of article - realisation of gravity. We need to calculate distance beetwen two points and force that acts on two objects and farther with using cos and sin calculate Fy and Fx movement.
def attraction(self, other):
# calculating R(distance between 2 points)
other_x, other_y = other.x, other.y
distance_x = other_x - self.x
distance_y = other_y - self.y
distance = math.sqrt(distance_x ** 2 + distance_y ** 2)
if other.sun:
self.distance_to_sun = distance
# calculating force_x and force_y to define how to move object
force = self.G * self.mass * other.mass / distance ** 2
angle = math.atan2(distance_y, distance_x)
force_x = math.cos(angle) * force
force_y = math.sin(angle) * force
return force_x, force_yNext step is setup start y velocity for planets in main function and add new if your sreen is quite big for this :) Again minus before number determines direction of movement(clockwise or counterclockwise).
# start y velocity mercury.y_velocity = -47.4 * 1000 venus.y_velocity = -35.02 * 1000 earth.y_velocity = 29.783 * 1000 mars.y_velocity = 24.077 * 1000 # new planet if your screen is big enough jupiter = Planet(5.2 * Planet.AU, 0, 10, ORANGE, 1.898 * 10**27) jupiter.y_velocity = -13.1 * 1000
Now we'll complete update function. Here we count x velocity for planets. You can test and delete starting values for y velocity, you'll see something like falling planets on the Sun without moving around it.
def update_position(self, planets):
total_force_x = total_force_y = 0
for planet in planets:
if self == planet:
continue
fx, fy = self.attraction(planet)
total_force_x += fx
total_force_y += fy
# F = m * a, Newton's second law of motion then a = F / m
self.x_velocity += total_force_x / self.mass * self.TIMESTEP
self.y_velocity += total_force_y / self.mass * self.TIMESTEP
self.x += self.x_velocity * self.TIMESTEP
self.y += self.y_velocity * self.TIMESTEP
self.orbit.append((self.x, self.y))
Of course add new function update_position to the main function!
planet.update_position(planets)
Let's add some orbits for moving planets. This code should be added to the draw function before drawing circle(planet).
if len(self.orbit) > 2:
updated_points = []
# we get all points during the movement and with it draw line
for point in self.orbit:
x, y = point
x = x * self.SCALE + WIDTH / 2
y = y * self.SCALE + HEIGHT / 2
updated_points.append((x, y))
# False means it isn't enclosed line
pygame.draw.lines(win, self.color, False, updated_points, 1)For now you get working simulation that looks pretty cool for me.
As always full code can be found at my GitHub: https://github.com/rastr-0/solar_system_simulation/blob/main/main.py