feat: transfer released physical drift profile to Rain EP

1 files changed, 61 insertions, 0 deletions

diff --git a/experiments/rain_ep_dillavou_smoke.py b/experiments/rain_ep_dillavou_smoke.py
index 9519210..39bf5fd 100644
--- a/experiments/rain_ep_dillavou_smoke.py
+++ b/experiments/rain_ep_dillavou_smoke.py
@@ -13,6 +13,7 @@ ROOT = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(ROOT))
 
 from sdil.rain_ep_adapter import (  # noqa: E402
+    DillavouBiasProfile,
     DillavouUpdateCorrector,
     attach_dillavou_to_rain_estimator,
 )
@@ -21,6 +22,9 @@ from sdil.rain_ep_adapter import (  # noqa: E402
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--author-root", type=Path, required=True)
+    parser.add_argument(
+        "--profile-json", type=Path,
+        default=ROOT / "results/physical_bias/p0_state_dependence.json")
     return parser.parse_args()
 
 
@@ -92,6 +96,58 @@ def main() -> None:
     assert constant.debiaser.neutral_observations == 1
     assert innovation.debiaser.neutral_observations == 1
 
+    # Build the state-dependent shape only from the committed analysis of the
+    # released physical traces. With matched neutral observations, an affine
+    # predictor must generalize across local parameter states better than an
+    # intercept-only predictor.
+    import json
+    profile = DillavouBiasProfile.from_state_dependence_report(
+        json.loads(args.profile_json.read_text()),
+        source=str(args.profile_json.resolve()),
+    )
+    profile_constant = DillavouUpdateCorrector(
+        mode="constant", bias_ratio=0.2, predictor_rate=0.2,
+        calibration_steps=1, neutral_cadence=1,
+        empirical_profile=profile, seed=67)
+    profile_innovation = DillavouUpdateCorrector(
+        mode="innovation", bias_ratio=0.2, predictor_rate=0.2,
+        calibration_steps=1, neutral_cadence=1,
+        empirical_profile=profile, seed=67)
+    parameter_scale = [
+        value.square().mean().sqrt().clamp_min(1e-6)
+        for value in parameters_a
+    ]
+    for _ in range(12):
+        for displacement in torch.linspace(-1.0, 1.0, 21):
+            state = [
+                value + displacement * scale
+                for value, scale in zip(parameters_a, parameter_scale)
+            ]
+            profile_constant.apply(clean_a, state)
+            profile_innovation.apply(clean_a, state)
+    assert (
+        profile_constant.debiaser.neutral_observations
+        == profile_innovation.debiaser.neutral_observations
+    )
+    profile_constant.neutral_cadence = 0
+    profile_innovation.neutral_cadence = 0
+    held_state = [
+        value - 0.55 * scale
+        for value, scale in zip(parameters_a, parameter_scale)
+    ]
+    held_constant = profile_constant.apply(clean_b, held_state)
+    held_innovation = profile_innovation.apply(clean_b, held_state)
+    held_constant_error = sum(
+        float((actual - target).square().sum())
+        for actual, target in zip(held_constant, clean_b)
+    )
+    held_innovation_error = sum(
+        float((actual - target).square().sum())
+        for actual, target in zip(held_innovation, clean_b)
+    )
+    assert held_innovation_error < 0.05 * held_constant_error, (
+        held_innovation_error, held_constant_error)
+
     # Integration check: the corruption is attached after Rain's hand-written
     # local EP estimator and introduces no autograd graph.
     energy, network, cost, augmented, minimizer, estimator = build_estimator(
@@ -122,6 +178,11 @@ def main() -> None:
         "innovation_relative_error": innovation_error,
         "integrated_bias_to_clean_update_rms": observed_ratio,
         "neutral_observations": constant.debiaser.neutral_observations,
+        "released_profile_normalized_offsets": profile.normalized_offsets,
+        "released_profile_normalized_state_variations": (
+            profile.normalized_state_variations),
+        "released_profile_heldout_mse_ratio_affine_over_constant": (
+            held_innovation_error / held_constant_error),
         "autodiff_used_for_learning": False,
     })