Blame for scripts/questions.py

#!/usr/bin/env python3
"""CLEVRER-style question generation for the pendulum dataset.

Four CLEVRER question types, adapted from collision scenes to pendulum physics:
  descriptive    - observable attributes/events           (open answer)
  explanatory    - what is responsible for what is seen   (4-way multiple choice)
  predictive     - what happens next / eventually         (4-way multiple choice)
  counterfactual - what if a physical factor changed      (4-way multiple choice)

Every answer and every gating condition is computed from the ground-truth
RK4 trajectory (336 frames = 28 s), so answers about the future have real
ground truth. `pendulum_gen.simulate` provides rows; `analyze` builds facts.
"""
from __future__ import annotations

import math
import numpy as np

FPS = 12
FRAMES = 160
PIVOT_X = 10.0

# ---------------------------------------------------------------------------
# Trajectory analysis -> fact table used by question templates
# ---------------------------------------------------------------------------

def _peaks(theta, omega):
    """Indices/times/amplitudes of turning points (|theta| local maxima)."""
    idx, times, amps = [], [], []
    for i in range(len(omega) - 1):
        if omega[i] == 0.0 or omega[i] * omega[i + 1] < 0:
            frac = abs(omega[i]) / (abs(omega[i]) + abs(omega[i + 1]) + 1e-12)
            idx.append(i)
            times.append((i + frac) / FPS)
            amps.append(max(abs(theta[i]), abs(theta[i + 1])))
    return np.array(idx, dtype=int), np.array(times), np.array(amps)


def _zero_crossings(theta):
    """Times at which the bob passes directly below the pivot (theta = 0)."""
    out = []
    for i in range(len(theta) - 1):
        if theta[i] == 0.0 or theta[i] * theta[i + 1] < 0:
            frac = abs(theta[i]) / (abs(theta[i]) + abs(theta[i + 1]) + 1e-12)
            out.append((i + frac) / FPS)
    return np.array(out)


def _shadow_series(theta, light_deg):
    phi = np.radians(light_deg)
    cot = np.cos(phi) / np.sin(phi)
    x = PIVOT_X + 8.0 * np.sin(theta)
    y = 10.5 - 8.0 * np.cos(theta)
    pivot_sh = PIVOT_X + (-0.5 - 10.5) * cot
    bob_sh = x + (-0.5 - y) * cot
    hw = 1.5 / np.abs(np.sin(phi))
    left = np.minimum(pivot_sh, bob_sh - hw)
    right = np.maximum(pivot_sh, bob_sh + hw)
    return (left + right) / 2


