feat: transfer released physical drift profile to Rain EP

1 files changed, 111 insertions, 8 deletions

diff --git a/sdil/rain_ep_adapter.py b/sdil/rain_ep_adapter.py
index 3711a6a..b16e96d 100644
--- a/sdil/rain_ep_adapter.py
+++ b/sdil/rain_ep_adapter.py
@@ -8,8 +8,10 @@ mode is introduced here.
 
 from __future__ import annotations
 
+from dataclasses import dataclass
+import math
 from types import MethodType
-from typing import Iterable
+from typing import Iterable, Mapping
 
 import torch
 
@@ -22,6 +24,70 @@ from sdil.two_state_debias import (
 Tensor = torch.Tensor
 
 
+@dataclass(frozen=True)
+class DillavouBiasProfile:
+    """Dimensionless per-edge affine shapes fitted from released drift data."""
+
+    normalized_offsets: tuple[float, ...]
+    normalized_state_variations: tuple[float, ...]
+    source: str
+
+    def __post_init__(self) -> None:
+        if not self.normalized_offsets:
+            raise ValueError("Dillavou profile must contain at least one edge")
+        if len(self.normalized_offsets) != len(
+            self.normalized_state_variations
+        ):
+            raise ValueError("Dillavou profile arrays disagree")
+        values = self.normalized_offsets + self.normalized_state_variations
+        if not all(math.isfinite(value) for value in values):
+            raise ValueError("Dillavou profile contains a nonfinite value")
+
+    @classmethod
+    def from_state_dependence_report(
+        cls, report: Mapping, *, source: str
+    ) -> "DillavouBiasProfile":
+        offsets = []
+        state_variations = []
+        pairs = report.get("pairs", {})
+        for pair_name in sorted(pairs):
+            pair = pairs[pair_name]
+            reference = pair["reference_gate"]
+            affine = pair["local_affine_model"]
+            pair_offsets = affine["bias_at_reference_v_per_s"]
+            pair_slopes = affine["local_slopes_per_s"]
+            gates_by_edge = [[], []]
+            for trace in pair["traces"]:
+                gates_by_edge[0].extend(trace["retained_gate_minus"])
+                gates_by_edge[1].extend(trace["retained_gate_plus"])
+            for edge in range(2):
+                observed_range = max(
+                    abs(float(value) - float(reference[edge]))
+                    for value in gates_by_edge[edge]
+                )
+                offsets.append(float(pair_offsets[edge]))
+                state_variations.append(
+                    float(pair_slopes[edge]) * observed_range)
+        offset_rms = math.sqrt(
+            sum(value * value for value in offsets) / len(offsets))
+        if offset_rms <= 0.0:
+            raise ValueError("released Dillavou offsets have zero RMS")
+        return cls(
+            normalized_offsets=tuple(value / offset_rms for value in offsets),
+            normalized_state_variations=tuple(
+                value / offset_rms for value in state_variations),
+            source=source,
+        )
+
+    def as_dict(self) -> dict:
+        return {
+            "normalized_offsets": list(self.normalized_offsets),
+            "normalized_state_variations": list(
+                self.normalized_state_variations),
+            "source": self.source,
+        }
+
+
 class LocalStructuredBias:
     """Fixed per-element bias affine in a local first-state measurement."""
 
@@ -239,12 +305,16 @@ class DillavouUpdateCorrector:
         calibration_steps: int = 1,
         neutral_cadence: int = 1,
         drift_ratio: float = 0.0,
+        empirical_profile: DillavouBiasProfile | None = None,
         seed: int = 1729,
     ) -> None:
         if mode not in self.MODES:
             raise ValueError(f"unrecognized correction mode {mode}")
         if bias_ratio < 0.0 or drift_ratio < 0.0:
             raise ValueError("Dillavou bias ratios must be nonnegative")
+        if empirical_profile is not None and drift_ratio != 0.0:
+            raise ValueError(
+                "free drift ratio cannot be combined with an empirical profile")
         if not 0.0 < predictor_rate <= 1.0:
             raise ValueError("predictor rate must lie in (0, 1]")
         if neutral_cadence < 0:
@@ -257,6 +327,7 @@ class DillavouUpdateCorrector:
         self.calibration_steps = calibration_steps
         self.neutral_cadence = neutral_cadence
         self.drift_ratio = drift_ratio
+        self.empirical_profile = empirical_profile
         self.seed = seed
         self.steps = 0
         self.debiaser: LocalAffineDebiaser | None = None
@@ -293,14 +364,41 @@ class DillavouUpdateCorrector:
         ):
             gradient_scale = gradient.square().mean().sqrt().clamp_min(1e-12)
             parameter_scale = parameter.square().mean().sqrt().clamp_min(1e-6)
-            offset_pattern = self._unit_pattern(
-                gradient, float(self.seed + 97 * index))
-            slope_pattern = self._unit_pattern(
-                gradient, float(self.seed + 193 * index + 41))
+            if self.empirical_profile is None:
+                offset_pattern = self._unit_pattern(
+                    gradient, float(self.seed + 97 * index))
+                slope_pattern = self._unit_pattern(
+                    gradient, float(self.seed + 193 * index + 41))
+                offset_gain = self.bias_ratio
+                slope_gain = self.drift_ratio
+            else:
+                flat_index = torch.arange(
+                    gradient.numel(), dtype=torch.int64,
+                    device=gradient.device).reshape(gradient.shape)
+                profile_index = torch.remainder(
+                    flat_index + self.seed + 97 * index,
+                    len(self.empirical_profile.normalized_offsets))
+                offset_values = torch.as_tensor(
+                    self.empirical_profile.normalized_offsets,
+                    dtype=gradient.dtype, device=gradient.device)
+                slope_values = torch.as_tensor(
+                    self.empirical_profile.normalized_state_variations,
+                    dtype=gradient.dtype, device=gradient.device)
+                offset_pattern = offset_values[profile_index]
+                slope_pattern = slope_values[profile_index]
+                # Small tensors need not contain the four profiles equally.
+                # Renormalize only their common amplitude; the measured
+                # per-profile slope/offset ratios remain unchanged.
+                pattern_rms = offset_pattern.square().mean().sqrt().clamp_min(
+                    1e-30)
+                offset_pattern = offset_pattern / pattern_rms
+                slope_pattern = slope_pattern / pattern_rms
+                offset_gain = self.bias_ratio
+                slope_gain = self.bias_ratio
             self._offsets.append(
-                self.bias_ratio * gradient_scale * offset_pattern)
+                offset_gain * gradient_scale * offset_pattern)
             self._slopes.append(
-                self.drift_ratio * gradient_scale * slope_pattern)
+                slope_gain * gradient_scale * slope_pattern)
             self._centers.append(parameter.clone())
             self._scales.append(parameter_scale.clone())
             feature_scales.append(torch.ones_like(parameter_scale))
@@ -395,10 +493,15 @@ class DillavouUpdateCorrector:
         self.last_diagnostics = {
             "step": self.steps,
             "bias_model": (
-                "dillavou_constant_update"
+                "dillavou_released_affine_update"
+                if self.empirical_profile is not None
+                else "dillavou_constant_update"
                 if self.drift_ratio == 0.0
                 else "dillavou_plus_local_state_drift"
             ),
+            "bias_profile_source": (
+                None if self.empirical_profile is None
+                else self.empirical_profile.source),
             "clean_update_rms": clean_rms,
             "bias_update_rms": bias_rms,
             "residual_update_rms": residual_rms,