From 4ad699ef9709943fcd01c5c3f6022997279f62ba Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Tue, 14 Jul 2026 13:38:32 +0800
Subject: [PATCH] Report fidelity pilot covariate shift

---
 .../fidelity-headroom/analyze_strong_pilot.py | 56 +++++++++++++++++++
 runs/fidelity-headroom/test_strong_pilot.py   | 10 ++++
 2 files changed, 66 insertions(+)

diff --git a/runs/fidelity-headroom/analyze_strong_pilot.py b/runs/fidelity-headroom/analyze_strong_pilot.py
index 57e79c3..de19dce 100644
--- a/runs/fidelity-headroom/analyze_strong_pilot.py
+++ b/runs/fidelity-headroom/analyze_strong_pilot.py
@@ -121,6 +121,45 @@ def predict_model(
     return _sigmoid(matrix @ model["weights"])
 
 
+def covariate_shift(
+    training_examples: list[PrefixExample],
+    training_simulator: list[tuple[float, ...]],
+    pilot_examples: list[PrefixExample],
+    pilot_simulator: list[tuple[float, ...]],
+    *,
+    instrumentation_aware: bool,
+) -> dict[str, Any]:
+    def matrix(
+        examples: list[PrefixExample], simulator: list[tuple[float, ...]]
+    ) -> np.ndarray:
+        rows = []
+        for example, simulator_features in zip(examples, simulator):
+            values = example.outcome + simulator_features
+            if instrumentation_aware:
+                values += example.instrumentation
+            rows.append(values)
+        return np.asarray(rows, dtype=np.float64)
+
+    training = matrix(training_examples, training_simulator)
+    pilot = matrix(pilot_examples, pilot_simulator)
+    mean = training.mean(axis=0)
+    standard_deviation = training.std(axis=0)
+    standard_deviation[standard_deviation < 1e-8] = 1.0
+    absolute_z = np.abs((pilot - mean) / standard_deviation)
+    names = [*OUTCOME_FEATURES, *SIMULATOR_FEATURES]
+    if instrumentation_aware:
+        names.extend(INSTRUMENTATION_FEATURES)
+    return {
+        "values": numeric(absolute_z.ravel().tolist()),
+        "count_gt_3": int(np.sum(absolute_z > 3.0)),
+        "count_gt_5": int(np.sum(absolute_z > 5.0)),
+        "total_feature_values": int(absolute_z.size),
+        "per_feature_max_abs_z": {
+            name: float(value) for name, value in zip(names, absolute_z.max(axis=0))
+        },
+    }
+
+
 def comparison(
     training_examples: list[PrefixExample],
     training_simulator: list[tuple[float, ...]],
@@ -238,6 +277,22 @@ def analyze(
         red_flags.append("pilot_single_label")
     if len(set(simulator_pass_rates)) <= 1:
         red_flags.append("pilot_simulator_results_identical")
+    covariate_diagnostics = {
+        "sim_plus_outcome": covariate_shift(
+            training_examples,
+            training_simulator,
+            pilot_examples,
+            pilot_simulator,
+            instrumentation_aware=False,
+        ),
+        "sim_plus_outcome_plus_instrumentation": covariate_shift(
+            training_examples,
+            training_simulator,
+            pilot_examples,
+            pilot_simulator,
+            instrumentation_aware=True,
+        ),
+    }
 
     if headline is None:
         decision = {
@@ -321,6 +376,7 @@ def analyze(
             }
             for detail, simulator in zip(pilot_details, pilot_simulator)
         ],
+        "covariate_shift_diagnostic": covariate_diagnostics,
         "decision": decision,
         "provenance": {
             "phase6_metrics": str(phase6_path.resolve()),
diff --git a/runs/fidelity-headroom/test_strong_pilot.py b/runs/fidelity-headroom/test_strong_pilot.py
index 921c2b3..c888438 100644
--- a/runs/fidelity-headroom/test_strong_pilot.py
+++ b/runs/fidelity-headroom/test_strong_pilot.py
@@ -11,6 +11,7 @@ from analyze_prefixes import PrefixExample
 from prepare_pilot_simulator import load_module as load_prepare_module
 from run_pilot_simulator import load_module as load_run_module
 from analyze_strong_pilot import (
+    covariate_shift,
     fit_model,
     load_pilot_simulator,
     predict_model,
@@ -46,6 +47,15 @@ def main() -> None:
         probability = predict_model(model, examples, simulator)
         assert probability.shape == (8,)
         assert np.all((probability >= 0.0) & (probability <= 1.0))
+        shift = covariate_shift(
+            examples,
+            simulator,
+            examples,
+            simulator,
+            instrumentation_aware=instrumentation_aware,
+        )
+        assert shift["values"]["min"] >= 0.0
+        assert shift["count_gt_3"] == 0
 
     payload = {
         "status": "PASS",