def analyze(cfg, A):
    th = A["theta_rad"]                      # 336 frames, 28 s
    om = A["omega_rad_s"]
    en = A["energy_j_per_kg"]
    light = A["light_angle_deg"]

    p_idx, p_t, p_a = _peaks(th, om)
    in_vid = p_idx < FRAMES
    pv_t, pv_a = p_t[in_vid], p_a[in_vid]
    cross = _zero_crossings(th[:FRAMES])

    f = {}
    f["theta0"] = th[0]
    f["release_side"] = "right" if th[0] > 0 else "left"
    f["num_bottom_passes"] = int(len(cross))
    f["peaks_times"], f["peaks_amps"] = pv_t, pv_a
    f["n_peaks_video"] = len(pv_a)
    if len(pv_a) >= 2:
        f["env_ratio"] = float(pv_a[-1] / pv_a[0])          # last / first peak
        f["first_peak_deg"] = float(np.degrees(pv_a[0]))
        f["last_peak_deg"] = float(np.degrees(pv_a[-1]))
        f["later_peak_max_ratio"] = float(pv_a[1:].max() / pv_a[0])
    emax = float(np.max(np.abs(om)))
    f["omega_max"] = emax

    # last-second motion
    omega_slice = om[FRAMES - 12:FRAMES]
    omax_last = float(np.max(np.abs(omega_slice)))
    f["last_sec_omega_max"] = omax_last
    i_last = FRAMES - 12 + int(np.argmax(np.abs(omega_slice)))
    f["last_sec_dir"] = ("right" if om[i_last] > 0 else "left") if omax_last >= 0.15 else None
    f["nearly_still_end"] = (omax_last < 0.08) or (f.get("last_peak_deg", 99) < 2.5)

    # after the video (frames 159.. use simulated future)
    om_after = om[FRAMES + 2]
    f["after_dir"] = ("right" if om_after > 0 else "left") if abs(om_after) >= 0.15 else None
    f["after_omega"] = float(om_after)

    # next peak after the video vs last one inside
    fut = p_idx >= FRAMES
    if len(pv_a) >= 1 and fut.any():
        f["next_peak_ratio"] = float(p_a[fut][0] / pv_a[-1])

    # bottom crossing in the next second?
    t_end = (FRAMES - 1) / FPS
    cross_ext = _zero_crossings(th)
    win = (cross_ext > t_end) & (cross_ext <= t_end + 1.0)
    marg = (cross_ext > t_end) & (cross_ext <= t_end + 1.25)
    if win.any():
        f["crossing_next_second"] = True
    elif not marg.any():
        f["crossing_next_second"] = False
    else:
        f["crossing_next_second"] = None                     # too close to call

    # energy shape
    e = en[:FRAMES]
    diffs = np.diff(e)
    if np.all(diffs <= 0.001 * e[0]):
        f["energy_shape"] = "steadily decreases"
    elif (e.max() - e.min()) < 0.05 * e.max():
        f["energy_shape"] = "stays nearly constant"
    else:
        f["energy_shape"] = "repeatedly rises and falls"

    # shadow attribution: bob term vs light term
    mid_real = _shadow_series(th[:FRAMES], light[:FRAMES])
    mid_bob = _shadow_series(th[:FRAMES], np.full(FRAMES, float(np.mean(light[:FRAMES]))))
    mid_light = _shadow_series(np.zeros(FRAMES), light[:FRAMES])
    f["shadow_bob_range"] = float(mid_bob.max() - mid_bob.min())
    f["shadow_light_range"] = float(mid_light.max() - mid_light.min())
    f["shadow_sweep_range"] = float(mid_real.max() - mid_real.min())
    f["bob_dominates_shadow"] = (f["shadow_bob_range"] > 2 * f["shadow_light_range"]
                                 and f["shadow_bob_range"] > 2.0)

    # light drift: light_x_world = pivot_x + (SUN_Y - pivot_y) * cot(phi)
    lx = 10.0 + (20.5 - 10.5) / np.tan(np.radians(light[:FRAMES]))
    f["light_x_range"] = float(lx.max() - lx.min())

    # physics labels
    f["gamma"] = cfg.damping_s_inv
    f["beta"] = cfg.quad_drag_beta
    f["g"] = cfg.gravity_m_s2
    f["driven"] = cfg.drive_amp_rad_s2 > 0
    f["quadratic"] = cfg.damping_model == "quadratic"
    f["damped_meaningfully"] = (cfg.damping_model == "linear" and cfg.damping_s_inv >= 0.1) or \
                               (cfg.damping_model == "quadratic" and cfg.quad_drag_beta >= 0.05)
    f["nearly_undamped"] = (not f["driven"] and cfg.damping_model == "linear"
                            and cfg.damping_s_inv < 0.05)

    # eventual fate over a long horizon (5 min). Only unambiguous cases are
    # allowed; weakly-damped linear cases in between return None -> skip.
    if f["driven"]:
        f["eventual"] = "keep swinging in a rhythmic grow-and-shrink pattern"
    elif f["quadratic"] and f["beta"] >= 0.02:
        f["eventual"] = "gradually settle until it hangs nearly still at the bottom"
    elif cfg.damping_model == "linear":
        if f["gamma"] >= 0.10:
            f["eventual"] = "gradually settle until it hangs nearly still at the bottom"
        elif f["gamma"] <= 0.01:
            f["eventual"] = "still be swinging at almost the same height"
        else:
            f["eventual"] = None                            # ambiguous band -> question skipped
    else:
        f["eventual"] = None
    return f


# ---------------------------------------------------------------------------
# Template machinery
# ---------------------------------------------------------------------------

def _mc(qtype, qtext, correct, wrongs, rng, concept):
    opts = [correct] + list(wrongs)[:3]
    order = rng.permutation(len(opts))
    shuffled = [opts[i] for i in order]
    return {
        "question_type": qtype,
        "question": qtext,
        "choices": shuffled,
        "answer": int(shuffled.index(correct)),
        "concept": concept,
    }


def _de(qtext, answer, concept):
    return {"question_type": "descriptive", "question": qtext,
            "answer": str(answer), "concept": concept}


# Each template returns a question dict or None (gate not satisfied) ----------

def q_release_side(f, rng):
    if abs(f["theta0"]) < np.radians(5):
        return None
    side = f["release_side"]
    q = rng.choice([
        "Which side of the pivot was the bob released from at the start of the video?",
        "At the beginning of the clip, is the bob let go from the left side or the right side?",
    ])
    return _de(q, side, "initial conditions")


def q_bottom_passes(f, rng):
    n = f["num_bottom_passes"]
    if n < 1:
        return None
    q = rng.choice([
        "How many times does the bob pass directly below the pivot (through the lowest point) during the clip?",
        "Count the crossings of the bob through the bottom of its arc, directly under the pivot.",
    ])
    return _de(q, n, "oscillation counting")


def q_end_motion(f, rng):
    q = rng.choice([
        "During the final second of the video, how is the bob moving?",
        "In the last second of the clip, what is the bob's motion like?",
    ])
    if f["nearly_still_end"]:
        return _mc("descriptive", q, "it has almost stopped and hangs near the bottom",
                   ["it is swinging strongly toward the left",
                    "it is swinging strongly toward the right",
                    "it is speeding up"], rng, "reading motion state")
    if f["last_sec_dir"] is None:
        return None
    wrong = "left" if f["last_sec_dir"] == "right" else "right"
    return _mc("descriptive", q, f"swinging to the {f['last_sec_dir']}",
               [f"swinging to the {wrong}", "sitting motionless at the bottom",
                "moving straight upward"], rng, "reading motion state")


def q_envelope(f, rng):
    r = f.get("env_ratio")
    if r is None:
        return None
    q = rng.choice([
        "Comparing the end of the video with the start, the pendulum's swing has ___",
        "How does the size of the swings at the end of the clip compare with the first swing?",
    ])
    if r < 0.30:
        ans, bad = "much smaller, having mostly died out", ["grown noticeably larger",
                                                            "stayed almost the same"]
    elif r < 0.65:
        ans, bad = "noticeably smaller but still swinging", ["mostly died out and stopped",
                                                             "grown noticeably larger"]
    elif r > 0.85:
        ans, bad = "stayed almost the same", ["mostly died out and stopped",
                                              "grown much larger"]
    else:
        return None                                              # ambiguous band
    return _mc("descriptive", q, ans, bad, rng, "damping envelope")


def q_shadow_direction(f, rng):
    q = "When the bob swings toward the left, in which direction does its shadow on the ground move?"
    return _de(q, "to the left as well - the shadow always moves in the same direction as the bob",
               "light-projection geometry")


def q_light_motion(f, rng, cfg):
    a = cfg.light_amplitude_deg
    q = "Does the orange light source stay in one place, or drift back and forth during the video?"
    if a >= 7.0:
        return _de(q, "it drifts back and forth", "light kinematics")
    if a <= 2.0:
        return _de(q, "it stays essentially in one place", "light kinematics")
    return None


def q_energy_shape(f, rng):
    q = "Over the full clip, the pendulum's total mechanical energy ___"
    ans = f["energy_shape"]
    opts = ["steadily decreases", "stays nearly constant", "repeatedly rises and falls"]
    return _mc("descriptive", q, ans, [o for o in opts if o != ans], rng, "energy bookkeeping")


def q_fastest_point(f, rng):
    q = "At which part of its arc is the bob moving fastest?"
    return _de(q, "at the bottom of the arc, directly below the pivot",
               "energy conversion")


# --- explanatory ------------------------------------------------------------

def q_why_decay(f, rng):
    if f.get("env_ratio", 1.0) > 0.8:
        return None
    q = rng.choice([
        "The swings get smaller and smaller as the video goes on. What is responsible for this?",
        "Why does the pendulum's swing shrink over time?",
    ])
    return _mc("explanatory", q,
               "friction / air resistance slowly drains the pendulum's energy",
               ["gravity is getting weaker over time",
                "the rod is slowly stretching longer",
                "the moving light is pushing on the bob"], rng, "dissipative force")


