Blame for scripts/pendulum_gen.py

#!/usr/bin/env python3
"""Pendulum core library: physics, rendering, and regime samplers.

Renderer and base physics re-implemented to match
`jimchen2/pendulum-cjepa-160f` (MIT), extended with OOD physics knobs:
quadratic drag and periodic driving.

ODE (generalized):
    theta'' = -(g/L) sin(theta) - gamma*theta' - beta*theta'|theta'| + drive(t)
    drive(t) = drive_amp * cos(2*pi*drive_freq*t + drive_phase)
"""
from __future__ import annotations

import csv
import math
import shutil
from dataclasses import dataclass, field, asdict
from pathlib import Path

import imageio.v2 as imageio
import numpy as np
from PIL import Image, ImageDraw

# --- renderer constants (identical to original dataset) --------------------
PIVOT = (10.0, 10.5)
DRAW_LENGTH = 8.0          # world units (physical length separately in metres)
BOB_RADIUS = 1.5
GROUND_Y, SUN_Y, SUN_RADIUS = -0.5, 20.5, 3.0
WORLD = (-5.0, 25.0, -3.0, 27.0)

FPS = 12
FRAMES = 160               # exported frames (matches C-JEPA TARGET_LEN)
EXT_FRAMES = 336           # simulated frames (28 s) -> ground truth for predictive Qs
SUBSTEPS = 8
SIZE_PX = 196

G_EARTH = 9.81
LENGTH_M = 1.0

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

@dataclass
class Configuration:
    name: str = "pendulum"
    theta0_deg: float = 35.0
    omega0_rad_s: float = 0.0
    damping_s_inv: float = 0.12
    light_mean_deg: float = 65.0
    light_amplitude_deg: float = 0.0
    light_frequency_hz: float = 0.08
    light_phase_rad: float = -math.pi / 2
    # physics extensions (all defaults == original C-JEPA pendulum)
    gravity_m_s2: float = G_EARTH
    length_m: float = LENGTH_M
    damping_model: str = "linear"          # "linear" | "quadratic"
    quad_drag_beta: float = 0.0            # beta in  -beta * w*|w|
    drive_amp_rad_s2: float = 0.0          # periodic forcing amplitude
    drive_freq_hz: float = 0.0
    drive_phase_rad: float = 0.0


# ---------------------------------------------------------------------------
# Physics
# ---------------------------------------------------------------------------

def angular_acceleration(theta, omega, t, cfg: Configuration):
    a = -(cfg.gravity_m_s2 / cfg.length_m) * math.sin(theta)
    if cfg.damping_model == "quadratic":
        a -= cfg.quad_drag_beta * omega * abs(omega)
    else:
        a -= cfg.damping_s_inv * omega
    if cfg.drive_amp_rad_s2 != 0.0:
        a += cfg.drive_amp_rad_s2 * math.cos(
            2 * math.pi * cfg.drive_freq_hz * t + cfg.drive_phase_rad)
    return a


def rk4_step(theta, omega, t, dt, cfg: Configuration):
    def f(a, b, tt):
        return b, angular_acceleration(a, b, tt, cfg)
    k1 = f(theta, omega, t)
    k2 = f(theta + dt*k1[0]/2, omega + dt*k1[1]/2, t + dt/2)
    k3 = f(theta + dt*k2[0]/2, omega + dt*k2[1]/2, t + dt/2)
    k4 = f(theta + dt*k3[0], omega + dt*k3[1], t + dt)
    return (theta + dt*(k1[0]+2*k2[0]+2*k3[0]+k4[0])/6,
            omega + dt*(k1[1]+2*k2[1]+2*k3[1]+k4[1])/6)


