arcade_solar_system_01.py

"""
Solar system simulation. Visualisation with Arcade.
"""
import collections
import dataclasses
import math
import time
import arcade


@dataclasses.dataclass
class Vec2d:
    """Generic 2-dim vector (x, y)."""
    x: float = 0
    y: float = 0

    def length(self) -> float:
        return math.sqrt(self.x ** 2 + self.y ** 2)
    
    def normalised(self) -> "Vec2d":
        return self / self.length()

    def __add__(self, other: "Vec2d") -> "Vec2d":
        return self.__class__(self.x + other.x, self.y + other.y)
        
    def __sub__(self, other: "Vec2d") -> "Vec2d":
        return self.__class__(self.x - other.x, self.y - other.y)

    def __rmul__(self, factor: float) -> "Vec2d":
        return self.__class__(factor * self.x, factor * self.y)

    def __truediv__(self, factor: float) -> "Vec2d":
        return self.__class__(self.x / factor, self.y / factor)


@dataclasses.dataclass
class OrbitPoint:
    time: float
    pos: Vec2d
    vel: Vec2d


@dataclasses.dataclass
class Body:
    """A physical body, e.g. star or planet."""
    name: str
    mass: float = 1
    pos: Vec2d = dataclasses.field(default_factory=lambda: Vec2d(0, 0))
    vel: Vec2d = dataclasses.field(default_factory=lambda: Vec2d(0, 0))
    orbit: collections.deque[OrbitPoint] = dataclasses.field(default_factory=lambda: collections.deque(maxlen=1000))
    radius: float = 1
    color: arcade.types.Color = arcade.color.PINK



class NBodySystem:
    # Gravitational constant
    G: float = 1

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

    def update(self, dt: float):
        """Update system by time step `dt`."""
        # Update velocity of each body
        for body in self.bodies:
            f = self.force_gravity(body)
            # Apply law of motion F = m * a
            a = f /  body.mass
            body.vel += dt * a
        
        # Update position of each body
        for body in self.bodies:
            body.pos += dt * body.vel

        # Store new point in orbit
        t = time.time()
        for body in self.bodies:
            point = OrbitPoint(time=t, pos=body.pos, vel=body.vel)
            body.orbit.append(point)

    def force_gravity(self, body: Body):
        """Compute total gravitational attraction force on the body.
        
        For each other body: F = - G * m * M / r ** 2
        """
        f = Vec2d(0, 0)
        for b in self.bodies:
            if body is not b:
                r = body.pos - b.pos
                f += - self.G * body.mass * b.mass / r.length() ** 2 * r.normalised()
        return f

    # TODO: implement center of mass computation
    # TODO: implement total momentum computation
    # TODO: implement total kinetic and potential energy computation
    # TODO: implement checks on 

class SystemView(arcade.View):

    # TODO: Maybe create a separate `Clock` to run the simulation that's separate from playback?
    # TODO: implement option to have camera or coords centered on center of mass / momentum
    def __init__(self, *, system: NBodySystem, radius: float = 10, dt_speedup = 0.1):
        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=2 * radius, height=2 * radius)
        )
        self.dt_speedup = dt_speedup

    def on_update(self, dt):
        self.system.update(self.dt_speedup * dt)


    def on_draw(self):
        self.clear()
        self.camera.use()
        for body in self.system.bodies:
            # Draw orbit history
            for point in body.orbit:
                # TODO: Create fadeaway effect for color, or transparency, or getting smaller?
                arcade.draw_circle_filled(
                    point.pos.x, point.pos.y, radius=body.radius/10, color=body.color, num_segments=30,
                )

            # Draw current positions
            arcade.draw_circle_filled(
                body.pos.x, body.pos.y, radius=body.radius, color=body.color, num_segments=30,
            )

def make_solar_system():
    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, 10), radius=0.1, color=arcade.color.BLUE)
    mars = Body(name="Mars", mass=1, pos=Vec2d(2, 0), vel=Vec2d(0, 7), radius=0.1, color=arcade.color.RED)
    return NBodySystem([sun, earth, mars])


def test_vec():
    system = make_solar_system()
    system.update(0.1)
    print(system.bodies[0].pos)
    print(system.bodies)


def main():
    system = make_solar_system()
    window = arcade.Window(800, 800, "Solar system simulation")
    system_view = SystemView(system=system, radius=5, dt_speedup=0.2)
    window.show_view(system_view)
    arcade.run()

    # 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__":
    test_vec()
    main()