def q_why_stop_top(f, rng):
    if f["n_peaks_video"] < 2:
        return None
    q = "Why does the bob momentarily come to rest at the highest point of every swing?"
    return _mc("explanatory", q,
               "it is reversing direction there, so its velocity passes through zero (all energy is potential)",
               ["gravity is weakest at the top of the arc",
                "the string goes slack at the top",
                "air resistance is strongest at the top"], rng, "turning point kinematics")


def q_shadow_cause(f, rng):
    if not f["bob_dominates_shadow"]:
        return None
    q = "The long shadow on the ground sweeps back and forth across the floor. What is mainly responsible?"
    return _mc("explanatory", q,
               "the bob swinging back and forth changes where the light is blocked",
               ["the light source drifting back and forth",
                "the camera moving sideways",
                "the ground tilting from side to side"], rng, "shadow causality")


def q_no_return(f, rng):
    if not (f.get("env_ratio", 1) < 0.7 and f["damped_meaningfully"]
            and f.get("later_peak_max_ratio", 9) < 0.97):
        return None
    q = "After the first swing, the bob never again reaches the height it was released from. Why not?"
    return _mc("explanatory", q,
               "damping removes mechanical energy every swing, so there is not enough energy to climb back",
               ["gravity pulls harder when the bob moves upward",
                "the rod quietly became shorter",
                "the light absorbed the extra height"], rng, "energy conservation loss")


def q_fastest_bottom_why(f, rng):
    q = "Why does the bob travel fastest at the bottom of its arc and slowest at the top?"
    return _mc("explanatory", q,
               "gravity converts potential energy to kinetic as the bob descends, and back as it climbs",
               ["friction pushes the bob along the bottom of the arc",
                "the rod pulls extra hard at the bottom",
                "the shadow drags the bob faster at the bottom"], rng, "energy exchange")


def q_beats_cause(f, rng):
    if not f["driven"] or f["energy_shape"] != "repeatedly rises and falls":
        return None
    q = "The swings repeatedly grow and shrink again, even though the pendulum is damped. What is causing this pattern?"
    return _mc("explanatory", q,
               "a rhythmic external push is repeatedly feeding energy in and sometimes fighting the swing",
               ["gravity pulsing stronger and weaker",
                "the rod sliding up and down at the pivot",
                "the shadow periodically pulling the bob"], rng, "driven oscillator")


def q_large_angle_slow(f, rng):
    if abs(np.degrees(f["theta0"])) < 30:
        return None
    q = ("Compared with the textbook small-angle formula T = 2*pi*sqrt(L/g), this pendulum's real "
         "period is slightly longer. Why?")
    return _mc("explanatory", q,
               "at large swing angles the restoring force follows sin(theta), which is weaker than the small-angle approximation assumes",
               ["a bigger swing makes friction larger, slowing the clock",
                "the bob leaves the circular path at large angles",
                "gravity points slightly sideways at large angles"], rng, "nonlinear period")


# --- predictive -------------------------------------------------------------

def q_next_peak(f, rng):
    r = f.get("next_peak_ratio")
    if r is None:
        return None
    if r < 0.97:
        ans, bad = "lower than the previous peak", ["higher than the previous peak",
                                                    "exactly the same height"]
    elif r > 1.03:
        ans, bad = "higher than the previous peak", ["lower than the previous peak",
                                                     "exactly the same height"]
    else:
        ans, bad = "about the same height as the previous peak", ["much lower",
                                                                  "much higher"]
    q = "If the video kept playing, the next swing peak would be ___"
    return _mc("predictive", q, ans, bad, rng, "extrapolating envelope")


def q_eventual(f, rng):
    if f["eventual"] is None:
        return None
    q = "If we kept filming this same pendulum for another five minutes, it would most likely ___"
    wrongs = ["swing higher and higher until it flips over the top",
              "start moving the light source around",
              "suddenly freeze upright at the top"]
    return _mc("predictive", q, f["eventual"], wrongs, rng, "long-term dynamics")