def scene_geometry(theta, light_deg):
    phi = math.radians(light_deg)
    cot = math.cos(phi) / math.sin(phi)
    x = PIVOT[0] + DRAW_LENGTH * math.sin(theta)
    y = PIVOT[1] - DRAW_LENGTH * math.cos(theta)
    pivot_shadow = PIVOT[0] + (GROUND_Y - PIVOT[1]) * cot
    bob_shadow = x + (GROUND_Y - y) * cot
    half_width = BOB_RADIUS / abs(math.sin(phi))
    left = min(pivot_shadow, bob_shadow - half_width)
    right = max(pivot_shadow, bob_shadow + half_width)
    return {
        "bob_x_world": x, "bob_y_world": y,
        "light_x_world": PIVOT[0] + (SUN_Y - PIVOT[1]) * cot,
        "shadow_left_world": left, "shadow_right_world": right,
        "shade": right - left, "mid": (left + right) / 2,
    }


def simulate(cfg: Configuration, frames=EXT_FRAMES, fps=FPS, substeps=SUBSTEPS):
    """Full time-series; row n samples time n/fps (same convention as original)."""
    theta, omega = math.radians(cfg.theta0_deg), cfg.omega0_rad_s
    g, L = cfg.gravity_m_s2, cfg.length_m
    dt = 1.0 / (fps * substeps)
    rows = []
    for frame in range(frames):
        t = frame / fps
        light = cfg.light_mean_deg + cfg.light_amplitude_deg * math.sin(
            2*math.pi*cfg.light_frequency_hz*t + cfg.light_phase_rad)
        row = {
            "frame": frame, "time_s": t,
            "theta_rad": theta, "theta_deg": math.degrees(theta),
            "omega_rad_s": omega,
            "alpha_rad_s2": angular_acceleration(theta, omega, t, cfg),
            "light_angle_deg": light,
            "bob_x_m": L * math.sin(theta),
            "bob_y_m": -L * math.cos(theta),
            "energy_j_per_kg": 0.5*L**2*omega**2 + g*L*(1-math.cos(theta)),
        }
        row.update(scene_geometry(theta, light))
        rows.append(row)
        if frame < frames - 1:
            for _ in range(substeps):
                theta, omega = rk4_step(theta, omega, t, dt, cfg)
    return rows


def validate_scene(row):
    xmin, xmax, ymin, ymax = WORLD
    return (xmin < row["shadow_left_world"] < row["shadow_right_world"] < xmax
            and xmin <= row["bob_x_world"]-BOB_RADIUS
            and row["bob_x_world"]+BOB_RADIUS <= xmax
            and ymin <= row["bob_y_world"]-BOB_RADIUS
            and row["bob_y_world"]+BOB_RADIUS <= ymax
            and xmin <= row["light_x_world"]-SUN_RADIUS
            and row["light_x_world"]+SUN_RADIUS <= xmax)


# ---------------------------------------------------------------------------
# Rendering (identical geometry/colors to the original)
# ---------------------------------------------------------------------------

def render_scene(row, size=SIZE_PX, supersample=3):
    high = size * supersample
    image = Image.new("RGB", (high, high), "white")
    draw = ImageDraw.Draw(image)
    xmin, xmax, ymin, ymax = WORLD
    scale = (high - 1) / (xmax - xmin)
    def point(x, y):
        return (scale*(x-xmin), scale*(ymax-y))
    def circle(x, y, radius, fill):
        draw.ellipse([point(x-radius, y+radius), point(x+radius, y-radius)], fill=fill)
    circle(row["light_x_world"], SUN_Y, SUN_RADIUS, "#ffa500")
    draw.line([point(xmin, GROUND_Y), point(xmax, GROUND_Y)],
              fill="#cbd5e1", width=max(1, round(0.08*scale)))
    draw.rectangle([point(row["shadow_left_world"], GROUND_Y+0.14),
                    point(row["shadow_right_world"], GROUND_Y-0.14)], fill="#171717")
    draw.line([point(*PIVOT), point(row["bob_x_world"], row["bob_y_world"])],
              fill="#171717", width=max(1, round(0.30*scale)))
    circle(*PIVOT, 0.18, "#171717")
    circle(row["bob_x_world"], row["bob_y_world"], BOB_RADIUS, "#b22222")
    return image.resize((size, size), Image.Resampling.LANCZOS)


