diff options
Diffstat (limited to 'sdil')
| -rw-r--r-- | sdil/rain_ep_adapter.py | 119 |
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, |