Kepler orbit simulation

orbit.py

"""Simulation of a Kepler orbit.

I wanted to make some reusable components to simulate Kepler orbits
for use in context of games. Rather than running a physics simulation,
I figured it would be for more efficient to use the mathmatical models
of Kepler orbits.

Required Packages:
- pygame
- esper
"""
from dataclasses import dataclass as component
from math import cos, pi, sin

import esper
import pygame
from pygame.locals import *

DISPLAY_WIDTH = 500
DISPLAY_HEIGHT = 500

dirty_rects = []


@component
class Position:
    """A component for object which exist in 2D space.
    Attributes:
        x (float): The position on the X axis.
        y (float): The position on the y axis.
    """
    x: float = 0.0
    y: float = 0.0


@component
class Renderable:
    """A component for renderable objects.
    Attributes:
        surface (:obj:`pygame.Surface`): The renderable surface.
    """
    surface: pygame.Surface = None


@component
class Orbit:
    """Structured data describing a simplified orbit.
    
    For more info, see:
    https://en.wikipedia.org/wiki/Orbital_elements

    Attributes:
        argument_of_periapsis (float): Defines the orientation of the ellipse
            in the orbital plane, as an angle measured from the ascending node
            to the periapsis.
        center (:obj:`Position`): The center point of the orbit.
        eccentricity (float): Shape of the ellipse, describing how much it is
            elongated compared to a circle.
            https://en.wikipedia.org/wiki/Orbital_eccentricity
        period (float): The time it takes to complete one full orbit.
        semimajor_axis (float): The sum of the periapsis and apoapsis distances
            divided by two.
        true_anomaly (float): Defines the position of the orbiting body along
            the ellipse at a specific time (the "epoch").
    """
    argument_of_periapsis: float = 0.0
    center: Position = Position(0.0, 0.0)
    eccentricity: float = 0.0
    period: float = 0.0
    semimajor_axis: float = 1.0
    true_anomaly: float = 0.0


class OrbitProcessor(esper.Processor):
    """A processor for :obj:`Orbit` components.

    Moves entities in 2D space according to a model of ideal orbits.

    For more info, see:
    https://en.wikipedia.org/wiki/Elliptic_orbit
    https://en.wikipedia.org/wiki/Kepler_orbit
    http://www.physics.csbsju.edu/orbit/orbit.2d.html
    https://math.stackexchange.com/a/3711214
    """
    def process(self, pygame_screen, background, time):
        for entity, (orbit, position, renderable) \
                in self.world.get_components(Orbit,
                                             Position,
                                             Renderable):
            # Add current position to dirty_rects and fill it with the background image
            extents = renderable.surface.get_size()
            old_rect = makeRect(
                position,
                Position(x=extents[0], y=extents[1])
            )
            pygame_screen.blit(background, (position.x, position.y), old_rect)
            dirty_rects.append(old_rect)

            # Set position to the calculated posiiton on the orbit elipse7
            x = orbit.semimajor_axis * (cos((time/orbit.period)+orbit.true_anomaly) - orbit.eccentricity)
            y = orbit.semimajor_axis * (1 - orbit.eccentricity**2)**0.5 * sin((time/orbit.period)+orbit.true_anomaly)

            position.x = x * cos(orbit.argument_of_periapsis) \
                         - y * sin(orbit.argument_of_periapsis) \
                         + orbit.center.x
            position.y = x * sin(orbit.argument_of_periapsis) \
                         + y * cos(orbit.argument_of_periapsis) \
                         + orbit.center.y

            # Add new position to dirty_rects
            dirty_rects.append(makeRect(
                position,
                Position(x=extents[0], y=extents[1])
            ))


class RenderProcessor(esper.Processor):
    """A proccessor for Renderable components."""
    def process(self, pygame_screen, background, time):
        for entity, (renderable, position) \
                in self.world.get_components(Renderable, Position):
            pygame_screen.blit(renderable.surface,
                               (position.x, position.y))


def makeDotSurface(size: int) -> pygame.Surface:
    """Create a surface with a dot on it.
    Args:
        size (int): Diameter of the dot and size of the surface.
    """
    surface = pygame.Surface((size, size))
    pygame.draw.circle(surface, (255,255,255), (size/2,size/2), size/2)
    return surface


def makeRect(origin: Position, extent: Position) -> pygame.Rect:
    """Create a rect using two :obj:`Position` objects.
    Args:
        origin (:obj:`Position`): The coordinates of the top-left corner.
        extent (:obj:`Position`): The size of the rect.
    """
    return pygame.Rect(origin.x, origin.y, extent.x, extent.y)


if __name__ == '__main__':
    # Init pygame
    pygame.init()
    screen = pygame.display.set_mode((DISPLAY_WIDTH, DISPLAY_HEIGHT),
                                     pygame.SCALED)
    pygame.display.set_caption("Orbit Test")
    pygame.mouse.set_visible(True)
    background = pygame.Surface(screen.get_size())
    background = background.convert()
    background.fill((0, 0, 0))
    clock = pygame.time.Clock()
    time = 0

    # Init Esper
    world = esper.World()
    world.add_processor(OrbitProcessor())
    world.add_processor(RenderProcessor())

    # Init planet
    celestial_body = world.create_entity(
        Renderable(makeDotSurface(10)),
        Position(x=DISPLAY_WIDTH/2, y=DISPLAY_HEIGHT/2),
        Orbit(argument_of_periapsis=0.0,
              center=Position(DISPLAY_WIDTH/2, DISPLAY_HEIGHT/2),
              eccentricity=0.0,
              period=1000,
              semimajor_axis=50.0,
              true_anomaly=0.0)
    )

    # Init star
    star = world.create_entity(
        Renderable(makeDotSurface(10)),
        Position(x=DISPLAY_WIDTH/2, y=DISPLAY_HEIGHT/2)
    )

    # Initial BG fill
    screen.blit(background, (0, 0))

    quit = False
    while not quit:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                quit = True

        dirty_rects = []
        world.process(screen, background, time)
        pygame.display.update(dirty_rects)
        time += clock.tick(60)
        
    pygame.quit()