Add CTT source-evidence dominance features

workspace/scripts/eval_learned_dominance_selector.py

@@ -36,6 +36,36 @@ BASIC_FEATURE_NAMES = [
 ]
 FEATURE_NAMES = BASIC_FEATURE_NAMES
 CONTEXT_HASH_WIDTH = 8
+SOURCE_EVIDENCE_NAMES = [
+    "source_chart_found",
+    "source_positive_count_log",
+    "source_negative_count_log",
+    "source_nonbase_count_log",
+    "source_positive_rate",
+    "source_best_delta_utility",
+    "source_mean_delta_utility",
+    "source_best_utility",
+    "source_mean_utility",
+    "source_best_success",
+    "source_mean_success",
+    "source_best_progress",
+    "source_mean_progress",
+    "source_positive_safety_known_rate",
+    "source_positive_unsafe_rate_known",
+    "generated_to_source_positive_min_rms",
+    "generated_to_source_negative_min_rms",
+    "generated_source_pos_closer_than_neg",
+]
+FEATURE_SET_CHOICES = (
+    "basic",
+    "tangent",
+    "context",
+    "context_tangent",
+    "source_evidence",
+    "tangent_source_evidence",
+    "context_source_evidence",
+    "context_tangent_source_evidence",
+)
 
 
 def main(argv: list[str] | None = None) -> int:
@@ -49,6 +79,15 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--calibration-target-index", type=Path, required=True)
     parser.add_argument("--eval-input", type=Path, required=True)
     parser.add_argument("--eval-target-index", type=Path, required=True)
+    parser.add_argument(
+        "--source-index",
+        type=Path,
+        default=None,
+        help=(
+            "Train split chart index used for source-evidence features. "
+            "Defaults to --calibration-target-index."
+        ),
+    )
     parser.add_argument(
         "--checkpoint-template",
         default="runs/ctt_residual_full_seed{seed}/model.pt",
@@ -58,7 +97,7 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--ridge-lambdas", default="0,0.01,0.1,1,10,100")
     parser.add_argument(
         "--feature-set",
-        choices=("basic", "tangent", "context", "context_tangent"),
+        choices=FEATURE_SET_CHOICES,
         default="basic",
         help="Deployment-visible feature family for candidate-level dominance fitting.",
     )
@@ -96,6 +135,10 @@ def main(argv: list[str] | None = None) -> int:
         args.eval_target_index,
         chart_feature_mode=chart_feature_mode,
     )
+    source_index_path = _resolve_index_path(args.source_index or args.calibration_target_index)
+    source_evidence, source_index = (
+        _source_evidence_map(source_index_path) if _uses_source_evidence(args.feature_set) else ({}, {})
+    )
 
     calibration_dataset = _candidate_dataset(
         calibration_rows,
@@ -104,6 +147,7 @@ def main(argv: list[str] | None = None) -> int:
         k=args.k,
         feature_set=args.feature_set,
         target=args.target,
+        source_evidence=source_evidence,
     )
     eval_dataset = _candidate_dataset(
         eval_rows,
@@ -112,6 +156,7 @@ def main(argv: list[str] | None = None) -> int:
         k=args.k,
         feature_set=args.feature_set,
         target=args.target,
+        source_evidence=source_evidence,
     )
     best = _fit_select_ridge(calibration_dataset, lambdas=lambdas)
     eval_cases = _evaluate_dataset(eval_dataset, best["weights"], best["mean"], best["std"], tau=best["tau"])
@@ -157,12 +202,15 @@ def main(argv: list[str] | None = None) -> int:
         "feature_std": best["std"].tolist(),
         "calibration_input": str(args.calibration_input),
         "eval_input": str(args.eval_input),
+        "source_index": str(source_index_path) if _uses_source_evidence(args.feature_set) else None,
         "data_hash": eval_index.get("content_hash"),
         "split_hash": eval_index.get("split_hash"),
         "calibration_target_content_hash": calibration_index.get("content_hash"),
         "calibration_target_split_hash": calibration_index.get("split_hash"),
         "eval_target_content_hash": eval_index.get("content_hash"),
         "eval_target_split_hash": eval_index.get("split_hash"),