def video_writer(path, fps=FPS):
    return imageio.get_writer(str(path), fps=fps, codec="libx264", quality=8,
                              pixelformat="yuv420p", macro_block_size=1,
                              ffmpeg_log_level="error",
                              ffmpeg_params=["-movflags", "+faststart"])


def write_csv(path, rows):
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


# ---------------------------------------------------------------------------
# Regime samplers
# ---------------------------------------------------------------------------

ID = "train"

def _sample_light(rng, cfg):
    cfg.light_mean_deg = rng.uniform(70.0, 110.0)
    cfg.light_amplitude_deg = rng.uniform(0.0, 15.0)
    cfg.light_frequency_hz = rng.uniform(0.04, 0.12)
    cfg.light_phase_rad = rng.uniform(0.0, 2*math.pi)

def _sample_ic(rng, cfg, th_lo=15.0, th_hi=40.0):
    cfg.theta0_deg = rng.uniform(th_lo, th_hi) * rng.choice([-1.0, 1.0])
    cfg.omega0_rad_s = rng.uniform(-0.2, 0.2)

def sample_config(regime: str, rng) -> Configuration:
    cfg = Configuration(name=f"{regime}")
    _sample_ic(rng, cfg)
    _sample_light(rng, cfg)
    if regime in (ID, "val_low_damp_baseline"):
        cfg.damping_s_inv = rng.uniform(0.15, 0.50)
    elif regime == "val_high_damp":
        cfg.damping_s_inv = rng.uniform(0.55, 0.90)
    elif regime == "val_low_damp":
        cfg.damping_s_inv = rng.uniform(0.0, 0.08)
    elif regime == "val_large_angle":
        _sample_ic(rng, cfg, 45.0, 80.0)
        cfg.damping_s_inv = rng.uniform(0.15, 0.50)
    elif regime == "val_gravity":
        cfg.gravity_m_s2 = float(rng.choice([1.62, 3.71, 24.79]))  # Moon, Mars, Jupiter
        cfg.damping_s_inv = rng.uniform(0.15, 0.50)
    elif regime == "val_quad_damping":
        cfg.damping_model = "quadratic"
        cfg.damping_s_inv = 0.0
        cfg.quad_drag_beta = rng.uniform(0.10, 0.35)
    elif regime == "val_driven":
        cfg.damping_s_inv = rng.uniform(0.15, 0.35)
        cfg.drive_amp_rad_s2 = rng.uniform(0.35, 1.0)
        cfg.drive_freq_hz = rng.uniform(0.35, 0.65)   # natural ~0.50 Hz
        cfg.drive_phase_rad = rng.uniform(0.0, 2*math.pi)
    else:
        raise ValueError(f"unknown regime {regime}")
    return cfg


REGIMES_VAL = ["val_high_damp", "val_low_damp", "val_large_angle",
               "val_gravity", "val_quad_damping", "val_driven"]

REGIME_DESCRIPTIONS = {
    "train": "in-distribution: theta0=±(15–40)°, omega0=±0.2 rad/s, gamma=U(0.15,0.50)/s, Earth gravity, linear damping, no forcing (matches jimchen2/pendulum-cjepa-160f)",
    "val_high_damp": "OOD physics: gamma=U(0.55,0.90)/s — much stronger air resistance than any training video (fast envelope decay, tau as low as 2.2 s)",
    "val_low_damp": "OOD physics: gamma=U(0.0,0.08)/s — almost un-damped; swings barely decay, energy nearly conserved",
    "val_large_angle": "OOD physics: theta0=±(45–80)° — strongly nonlinear regime; period elongated up to ~30% above small-angle value",
    "val_gravity": "OOD physics: g in {1.62 (Moon), 3.71 (Mars), 24.79 (Jupiter)} m/s^2 — periods shift by up to 2.5x either way",
    "val_quad_damping": "OOD physics (new law): drag = -beta*w*|w|, beta=U(0.10,0.35); envelope decays ~linearly instead of exponentially",
    "val_driven": "OOD physics (new forcing): periodic torque A*cos(2*pi*f*t+phi), A=U(0.35,1.0) rad/s^2, f=U(0.35,0.65) Hz near the ~0.5 Hz natural frequency — beats and resonant growth",
}


