pygame-law-of-universal-gravitation

import pygame, math
from pygame.locals import *

SCALE = 10 ** -6.58
REAL_RADIUS_OF_EARTH = 6_371 * 10 ** 3
REAL_RADIUS_OF_SUN = 696_340 * 10 ** 3
G =  9.807

pygame.init()

display = pygame.display.set_mode((600, 600))
game_is_running = True
clock = pygame.time.Clock()

class Body:
  def __init__(self, position, radius = 50.0, density = 0.89) -> None:
    self.position = position
    self.radius = radius
    self.density = density
    self.mass = self.radius ** 2 * math.pi * self.density

    self.acceleration = [0, 0]
    self.velocity = [0, 0]

  def draw(self) -> None:
    pygame.draw.circle(display, (255, 255, 255), self.position, self.radius)

class LawOfUniversalGravitatoin:
  def __init__(self, body_a: Body, body_b: Body) -> None:
    self.body_a = body_a
    self.body_b = body_b

  def step(self) -> None:
    angle = math.atan2(self.body_b.position[1] - self.body_a.position[1], self.body_b.position[0] - self.body_a.position[0])
    magnitude = math.sqrt((self.body_b.position[0] - self.body_a.position[0]) ** 2 + (self.body_b.position[1] - self.body_a.position[1]) ** 2)

    F = G * (self.body_a.mass * self.body_b.mass) / magnitude

    self.body_a.acceleration[0] += math.cos(angle) * F / self.body_a.mass * SCALE
    self.body_a.velocity[0] = self.body_a.acceleration[0]
    self.body_a.acceleration[1] += math.sin(angle) * F / self.body_a.mass * SCALE
    self.body_a.velocity[1] = self.body_a.acceleration[1]
    self.body_a.position[1] += self.body_a.velocity[1]
    self.body_a.position[0] += self.body_a.velocity[0]

    self.body_b.acceleration[0] += math.cos(angle) * F / self.body_b.mass * SCALE
    self.body_b.velocity[0] = self.body_b.acceleration[1]
    self.body_b.acceleration[1] += math.sin(angle) * F / self.body_b.mass * SCALE
    self.body_b.velocity[1] = self.body_b.acceleration[1]
    self.body_b.position[1] += self.body_b.velocity[1]
    self.body_b.position[0] += self.body_b.velocity[0]

earth_body = Body([400, 100], REAL_RADIUS_OF_EARTH * SCALE)
earth_body.acceleration[0] = -0.5

sun_body = Body([400, 400], REAL_RADIUS_OF_SUN * SCALE)

law_of_universal_gravitation = LawOfUniversalGravitatoin(earth_body, sun_body)

while game_is_running:
  law_of_universal_gravitation.step()

  for event in pygame.event.get():
    if event.type == pygame.QUIT:
      game_is_running = False

  display.fill([20, 20, 20])

  earth_body.draw()
  sun_body.draw()

  pygame.display.flip()
  clock.tick(120)

pygame.quit()