+        "source_content_hash": source_index.get("content_hash"),
+        "source_split_hash": source_index.get("split_hash"),
         "num_calibration_rows": len(calibration_rows),
         "num_eval_rows": len(eval_rows),
         "num_calibration_candidates": len(calibration_dataset["samples"]),
@@ -217,7 +265,9 @@ def _candidate_dataset(
     k: int,
     feature_set: str = "basic",
     target: str = "utility_margin",
+    source_evidence: dict[str, dict[str, Any]] | None = None,
 ) -> dict[str, Any]:
+    source_evidence = source_evidence or {}
     samples: list[dict[str, Any]] = []
     by_row: dict[int, list[int]] = defaultdict(list)
     for row_index, row in enumerate(rows):
@@ -231,6 +281,7 @@ def _candidate_dataset(
         tangents = row.get("generated_tangents", [])[:k]
         candidate_types = row.get("candidate_types", [])[:k]
         source_task_ids = row.get("candidate_source_task_ids", [])[:k]
+        source_chart_ids = row.get("candidate_source_chart_ids", [])[:k]
         if not scores or len(scores) != len(utilities):
             continue
         score_mean = sum(scores) / len(scores)
@@ -255,6 +306,14 @@ def _candidate_dataset(
                     else "",
                 },
                 tangent=tangent,
+                source_evidence=_source_evidence_feature(
+                    source_evidence.get(
+                        str(source_chart_ids[candidate_index])
+                        if candidate_index < len(source_chart_ids)
+                        else ""
+                    ),
+                    tangent=tangent,
+                ),
                 num_candidates=len(scores),
                 feature_set=feature_set,
             )
@@ -308,12 +367,15 @@ def _feature_names(feature_set: str) -> list[str]:
         *[f"tangent_{index:02d}" for index in range(21)],
         *[f"abs_tangent_{index:02d}" for index in range(21)],
     ]
-    if feature_set == "tangent":
-        return [*BASIC_FEATURE_NAMES, *tangent_names]
-    if feature_set == "context":
-        return [*BASIC_FEATURE_NAMES, *context_names]
-    if feature_set == "context_tangent":
-        return [*BASIC_FEATURE_NAMES, *context_names, *tangent_names]
+    names = list(BASIC_FEATURE_NAMES)
+    if _uses_context(feature_set):
+        names.extend(context_names)
+    if _uses_tangent(feature_set):
+        names.extend(tangent_names)
+    if _uses_source_evidence(feature_set):
+        names.extend(SOURCE_EVIDENCE_NAMES)
+    if feature_set in FEATURE_SET_CHOICES:
+        return names
     raise ValueError(f"unknown feature_set: {feature_set}")
 
 
@@ -329,6 +391,7 @@ def _candidate_feature(
     num_candidates: int,
     feature_set: str,
     context: dict[str, Any] | None = None,
+    source_evidence: np.ndarray | None = None,
 ) -> np.ndarray:
     tangent = np.asarray(tangent, dtype=float).reshape(-1)
     if tangent.size < 21:
@@ -354,22 +417,161 @@ def _candidate_feature(
     )
     if feature_set == "basic":
         return basic
-    if feature_set == "tangent":
-        return np.concatenate([basic, tangent.astype(float), np.abs(tangent).astype(float)])
-    if feature_set == "context":
-        return np.concatenate([basic, _context_feature(context or {})])
-    if feature_set == "context_tangent":
-        return np.concatenate(
-            [
-                basic,
-                _context_feature(context or {}),
-                tangent.astype(float),
-                np.abs(tangent).astype(float),
-            ]
-        )
+    parts = [basic]
+    if _uses_context(feature_set):
+        parts.append(_context_feature(context or {}))
+    if _uses_tangent(feature_set):
+        parts.extend([tangent.astype(float), np.abs(tangent).astype(float)])
+    if _uses_source_evidence(feature_set):
+        if source_evidence is None:
+            source_evidence = np.zeros(len(SOURCE_EVIDENCE_NAMES), dtype=float)
+        parts.append(np.asarray(source_evidence, dtype=float).reshape(-1))
+    if feature_set in FEATURE_SET_CHOICES:
+        return np.concatenate(parts)
     raise ValueError(f"unknown feature_set: {feature_set}")
 
 