def q_after_direction(f, rng):
    if f["nearly_still_end"]:
        ans = "barely moving at all, hanging near the bottom"
        wrongs = ["moving to the left", "moving to the right",
                  "moving straight up off the screen"]
    elif f["after_dir"] is not None:
        ans = f"moving to the {f['after_dir']}"
        wrong = "left" if f["after_dir"] == "right" else "right"
        wrongs = [f"moving to the {wrong}",
                  "moving straight up off the screen",
                  "instantly stopped mid-air"]
    else:
        return None
    q = "In the first instant after the video ends, the bob will be ___"
    return _mc("predictive", q, ans, wrongs[:3], rng, "state extrapolation")


def q_cross_next_second(f, rng):
    v = f.get("crossing_next_second")
    if v is None:
        return None
    q = "If the video continued for one more second, would the bob pass through the lowest point of its arc?"
    if v:
        return _mc("predictive", q, "yes, it would swing through the bottom within that second",
                   ["no, it would turn around before reaching the bottom",
                    "no, it would have already stopped by then"], rng, "event timing")
    return _mc("predictive", q, "no, it would not reach the bottom within that second",
               ["yes, it would definitely swing through the bottom",
                "yes, it would pass through the bottom twice"], rng, "event timing")


def q_shadow_future(f, rng):
    q = "If the video kept playing, the shadow on the ground would ___"
    if f["nearly_still_end"]:
        ans = "settle to one spot and almost stop moving"
        bad = ["keep sweeping back and forth widely",
               "swing faster and faster",
               "lift off the ground"]
    else:
        ans = "keep sweeping back and forth as the bob keeps swinging"
        bad = ["freeze in place immediately",
               "start sweeping in the opposite phase to the bob",
               "shrink to a point and vanish"]
    return _mc("predictive", q, ans, bad, rng, "coupled prediction")


# --- counterfactual ---------------------------------------------------------

def q_cf_no_damping(f, rng):
    if not (f["gamma"] >= 0.15 and not f["quadratic"] and not f["driven"]):
        return None
    q = "If all friction and air resistance were removed, this pendulum would ___"
    return _mc("counterfactual", q,
               "return to its release height every swing and never slow down",
               ["still run out of energy and stop",
                "swing even higher than its release point",
                "immediately hang straight down"], rng, "conservation of energy")


def q_cf_add_damping(f, rng):
    if not f["nearly_undamped"]:
        return None
    q = "If this pendulum had the usual air resistance (like a real one in a room), its motion would ___"
    return _mc("counterfactual", q,
               "visibly shrink swing by swing as energy drained away",
               ["keep swinging at exactly the same height forever",
                "swing higher with every pass",
                "reverse direction randomly"], rng, "dissipative force")


def q_cf_gravity(f, rng):
    q = "If this pendulum were moved somewhere with gravity twice as strong, each full swing would take ___"
    return _mc("counterfactual", q, "less time (the pendulum would swing faster)",
               ["more time (the pendulum would swing slower)",
                "exactly the same time"], rng, "g dependence of period")


def q_cf_rod(f, rng):
    q = "If the rod were twice as long but the bob was released the same way, the pendulum would swing ___"
    return _mc("counterfactual", q, "more slowly, with a longer period",
               ["more quickly, with a shorter period",
                "with exactly the same period"], rng, "length dependence of period")


def q_cf_mass(f, rng):
    q = "If the bob were replaced by one twice as heavy (same size, same release), the time of each swing would ___"
    return _mc("counterfactual", q, "stay the same - the period does not depend on the bob's mass",
               ["become twice as long",
                "become half as long"], rng, "mass independence")


def q_cf_light(f, rng):
    q = "If the orange light were moved to a different place or made brighter, the pendulum's swinging would ___"
    return _mc("counterfactual", q,
               "stay exactly the same - the light only changes the illumination and the shadow",
               ["slow down, because brighter light heats the air",
                "speed up, because the light pushes the bob",
                "become irregular, because of the moving shadow"], rng, "interventional invariance")


def q_cf_steeper(f, rng):
    q = ("If the bob had been released from rest at an even steeper angle, the time needed for one "
         "complete swing would be ___")
    return _mc("counterfactual", q,
               "slightly longer - a pendulum's period grows a little as the swing angle increases",
               ["slightly shorter - released higher it has more speed",
                "exactly identical - the period never depends on amplitude"], rng, "amplitude-period coupling")


