solar_system_arcade.py

"""
Solar system simulation. Visualisation with Arcade.
"""

import arcade


class Vec2d:
    def __init__(self, x: float, y: float):
        self.x = x
        self.y = y


class Body:
    def __init__(
        self,
        *,
        name: str,
        mass: float = 1,
        pos: Vec2d = Vec2d(0, 0),
        vel: Vec2d = Vec2d(0, 0),
        radius: float = 1,
        color: arcade.types.Color = arcade.color.PINK,
    ):
        self.name = name
        self.mass = mass
        self.pos = pos
        self.vel = vel
        self.radius = radius
        self.color = color


class NBodySystem:
    def __init__(self, bodies: list[Body]):
        self.bodies = bodies
        self.orbits = []


class SystemView(arcade.View):
    def __init__(self, system: NBodySystem):
        super().__init__()
        self.system = system
        self.background_color = arcade.color.BLACK
        self.camera = arcade.Camera2D(
            position=(0, 0),
            projection=arcade.XYWH(x=0, y=0, width=8, height=8)
        )

    def on_draw(self):
        self.clear()
        self.camera.use()
        for body in self.system.bodies:
            arcade.draw_circle_filled(
                body.pos.x, body.pos.y, radius=body.radius, color=body.color, num_segments=30,
            )


def main():
    sun = Body(name="Sun", mass=100, pos=Vec2d(0, 0), vel=Vec2d(0, 0), radius=0.2, color=arcade.color.YELLOW)
    earth = Body(name="Earth", mass=1, pos=Vec2d(1, 0), vel=Vec2d(0, 1), radius=0.1, color=arcade.color.BLUE)
    test = Body(name="Mars", mass=1, pos=Vec2d(3, 3), vel=Vec2d(0, 1), radius=0.1, color=arcade.color.RED)
    system = NBodySystem([sun, earth, test])

    # TODO: implement camera to select what to show or to follow bodies around
    window = arcade.Window(800, 800, "Solar system simulation")
    system_view = SystemView(system)
    window.show_view(system_view)
    arcade.run()

    # TODO: Decide where to store the orbits (on Body or NBodySystem)
    # and what data structure to use (list[Body] or deparate lists)
    # Add in timestamps to the simulation! Multi-time resolution?

    # system.simulate(n_steps=100, dt=0.1)

    # TODO: add Arcade visualisation and playback
    # Simulation timestep `dt` should be separate from visualisation playback speed (and timestep?)
    # Allowing detailed simulation and fast playback.

    # TODO: Implement body lookup by "name" and state lookup by "timestamp" (nearest neighbor)


if __name__ == "__main__":
    main()