+def _uses_context(feature_set: str) -> bool:
+    return feature_set in {
+        "context",
+        "context_tangent",
+        "context_source_evidence",
+        "context_tangent_source_evidence",
+    }
+
+
+def _uses_tangent(feature_set: str) -> bool:
+    return feature_set in {
+        "tangent",
+        "context_tangent",
+        "tangent_source_evidence",
+        "context_tangent_source_evidence",
+    }
+
+
+def _uses_source_evidence(feature_set: str) -> bool:
+    return feature_set in {
+        "source_evidence",
+        "tangent_source_evidence",
+        "context_source_evidence",
+        "context_tangent_source_evidence",
+    }
+
+
+def _source_evidence_map(index_path: Path) -> tuple[dict[str, dict[str, Any]], dict[str, Any]]:
+    index_path = _resolve_index_path(index_path)
+    index = json.loads(index_path.read_text())
+    if not index.get("include_outcomes", False):
+        raise SystemExit(f"{index_path} must include train outcomes for source evidence")
+    grouped: dict[str, dict[str, Any]] = {}
+    for shard in index.get("shards", []):
+        shard_path = index_path.parent / shard["path"]
+        with np.load(shard_path, allow_pickle=False) as data:
+            chart_ids = data["chart_id"]
+            labels = data["label"]
+            is_base = data["is_base_branch"]
+            tangents = data["spline_tangent_code"]
+            utilities = data["utility"]
+            delta_utilities = data["delta_utility"]
+            outcomes = data["outcome_vector"]
+            for row in range(chart_ids.shape[0]):
+                if bool(is_base[row]):
+                    continue
+                chart_id = str(chart_ids[row])
+                item = grouped.setdefault(
+                    chart_id,
+                    {
+                        "positive_tangents": [],
+                        "negative_tangents": [],
+                        "num_nonbase": 0,
+                        "positive_delta_utilities": [],
+                        "positive_utilities": [],
+                        "positive_success": [],
+                        "positive_progress": [],
+                        "positive_safety": [],
+                    },
+                )
+                item["num_nonbase"] += 1
+                label = str(labels[row])
+                if label == "positive":
+                    item["positive_tangents"].append(tangents[row].astype("float32"))
+                    item["positive_delta_utilities"].append(float(delta_utilities[row]))
+                    item["positive_utilities"].append(float(utilities[row]))
+                    outcome = np.asarray(outcomes[row], dtype=float)
+                    item["positive_success"].append(float(outcome[0]) if outcome.size > 0 else math.nan)
+                    item["positive_progress"].append(float(outcome[1]) if outcome.size > 1 else math.nan)
+                    item["positive_safety"].append(float(outcome[3]) if outcome.size > 3 else math.nan)
+                elif label == "negative":
+                    item["negative_tangents"].append(tangents[row].astype("float32"))
+    return grouped, index
+
+
+def _source_evidence_feature(source: dict[str, Any] | None, *, tangent: np.ndarray) -> np.ndarray:
+    if not source:
+        return np.zeros(len(SOURCE_EVIDENCE_NAMES), dtype=float)
+    positives = np.asarray(source.get("positive_tangents") or [], dtype=float).reshape(-1, 21)
+    negatives = np.asarray(source.get("negative_tangents") or [], dtype=float).reshape(-1, 21)
+    positive_count = int(positives.shape[0])
+    negative_count = int(negatives.shape[0])
+    nonbase_count = int(source.get("num_nonbase") or (positive_count + negative_count))
+    positive_delta = _clean_array(source.get("positive_delta_utilities") or [])
+    positive_utility = _clean_array(source.get("positive_utilities") or [])
+    positive_success = _clean_array(source.get("positive_success") or [])
+    positive_progress = _clean_array(source.get("positive_progress") or [])
+    positive_safety_raw = np.asarray(source.get("positive_safety") or [], dtype=float)
+    known_safety = positive_safety_raw[np.isfinite(positive_safety_raw)]
+    tangent = np.asarray(tangent, dtype=float).reshape(-1)
+    if tangent.size < 21:
+        tangent = np.pad(tangent, (0, 21 - tangent.size))
+    elif tangent.size > 21:
+        tangent = tangent[:21]
+    pos_dist = _min_rms_distance(tangent, positives)
+    neg_dist = _min_rms_distance(tangent, negatives)
+    return np.asarray(
+        [
+            1.0,
+            math.log1p(positive_count),
+            math.log1p(negative_count),
+            math.log1p(nonbase_count),
+            positive_count / max(1.0, float(nonbase_count)),
+            _safe_max(positive_delta),
+            _safe_mean(positive_delta),
+            _safe_max(positive_utility),
+            _safe_mean(positive_utility),
+            _safe_max(positive_success),
+            _safe_mean(positive_success),
+            _safe_max(positive_progress),
+            _safe_mean(positive_progress),
+            float(known_safety.size) / max(1.0, float(positive_count)),
+            _safe_mean(known_safety),
+            pos_dist,
+            neg_dist,
+            float(pos_dist < neg_dist) if math.isfinite(pos_dist) and math.isfinite(neg_dist) else 0.0,
+        ],
+        dtype=float,
+    )
+
+
+def _clean_array(values: Any) -> np.ndarray:
+    array = np.asarray(values, dtype=float).reshape(-1)
+    return array[np.isfinite(array)]
+
+
+def _safe_mean(values: np.ndarray) -> float:
+    return float(values.mean()) if values.size else 0.0
+
+
+def _safe_max(values: np.ndarray) -> float:
+    return float(values.max()) if values.size else 0.0
+
+
+def _min_rms_distance(tangent: np.ndarray, candidates: np.ndarray) -> float:
+    if candidates.size == 0:
+        return 0.0
+    diff = candidates - tangent.reshape(1, -1)
+    return float(np.sqrt(np.mean(diff * diff, axis=1)).min())
+
+
 def _context_feature(context: dict[str, Any]) -> np.ndarray:
     target_task = str(context.get("target_task_id", ""))
     source_task = str(context.get("source_task_id", ""))