def resample_until_valid(regime: str, rng, max_tries=60):
    """Rejection-sample so every frame of the *simulated* rollout stays in-world."""
    for _ in range(max_tries):
        cfg = sample_config(regime, rng)
        rows = simulate(cfg)
        if all(validate_scene(r) for r in rows) and all(
                abs(r["theta_deg"]) <= 85.0 for r in rows):
            return cfg, rows
    raise RuntimeError(f"could not sample a valid {regime} config")


# ---------------------------------------------------------------------------
# Per-video generation job
# ---------------------------------------------------------------------------

def rows_to_np(rows):
    keys = ["theta_rad", "omega_rad_s", "energy_j_per_kg", "light_angle_deg",
            "bob_x_world", "bob_y_world", "shadow_left_world", "shadow_right_world",
            "mid", "bob_x_m"]
    return {k: np.array([r[k] for r in rows], dtype=np.float64) for k in keys}


def generate_one(job: dict, out_root: str):
    """Render one video + CSV; returns (index_row, analytics) for aggregation."""
    out = Path(out_root)
    split = job["split"]
    vid = job["video_id"]
    mp4_path = out / "videos" / split / f"{vid}.mp4"
    csv_path = out / "timeseries" / f"{vid}.csv"
    rel_mp4 = f"videos/{split}/{vid}.mp4"
    rel_csv = f"timeseries/{vid}.csv"

    rng = np.random.default_rng(np.random.SeedSequence([job["seed"], job["index"]]))
    cfg, rows = resample_until_valid(job["regime"], rng)

    if not mp4_path.exists() or mp4_path.stat().st_size < 1000:
        mp4_path.parent.mkdir(parents=True, exist_ok=True)
        tmp = mp4_path.with_suffix(".tmp.mp4")
        w = video_writer(tmp)
        try:
            for r in rows[:FRAMES]:
                w.append_data(np.asarray(render_scene(r)))
        finally:
            w.close()
        shutil.move(tmp, mp4_path)
    if not csv_path.exists():
        csv_path.parent.mkdir(parents=True, exist_ok=True)
        write_csv(csv_path, rows[:FRAMES])

    d = asdict(cfg)
    index_row = {
        "video_id": vid, "split": split, "regime": job["regime"],
        "video_path": rel_mp4, "csv_path": rel_csv,
        "frames": FRAMES, "duration_s": (FRAMES-1 + 1)/FPS,
        "fps": FPS, "size_px": SIZE_PX,
        "gravity_m_s2": cfg.gravity_m_s2, "length_m": cfg.length_m,
        "name": f"{job['regime']}_{job['index']:05d}",
        "theta0_deg": cfg.theta0_deg, "omega0_rad_s": cfg.omega0_rad_s,
        "damping_s_inv": cfg.damping_s_inv,
        "damping_model": cfg.damping_model, "quad_drag_beta": cfg.quad_drag_beta,
        "drive_amp_rad_s2": cfg.drive_amp_rad_s2, "drive_freq_hz": cfg.drive_freq_hz,
        "drive_phase_rad": cfg.drive_phase_rad,
        "light_mean_deg": cfg.light_mean_deg,
        "light_amplitude_deg": cfg.light_amplitude_deg,
        "light_frequency_hz": cfg.light_frequency_hz,
        "light_phase_rad": cfg.light_phase_rad,
    }
    return index_row, cfg, rows_to_np(rows)