solar_system.py

import numpy as np
import spatialstudio as splv
from tqdm import tqdm

"""Complete Solar System Animation – High-Resolution 512³ Voxel Animation
----------------------------------------------------------------------
This script generates a seamless looping animation of our complete solar system
with all 8 planets orbiting around a bright sun with realistic relative sizes
and orbital distances (scaled for visibility).

• Grid: 512 × 512 × 512
• Duration: 20 s @ 30 fps (600 frames)
• Output: ../outputs/solar_system.splv

All planets complete their orbits at different rates for realistic motion,
with the inner planets moving faster than outer planets.
"""

# Scene parameters
SIZE, FPS, SECONDS = 256, 30, 90  # 12-second loop for faster rendering
FRAMES = FPS * SECONDS

CENTER_X = SIZE // 2
CENTER_Y = SIZE // 2
CENTER_Z = SIZE // 2
OUT_PATH = "../outputs/solar_system.splv"

# Solar parameters
SUN_RADIUS = 20
FLARE_COUNT = 12

# Planet data: [radius, orbit_radius, orbit_speed_multiplier, color]
# Orbit speeds are relative - inner planets orbit faster
PLANET_DATA = {
    "mercury": [1, 40, 4.0, (169, 169, 169)],     # Gray
    "venus": [2, 48, 3.5, (255, 198, 73)],        # Yellow-orange
    "earth": [2, 55, 3.0, (100, 149, 237)],       # Blue
    "mars": [2, 65, 2.5, (205, 92, 92)],          # Red
    "jupiter": [6, 85, 1.5, (255, 140, 0)],       # Orange
    "saturn": [5, 100, 1.2, (255, 218, 185)],     # Pale yellow
    "uranus": [3, 115, 0.8, (64, 224, 208)],      # Turquoise
    "neptune": [3, 125, 0.6, (70, 130, 180)],     # Steel blue
}

def smooth_ease(t: float) -> float:
    """Smoothstep easing for nicer motion curves."""
    return t * t * (3.0 - 2.0 * t)

def add_voxel(vol: np.ndarray, x: int, y: int, z: int, color: tuple[int, int, int]):
    """Safely set RGBA on a voxel inside the volume."""
    if 0 <= x < SIZE and 0 <= y < SIZE and 0 <= z < SIZE:
        vol[x, y, z, :3] = color
        vol[x, y, z, 3] = 255

def generate_sun(vol: np.ndarray, t_norm: float):
    """Render the sun with pulsating surface activity and flares."""
    # Base sun colours
    sun_colours = [
        (255, 245, 120),  # bright yellow core
        (255, 235,  90),  # vivid yellow
        (255, 220,  40),  # golden yellow
        (255, 200,   0),  # strong yellow-orange
        (255, 175,   0),  # deeper orange-yellow
    ]

    # Surface pulsation amplitude
    pulsate = 1.0 + 0.15 * np.sin(t_norm * 2 * np.pi * 4)

    for dx in range(-SUN_RADIUS, SUN_RADIUS + 1):
        for dy in range(-SUN_RADIUS, SUN_RADIUS + 1):
            for dz in range(-SUN_RADIUS, SUN_RADIUS + 1):
                r = np.sqrt(dx * dx + dy * dy + dz * dz)
                if r <= SUN_RADIUS:
                    depth = 1.0 - r / SUN_RADIUS
                    colour_idx = int(depth * (len(sun_colours) - 1))
                    base_colour = sun_colours[colour_idx]
                    
                    # Rotate coordinates around Y-axis for visible sun spin
                    phi = t_norm * 2 * np.pi
                    rot_x = dx * np.cos(phi) - dz * np.sin(phi)
                    rot_z = dx * np.sin(phi) + dz * np.cos(phi)

                    # Surface activity noise using rotated coords
                    activity = 0.9 + 0.1 * np.sin(rot_x * 0.3 + dy * 0.4 + rot_z * 0.3 + t_norm * np.pi * 10)

                    # Final brightened colour
                    final_colour = tuple(min(255, int(c * activity * pulsate * 1.3)) for c in base_colour)
                    add_voxel(vol, CENTER_X + dx, CENTER_Y + dy, CENTER_Z + dz, final_colour)

    # Solar flares
    flare_base_radius = SUN_RADIUS + 6
    for i in range(FLARE_COUNT):
        angle = (i / FLARE_COUNT) * 2 * np.pi + t_norm * 2 * np.pi * 2
        fx = CENTER_X + int((flare_base_radius + 6 * np.sin(t_norm * 2 * np.pi * 6 + i)) * np.cos(angle))
        fz = CENTER_Z + int((flare_base_radius + 6 * np.sin(t_norm * 2 * np.pi * 6 + i)) * np.sin(angle))
        fy = CENTER_Y + int(6 * np.cos(angle * 3))
        add_voxel(vol, fx, fy, fz, (255, 240, 120))

