| |
| """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 |
|
|
| |
| PIVOT = (10.0, 10.5) |
| DRAW_LENGTH = 8.0 |
| 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 |
| EXT_FRAMES = 336 |
| SUBSTEPS = 8 |
| SIZE_PX = 196 |
|
|
| G_EARTH = 9.81 |
| LENGTH_M = 1.0 |
|
|
| |
| |
| |
|
|
| @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 |
| |
| gravity_m_s2: float = G_EARTH |
| length_m: float = LENGTH_M |
| damping_model: str = "linear" |
| quad_drag_beta: float = 0.0 |
| drive_amp_rad_s2: float = 0.0 |
| drive_freq_hz: float = 0.0 |
| drive_phase_rad: float = 0.0 |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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])) |
| 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) |
| 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") |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
