feat: add counted Rain neutral calibration

1 files changed, 20 insertions, 0 deletions

diff --git a/sdil/rain_ep_adapter.py b/sdil/rain_ep_adapter.py
index 81442aa..9c7ab8f 100644
--- a/sdil/rain_ep_adapter.py
+++ b/sdil/rain_ep_adapter.py
@@ -109,6 +109,20 @@ class RainGradientCorrector:
         )
 
     @torch.no_grad()
+    def observe_neutral(self, local_states: Iterable[Tensor]) -> None:
+        """Fit the local bias field from one instruction-off observation."""
+        if self.mode not in {"constant", "innovation"}:
+            raise ValueError(
+                "neutral predictor observations require constant or innovation mode")
+        local_states = list(local_states)
+        if any(value.requires_grad for value in local_states):
+            raise ValueError("Rain adapter received a requires-grad tensor")
+        bases, bias = self.bias.measure(local_states)
+        if self.debiaser is None:
+            self._initialize_debiaser(local_states)
+        self.debiaser.update_neutral(bases, bias, self.predictor_rate)
+
+    @torch.no_grad()
     def apply(self, clean: Iterable[Tensor], local_states: Iterable[Tensor]) -> list[Tensor]:
         clean = list(clean)
         local_states = list(local_states)
@@ -187,3 +201,9 @@ def attach_to_rain_estimator(estimator, corrector: RainGradientCorrector):
     estimator.sdil_corrector = corrector
     return estimator
 
+
+@torch.no_grad()
+def observe_rain_neutral(estimator, corrector: RainGradientCorrector) -> None:
+    """Expose one label-free Rain equilibrium to the local predictor."""
+    local_states = [updater.grad() for updater in estimator._param_updaters]
+    corrector.observe_neutral(local_states)