def generate_planet(vol: np.ndarray, t_norm: float, planet_name: str, planet_data: list):
    """Render a planet with orbital motion."""
    radius, orbit_radius, speed_mult, color = planet_data
    
    # Orbital position - different speeds for each planet
    theta_orbit = t_norm * 2 * np.pi * speed_mult
    px = CENTER_X + int(orbit_radius * np.cos(theta_orbit))
    pz = CENTER_Z + int(orbit_radius * np.sin(theta_orbit))
    
    # Small vertical oscillation for visual interest
    py = CENTER_Y + int(3 * np.sin(theta_orbit * 1.5))

    # Planet self-rotation
    spin_angle = t_norm * 2 * np.pi * (speed_mult + 2)  # Faster spin for inner planets

    for dx in range(-radius, radius + 1):
        for dy in range(-radius, radius + 1):
            for dz in range(-radius, radius + 1):
                r = np.sqrt(dx * dx + dy * dy + dz * dz)
                if r <= radius:
                    # Simple shading based on position
                    shade = 0.7 + 0.3 * np.cos(dx * 0.5 + dy * 0.3 + spin_angle)
                    
                    # Apply shading to base color
                    final_color = tuple(max(50, int(c * shade)) for c in color)
                    
                    add_voxel(vol, px + dx, py + dy, pz + dz, final_color)

def generate_saturn_rings(vol: np.ndarray, t_norm: float):
    """Generate Saturn's distinctive rings."""
    saturn_data = PLANET_DATA["saturn"]
    radius, orbit_radius, speed_mult, _ = saturn_data
    
    # Saturn's orbital position (must match generate_planet)
    theta_orbit = t_norm * 2 * np.pi * speed_mult
    sx = CENTER_X + int(orbit_radius * np.cos(theta_orbit))
    sz = CENTER_Z + int(orbit_radius * np.sin(theta_orbit))
    sy = CENTER_Y + int(3 * np.sin(theta_orbit * 1.5))
    
    ring_inner = radius + 2
    ring_outer = radius + 8
    ring_color = (200, 180, 140)  # Pale brown
    
    # Generate ring particles in a disk (reduced for performance)
    for angle in np.linspace(0, 2 * np.pi, 100):
        for ring_r in range(ring_inner, ring_outer, 3):
            # Ring particles with some random variation
            particle_angle = angle + 0.1 * np.sin(t_norm * 2 * np.pi + angle * 10)
            rx = sx + int(ring_r * np.cos(particle_angle))
            rz = sz + int(ring_r * np.sin(particle_angle))
            
            add_voxel(vol, rx, sy, rz, ring_color)

def generate_asteroid_belt(vol: np.ndarray, t_norm: float):
    """Generate asteroid belt between Mars and Jupiter."""
    belt_radius = 75  # Between Mars (65) and Jupiter (85)
    asteroid_count = 50  # Reduced for performance
    
    for i in range(asteroid_count):
        # Asteroid orbital position
        base_angle = (i / asteroid_count) * 2 * np.pi
        orbit_variation = 5 * np.sin(i * 0.3)  # Orbital radius variation
        current_radius = belt_radius + orbit_variation
        
        # Slow orbital motion
        asteroid_angle = base_angle + t_norm * 2 * np.pi * 0.5
        
        ax = CENTER_X + int(current_radius * np.cos(asteroid_angle))
        az = CENTER_Z + int(current_radius * np.sin(asteroid_angle))
        ay = CENTER_Y + int(2 * np.sin(asteroid_angle * 3 + i))
        
        # Small asteroid - just single voxels for performance
        asteroid_color = (100, 80, 60)  # Brown-gray
        add_voxel(vol, ax, ay, az, asteroid_color)

# Encoder setup & frame generation loop
enc = splv.SPLVencoder(
    width=SIZE,
    height=SIZE,
    depth=SIZE,
    framerate=FPS,
    motionVectors="off",
    outputPath=OUT_PATH,
)

for frame_idx in tqdm(range(FRAMES), desc="generating complete solar system"):
    vol = np.zeros((SIZE, SIZE, SIZE, 4), dtype=np.uint8)

    t_norm = frame_idx / FRAMES  # 0->1 over the full animation

    # Generate sun
    generate_sun(vol, t_norm)
    
    # Generate all planets
    for planet_name, planet_data in PLANET_DATA.items():
        generate_planet(vol, t_norm, planet_name, planet_data)
    
    # Generate Saturn's rings
    generate_saturn_rings(vol, t_norm)
    
    # Generate asteroid belt
    generate_asteroid_belt(vol, t_norm)

    enc.encode(splv.frame_from_numpy(vol))

enc.finish()
print(f"created {OUT_PATH}")