#!/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)