1 files changed, 129 insertions, 0 deletions
diff --git a/worldalign/synth_probes.py b/worldalign/synth_probes.py
new file mode 100644
index 0000000..180b96a
--- /dev/null
+++ b/worldalign/synth_probes.py
@@ -0,0 +1,129 @@
+"""Ground-truth factor probes for synthetic-world representations.
+
+Linear probes from a representation to the discrete scene factors measure
+which world variables survive the encoder and its pooling. Probes are
+fitted per modality on that modality's own training split and evaluated
+on held-out scenes; scene truth is generator metadata, so this is a
+within-modality diagnostic, not a cross-modal signal.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .common import read_json, write_json
+from .synth_world import COLORS, SHAPES
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-dir", default="artifacts/synth_v0")
+    parser.add_argument("--features", required=True)
+    parser.add_argument(
+        "--side", choices=["vision", "text"], required=True,
+        help="Selects the training split whose scenes fit the probes.",
+    )
+    parser.add_argument("--ridge", type=float, default=1.0)
+    parser.add_argument("--output", required=True)
+    return parser.parse_args()
+
+
+def scene_targets(scene: dict) -> dict[str, np.ndarray | float]:
+    color_presence = np.zeros(len(COLORS))
+    shape_presence = np.zeros(len(SHAPES))
+    for group in scene["groups"]:
+        color_presence[list(COLORS).index(group["color"])] = 1.0
+        shape_presence[SHAPES.index(group["shape"])] = 1.0
+    return {
+        "color_presence": color_presence,
+        "shape_presence": shape_presence,
+        "group_count": float(len(scene["groups"])),
+        "object_total": float(sum(g["count"] for g in scene["groups"])),
+        "relation_count": float(len(scene["relations"])),
+    }
+
+
+def ridge_fit(
+    x: np.ndarray, y: np.ndarray, ridge: float
+) -> tuple[np.ndarray, np.ndarray]:
+    x = np.concatenate([x, np.ones((len(x), 1))], axis=1)
+    gram = x.T @ x + ridge * np.eye(x.shape[1])
+    weights = np.linalg.solve(gram, x.T @ y)
+    return weights, x @ weights
+
+
+def evaluate(
+    features_train: np.ndarray,
+    features_test: np.ndarray,
+    train_targets: np.ndarray,
+    test_targets: np.ndarray,
+    ridge: float,
+    binary: bool,
+) -> dict:
+    weights, _ = ridge_fit(features_train, train_targets, ridge)
+    prediction = (
+        np.concatenate([features_test, np.ones((len(features_test), 1))], axis=1)
+        @ weights
+    )
+    if binary:
+        accuracy = float(((prediction > 0.5) == (test_targets > 0.5)).mean())
+        balanced = []
+        for column in range(test_targets.shape[1]):
+            truth = test_targets[:, column] > 0.5
+            if truth.any() and (~truth).any():
+                hit = (prediction[:, column] > 0.5) == truth
+                balanced.append(
+                    (hit[truth].mean() + hit[~truth].mean()) / 2.0
+                )
+        return {
+            "accuracy": accuracy,
+            "balanced_accuracy": float(np.mean(balanced)),
+        }
+    residual = prediction[:, 0] - test_targets[:, 0]
+    variance = test_targets[:, 0].var()
+    return {
+        "r2": float(1.0 - residual.var() / max(variance, 1e-9)),
+        "mae": float(np.abs(residual).mean()),
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    manifest = read_json(Path(args.data_dir, "manifest.json"))
+    scenes = read_json(Path(args.data_dir, "scenes.private.json"))["scenes"]
+    state = torch.load(args.features, map_location="cpu", weights_only=False)
+    lookup = {int(row): i for i, row in enumerate(state["rows"])}
+    features = state["features"].float().numpy()
+
+    split_key = "vision_only_train" if args.side == "vision" else "text_only_train"
+    train_rows = [row for row in manifest[split_key] if row in lookup][:8000]
+    test_rows = [row for row in manifest["test"] if row in lookup]
+
+    x_train = features[[lookup[row] for row in train_rows]]
+    x_test = features[[lookup[row] for row in test_rows]]
+    report = {"features": args.features, "side": args.side, "targets": {}}
+    for name in ("color_presence", "shape_presence"):
+        y_train = np.stack([scene_targets(scenes[row])[name] for row in train_rows])
+        y_test = np.stack([scene_targets(scenes[row])[name] for row in test_rows])
+        report["targets"][name] = evaluate(
+            x_train, x_test, y_train, y_test, args.ridge, binary=True
+        )
+    for name in ("group_count", "object_total", "relation_count"):
+        y_train = np.array(
+            [[scene_targets(scenes[row])[name]] for row in train_rows]
+        )
+        y_test = np.array([[scene_targets(scenes[row])[name]] for row in test_rows])
+        report["targets"][name] = evaluate(
+            x_train, x_test, y_train, y_test, args.ridge, binary=False
+        )
+    write_json(args.output, report)
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()