def q_cf_double_damping(f, rng):
    if not (f["gamma"] >= 0.2 and not f["quadratic"] and not f["driven"]):
        return None
    q = "If the air resistance were twice as strong, the swings would die out ___"
    return _mc("counterfactual", q, "about twice as fast, in roughly half the time",
               ["at exactly the same rate",
                "much more slowly",
                "not at all"], rng, "damping rate scaling")


def q_cf_drive_off(f, rng):
    if not f["driven"]:
        return None
    q = "If the rhythmic pushing force acting on this pendulum suddenly stopped, the pendulum would ___"
    return _mc("counterfactual", q,
               "gradually calm down and settle toward the bottom as damping drains its energy",
               ["keep its grow-and-shrink beat pattern forever",
                "stop dead within one frame",
                "speed up without the push holding it back"], rng, "forced vs free response")


def q_cf_vacuum_quad(f, rng):
    if not f["quadratic"]:
        return None
    q = "This pendulum swings through air that resists it harder the faster it moves. In a perfect vacuum it would ___"
    return _mc("counterfactual", q,
               "keep swinging at almost the same height for a very long time",
               ["stop even faster than in air",
                "swing at double the speed",
                "trace the same shrinking envelope"], rng, "drag law")


def q_cf_gravity_planet(f, rng):
    g = f["g"]
    if abs(g - 9.81) < 1e-6:
        return None
    q = ("This pendulum is swinging in gravity different from Earth's. If the identical pendulum "
         "were brought to Earth (g = 9.81 m/s^2), its swings would take ___")
    if g < 9.81:
        return _mc("counterfactual", q, "less time - Earth gravity would make it swing faster",
                   ["more time - it would swing more slowly",
                    "exactly the same time"], rng, "g dependence of period")
    return _mc("counterfactual", q, "more time - Earth gravity would make it swing more slowly",
               ["less time - it would swing faster",
                "exactly the same time"], rng, "g dependence of period")


# ---------------------------------------------------------------------------
# Per-video assembly
# ---------------------------------------------------------------------------

DESCRIPTIVE = [q_release_side, q_bottom_passes, q_end_motion, q_envelope,
               q_shadow_direction, q_light_motion, q_energy_shape, q_fastest_point]
EXPLANATORY = [q_why_decay, q_why_stop_top, q_shadow_cause, q_no_return,
               q_fastest_bottom_why, q_beats_cause, q_large_angle_slow]
PREDICTIVE = [q_next_peak, q_eventual, q_after_direction, q_cross_next_second,
              q_shadow_future]
COUNTERFACTUAL = [q_cf_no_damping, q_cf_add_damping, q_cf_gravity, q_cf_rod,
                  q_cf_mass, q_cf_light, q_cf_steeper, q_cf_double_damping,
                  q_cf_drive_off, q_cf_vacuum_quad, q_cf_gravity_planet]

POOLS = {"descriptive": DESCRIPTIVE, "explanatory": EXPLANATORY,
         "predictive": PREDICTIVE, "counterfactual": COUNTERFACTUAL}


def build_questions(cfg, A, rng, per_type=2, target=8):
    """Returns a list of question dicts (CLEVRER-style) for one video."""
    f = analyze(cfg, A)
    candidates = {}
    for qtype, pool in POOLS.items():
        cands = []
        for tpl in pool:
            try:
                q = tpl(f, rng, cfg) if tpl is q_light_motion else tpl(f, rng)
            except Exception:
                q = None
            if q is not None:
                cands.append(q)
        candidates[qtype] = cands

    chosen = []
    for qtype in POOLS:
        c = candidates[qtype]
        idx = rng.permutation(len(c))[:per_type]
        chosen += [c[i] for i in idx]

    # top up toward target from remaining candidates of any type
    leftovers = [q for qtype in POOLS for q in candidates[qtype] if q not in chosen]
    j = 0
    while len(chosen) < target and j < len(leftovers):
        if leftovers[j]["question"] not in {c["question"] for c in chosen}:
            chosen.append(leftovers[j])
        j += 1

    out = []
    for i, q in enumerate(chosen):
        q = dict(q)
        q["question_id"] = i
        out.append(q)
    return out