@@ -649,10 +851,12 @@ def _write_provenance(out_dir: Path, args: argparse.Namespace) -> None:
     (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
     hashes = {
         "calibration_input": _sha256(args.calibration_input),
-        "calibration_target_index": _sha256(args.calibration_target_index),
+        "calibration_target_index": _sha256(_resolve_index_path(args.calibration_target_index)),
         "eval_input": _sha256(args.eval_input),
-        "eval_target_index": _sha256(args.eval_target_index),
+        "eval_target_index": _sha256(_resolve_index_path(args.eval_target_index)),
     }
+    if getattr(args, "source_index", None) is not None:
+        hashes["source_index"] = _sha256(_resolve_index_path(args.source_index))
     (out_dir / "data_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n")
     (out_dir / "split_hash.txt").write_text(
         json.dumps(
@@ -668,7 +872,7 @@ def _write_provenance(out_dir: Path, args: argparse.Namespace) -> None:
 
 
 def _index_hash(path: Path) -> dict[str, Any]:
-    payload = json.loads(path.read_text())
+    payload = json.loads(_resolve_index_path(path).read_text())
     return {
         "split": payload.get("split"),
         "content_hash": payload.get("content_hash"),
@@ -685,6 +889,10 @@ def _sha256(path: Path) -> str:
     return h.hexdigest()
 
 
+def _resolve_index_path(path: Path) -> Path:
+    return path / "index.json" if path.is_dir() else path
+
+
 def _fmt(value: Any) -> str:
     if not isinstance(value, (int, float)) or not math.isfinite(float(value)):
         return "n/a"