1 files changed, 124 insertions, 0 deletions
diff --git a/experiments/rain_ep_dillavou_smoke.py b/experiments/rain_ep_dillavou_smoke.py
new file mode 100644
index 0000000..d4a634a
--- /dev/null
+++ b/experiments/rain_ep_dillavou_smoke.py
@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Mechanics checks for the post-estimator Dillavou update bias."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from sdil.rain_ep_adapter import (  # noqa: E402
+    DillavouUpdateCorrector,
+    attach_dillavou_to_rain_estimator,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--author-root", type=Path, required=True)
+    return parser.parse_args()
+
+
+def build_estimator(author_root: Path):
+    sys.path.insert(0, str(author_root))
+    from model.function.cost import SquaredError
+    from model.function.network import Network
+    from model.hopfield.minimizer import FixedPointMinimizer
+    from model.hopfield.network import DeepHopfieldEnergy
+    from training.sgd import AugmentedFunction, EquilibriumProp
+
+    energy = DeepHopfieldEnergy([(4,), (7,), (3,)], [0.5, 0.5])
+    energy.set_device("cpu")
+    network = Network(energy)
+    cost = SquaredError(energy.layers()[-1])
+    augmented = AugmentedFunction(energy, cost)
+    minimizer = FixedPointMinimizer(augmented, network.free_layers())
+    minimizer.mode = "asynchronous"
+    minimizer.num_iterations = 12
+    estimator = EquilibriumProp(
+        energy.params(), energy.layers(), augmented, cost, minimizer)
+    estimator.variant = "positive"
+    estimator.nudging = 0.25
+    return energy, network, cost, augmented, minimizer, estimator
+
+
+def main() -> None:
+    args = parse_args()
+    torch.manual_seed(20260807)
+
+    # The exact paper model has a fixed B_i after the estimator.  Changing the
+    # clean signal or parameter state must not change that field.
+    clean_a = [torch.randn(11, 7), torch.randn(7)]
+    clean_b = [torch.randn_like(value) for value in clean_a]
+    parameters_a = [torch.randn_like(value) for value in clean_a]
+    parameters_b = [value + 0.3 for value in parameters_a]
+    raw = DillavouUpdateCorrector(
+        mode="raw", bias_ratio=0.2, seed=41)
+    measured_a = raw.apply(clean_a, parameters_a)
+    measured_b = raw.apply(clean_b, parameters_b)
+    bias_a = [value - clean for value, clean in zip(measured_a, clean_a)]
+    bias_b = [value - clean for value, clean in zip(measured_b, clean_b)]
+    fixed_relative_error = max(
+        float((first - second).norm() / first.norm().clamp_min(1e-30))
+        for first, second in zip(bias_a, bias_b)
+    )
+    assert fixed_relative_error < 2e-6, fixed_relative_error
+
+    constant = DillavouUpdateCorrector(
+        mode="constant", bias_ratio=0.2, predictor_rate=1.0, seed=41)
+    innovation = DillavouUpdateCorrector(
+        mode="innovation", bias_ratio=0.2, predictor_rate=1.0, seed=41)
+    corrected_constant = constant.apply(clean_a, parameters_a)
+    corrected_innovation = innovation.apply(clean_a, parameters_a)
+    constant_error = max(
+        float((actual - target).norm() / target.norm().clamp_min(1e-30))
+        for actual, target in zip(corrected_constant, clean_a)
+    )
+    innovation_error = max(
+        float((actual - target).norm() / target.norm().clamp_min(1e-30))
+        for actual, target in zip(corrected_innovation, clean_a)
+    )
+    assert constant_error < 2e-7, constant_error
+    assert innovation_error < 2e-7, innovation_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(
+        args.author_root)
+    x = torch.randn(8, 4)
+    labels = torch.arange(8) % 3
+    network.set_input(x, reset=True)
+    cost.set_target(labels)
+    augmented.nudging = 0.0
+    minimizer.compute_equilibrium()
+    free = [layer.state.clone() for layer in minimizer._layers]
+    clean = [value.clone() for value in estimator.compute_gradient()]
+    for layer, state in zip(minimizer._layers, free):
+        layer.state = state.clone()
+    integrated = DillavouUpdateCorrector(
+        mode="raw", bias_ratio=0.2, seed=53)
+    attach_dillavou_to_rain_estimator(estimator, integrated)
+    measured = estimator.compute_gradient()
+    assert all(not value.requires_grad for value in measured)
+    assert integrated.last_diagnostics["bias_model"] == (
+        "dillavou_constant_update")
+    observed_ratio = integrated.last_diagnostics["bias_to_clean_update_rms"]
+    assert abs(observed_ratio - 0.2) < 2e-6, observed_ratio
+
+    print({
+        "fixed_bias_relative_error_after_state_change": fixed_relative_error,
+        "constant_calibration_relative_error": constant_error,
+        "innovation_relative_error": innovation_error,
+        "integrated_bias_to_clean_update_rms": observed_ratio,
+        "neutral_observations": constant.debiaser.neutral_observations,
+        "autodiff_used_for_learning": False,
+    })
+
+
+if __name__ == "__main__":
+    main()