auto-sync 2026-07-04T00:30:01Z workspace (part 6)

workspace/scripts/eval_learned_dominance_selector.py

@@ -110,6 +110,15 @@ def main(argv: list[str] | None = None) -> int:
             "candidate success to prioritize the lexicographic success/progress utility."
         ),
     )
+    parser.add_argument(
+        "--threshold-scope",
+        choices=("global", "task"),
+        default="global",
+        help=(
+            "Calibrate one global execute/fallback threshold or a Mondrian "
+            "threshold per visible task_id bucket using calibration rows only."
+        ),
+    )
     parser.add_argument("--bootstrap-samples", type=int, default=1000)
     args = parser.parse_args(argv)
 
@@ -158,7 +167,11 @@ def main(argv: list[str] | None = None) -> int:
         target=args.target,
         source_evidence=source_evidence,
     )
-    best = _fit_select_ridge(calibration_dataset, lambdas=lambdas)
+    best = _fit_select_ridge(
+        calibration_dataset,
+        lambdas=lambdas,
+        threshold_scope=args.threshold_scope,
+    )
     eval_cases = _evaluate_dataset(eval_dataset, best["weights"], best["mean"], best["std"], tau=best["tau"])
     calibration_cases = _evaluate_dataset(
         calibration_dataset,
@@ -192,6 +205,7 @@ def main(argv: list[str] | None = None) -> int:
         "k": args.k,
         "feature_set": args.feature_set,
         "target": args.target,
+        "threshold_scope": args.threshold_scope,
         "chart_feature_mode": chart_feature_mode,
         "feature_names": _feature_names(args.feature_set),
         "ridge_lambdas": lambdas,
@@ -235,7 +249,7 @@ def main(argv: list[str] | None = None) -> int:
     (out_dir / "train.log").write_text(
         "trained ridge dominance calibrator on calibration measured rows only\n"
         f"selected_lambda={best['lambda']}\n"
-        f"tau={best['tau']:.6f}\n"
+        f"tau={_format_tau(best['tau'])}\n"
         f"calibration_input={args.calibration_input}\n"
     )
     (out_dir / "eval.log").write_text(
@@ -609,7 +623,12 @@ def _target_value(
     raise ValueError(f"unknown target: {target}")
 
 
-def _fit_select_ridge(dataset: dict[str, Any], *, lambdas: list[float]) -> dict[str, Any]:
+def _fit_select_ridge(
+    dataset: dict[str, Any],
+    *,
+    lambdas: list[float],
+    threshold_scope: str = "global",
+) -> dict[str, Any]:
     samples = dataset["samples"]
     if not samples:
         raise ValueError("cannot fit learned dominance selector without candidates")
@@ -627,7 +646,11 @@ def _fit_select_ridge(dataset: dict[str, Any], *, lambdas: list[float]) -> dict[
         penalty[0, 0] = 0.0
         weights = np.linalg.pinv(x_norm.T @ x_norm + penalty) @ (x_norm.T @ y)
         predictions = x_norm @ weights
-        tau, selection = _choose_tau(dataset, predictions)
+        tau, selection = _choose_thresholds(
+            dataset,
+            predictions,
+            threshold_scope=threshold_scope,
+        )
         key = (
             float(selection["selected_success"]),
             float(selection["selected_utility"]),
@@ -648,6 +671,27 @@ def _fit_select_ridge(dataset: dict[str, Any], *, lambdas: list[float]) -> dict[
     return best
 
 
+def _choose_thresholds(
+    dataset: dict[str, Any],
+    predictions: np.ndarray,
+    *,
+    threshold_scope: str,
+) -> tuple[float | dict[str, float], dict[str, float | None]]:
+    if threshold_scope == "global":
+        return _choose_tau(dataset, predictions)
+    if threshold_scope != "task":
+        raise ValueError(f"unknown threshold_scope: {threshold_scope}")
+    global_tau, _global_summary = _choose_tau(dataset, predictions)
+    task_taus: dict[str, float] = {"__global__": float(global_tau)}
+    for task_id, row_indices in _rows_by_task(dataset).items():
+        subset = _subset_dataset_rows(dataset, row_indices)
+        subset_predictions = _predictions_for_subset(dataset, predictions, subset)
+        task_tau, _task_summary = _choose_tau(subset, subset_predictions)
+        task_taus[task_id] = float(task_tau)
+    cases = _evaluate_predictions(dataset, predictions, tau=task_taus)
+    return task_taus, _simple_summary(cases)
+
+
 def _choose_tau(dataset: dict[str, Any], predictions: np.ndarray) -> tuple[float, dict[str, float | None]]:
     candidates = sorted(float(value) for value in predictions)
     thresholds = [min(candidates) - 1.0, *candidates, max(candidates) + 1.0]
@@ -676,7 +720,7 @@ def _evaluate_dataset(
     mean: np.ndarray,
     std: np.ndarray,
     *,
-    tau: float,
+    tau: float | dict[str, float],
 ) -> list[dict[str, Any]]:
     x = np.stack([sample["feature"] for sample in dataset["samples"]], axis=0)
     x_norm = (x - mean) / std
@@ -684,14 +728,20 @@ def _evaluate_dataset(
     return _evaluate_predictions(dataset, x_norm @ weights, tau=tau)
 
 
-def _evaluate_predictions(dataset: dict[str, Any], predictions: np.ndarray, *, tau: float) -> list[dict[str, Any]]:
+def _evaluate_predictions(
+    dataset: dict[str, Any],
+    predictions: np.ndarray,
+    *,
+    tau: float | dict[str, float],
+) -> list[dict[str, Any]]:
     samples = dataset["samples"]
     rows: list[dict[str, Any]] = []
     for row_index, sample_indices in sorted(dataset["by_row"].items()):
         best_index = max(sample_indices, key=lambda index: float(predictions[index]))
         sample = samples[best_index]
         predicted_margin = float(predictions[best_index])
-        execute = predicted_margin > float(tau)
+        row_tau = _tau_for_sample(sample, tau)
+        execute = predicted_margin > row_tau
         selected_utility = (
             float(sample["candidate_utility"]) if execute else float(sample["base_utility"])
         )
@@ -708,7 +758,7 @@ def _evaluate_predictions(dataset: dict[str, Any], predictions: np.ndarray, *, t
                 "train_seed": sample["train_seed"],
                 "selected_candidate_index": int(sample["candidate_index"]),
                 "predicted_margin": predicted_margin,
-                "tau": float(tau),
+                "tau": row_tau,
                 "execute_generated": float(execute),
                 "coverage": float(execute),
                 "fallback_rate": float(not execute),
@@ -736,6 +786,56 @@ def _evaluate_predictions(dataset: dict[str, Any], predictions: np.ndarray, *, t
     return rows
 
 
+def _tau_for_sample(sample: dict[str, Any], tau: float | dict[str, float]) -> float:
+    if isinstance(tau, dict):
+        return float(tau.get(str(sample.get("task_id", "")), tau.get("__global__", 0.0)))
+    return float(tau)
+
+
+def _rows_by_task(dataset: dict[str, Any]) -> dict[str, list[int]]:
+    grouped: dict[str, list[int]] = defaultdict(list)
+    for row_index, sample_indices in dataset["by_row"].items():
+        if not sample_indices:
+            continue
+        task_id = str(dataset["samples"][sample_indices[0]].get("task_id", "unknown"))
+        grouped[task_id].append(int(row_index))
+    return grouped
+
+
+def _subset_dataset_rows(dataset: dict[str, Any], row_indices: list[int]) -> dict[str, Any]:
+    wanted = set(int(row) for row in row_indices)
+    old_to_new_sample: dict[int, int] = {}
+    samples: list[dict[str, Any]] = []
+    by_row: dict[int, list[int]] = {}
+    for old_row in sorted(wanted):
+        if old_row not in dataset["by_row"]:
+            continue
+        new_row = len(by_row)
+        by_row[new_row] = []
+        for old_sample_index in dataset["by_row"][old_row]:
+            old_to_new_sample[old_sample_index] = len(samples)
+            sample = dict(dataset["samples"][old_sample_index])
+            sample["row_index"] = new_row
+            by_row[new_row].append(len(samples))
+            samples.append(sample)
+    return {
+        "samples": samples,
+        "by_row": by_row,
+        "num_rows": len(by_row),
+        "_old_to_new_sample": old_to_new_sample,
+    }
+
+
+def _predictions_for_subset(
+    dataset: dict[str, Any],
+    predictions: np.ndarray,
+    subset: dict[str, Any],
+) -> np.ndarray:
+    old_to_new = subset.get("_old_to_new_sample", {})
+    ordered_old_indices = sorted(old_to_new, key=lambda old: old_to_new[old])
+    return np.asarray([float(predictions[old]) for old in ordered_old_indices], dtype=float)
+
+
 def _simple_summary(rows: list[dict[str, Any]]) -> dict[str, float | None]:
     keys = [
         "base_success",
@@ -819,7 +919,8 @@ def _report(metrics: dict[str, Any]) -> str:
         f"Calibration rows: `{metrics['num_calibration_rows']}`",
         f"Eval rows: `{metrics['num_eval_rows']}`",
         f"Selected ridge lambda: `{metrics['selected_lambda']}`",
-        f"Tau: `{metrics['tau']:.6f}`",
+        f"Tau: `{_format_tau(metrics['tau'])}`",
+        f"Threshold scope: `{metrics.get('threshold_scope', 'global')}`",
         f"Feature set: `{metrics['feature_set']}`",
         f"Target: `{metrics['target']}`",
         "",
@@ -899,6 +1000,14 @@ def _fmt(value: Any) -> str:
     return f"{float(value):.4f}"
 
 
+def _format_tau(value: Any) -> str:
+    if isinstance(value, dict):
+        return json.dumps({key: round(float(val), 6) for key, val in sorted(value.items())})
+    if isinstance(value, (int, float)) and math.isfinite(float(value)):
+        return f"{float(value):.6f}"
+    return str(value)
+
+
 def _run(command: list[str]) -> str:
     try:
         return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()

workspace/scripts/summarize_ctt_runs.py

@@ -41,6 +41,7 @@ def main(argv: list[str] | None = None) -> int:
     args = parser.parse_args(argv)
 
     rows = []
+    seen_metric_paths: set[Path] = set()
     reproduce = args.run_root / "reproduce_v0" / "metrics.json"
     if reproduce.exists():
         rows.extend(_reproduce_rows(reproduce))
@@ -53,6 +54,7 @@ def main(argv: list[str] | None = None) -> int:
         }:
             continue
         rows.append(_ctt_row(metrics_path))
+        seen_metric_paths.add(metrics_path)
     for metrics_path in sorted(args.run_root.glob("*memory*_val_proxy/metrics.json")):
         rows.append(_positive_memory_row(metrics_path))
     local_atlas = args.run_root / "local_atlas_val_proxy" / "metrics.json"
@@ -65,11 +67,21 @@ def main(argv: list[str] | None = None) -> int:
     if utility.exists():
         rows.append(_utility_row(utility))
     for metrics_path in sorted(args.run_root.glob("ctt_dominance*/metrics.json")):
+        if metrics_path in seen_metric_paths:
+            continue
         rows.append(_dominance_row(metrics_path))
-    for metrics_path in sorted(args.run_root.glob("ctt_learned_dominance*/metrics.json")):
-        rows.append(_dominance_row(metrics_path))
+        seen_metric_paths.add(metrics_path)
+    for pattern in ("ctt_learned_dominance*/metrics.json", "ctt*_learned_dominance*/metrics.json"):
+        for metrics_path in sorted(args.run_root.glob(pattern)):
+            if metrics_path in seen_metric_paths:
+                continue
+            rows.append(_dominance_row(metrics_path))
+            seen_metric_paths.add(metrics_path)
     for metrics_path in sorted(args.run_root.glob("ctt*_nonlinear_dominance*/metrics.json")):
+        if metrics_path in seen_metric_paths:
+            continue
         rows.append(_dominance_row(metrics_path))
+        seen_metric_paths.add(metrics_path)
 
     args.out_csv.parent.mkdir(parents=True, exist_ok=True)
     with args.out_csv.open("w", newline="") as handle: