feat: add local online least-squares residual predictor

2 files changed, 156 insertions, 7 deletions

diff --git a/sdil/rain_ep_adapter.py b/sdil/rain_ep_adapter.py
index b16e96d..7dd47eb 100644
--- a/sdil/rain_ep_adapter.py
+++ b/sdil/rain_ep_adapter.py
@@ -18,6 +18,7 @@ import torch
 from sdil.two_state_debias import (
     BatchedLocalAffineDebiaser,
     LocalAffineDebiaser,
+    LocalLeastSquaresDebiaser,
 )
 
 
@@ -306,6 +307,7 @@ class DillavouUpdateCorrector:
         neutral_cadence: int = 1,
         drift_ratio: float = 0.0,
         empirical_profile: DillavouBiasProfile | None = None,
+        predictor_kind: str = "nlms",
         seed: int = 1729,
     ) -> None:
         if mode not in self.MODES:
@@ -321,6 +323,8 @@ class DillavouUpdateCorrector:
             raise ValueError("neutral cadence must be nonnegative")
         if calibration_steps < 0:
             raise ValueError("calibration steps must be nonnegative")
+        if predictor_kind not in {"nlms", "ols"}:
+            raise ValueError("predictor kind must be nlms or ols")
         self.mode = mode
         self.bias_ratio = bias_ratio
         self.predictor_rate = predictor_rate
@@ -328,9 +332,10 @@ class DillavouUpdateCorrector:
         self.neutral_cadence = neutral_cadence
         self.drift_ratio = drift_ratio
         self.empirical_profile = empirical_profile
+        self.predictor_kind = predictor_kind
         self.seed = seed
         self.steps = 0
-        self.debiaser: LocalAffineDebiaser | None = None
+        self.debiaser: LocalAffineDebiaser | LocalLeastSquaresDebiaser | None = None
         self._offsets: list[Tensor] | None = None
         self._slopes: list[Tensor] | None = None
         self._centers: list[Tensor] | None = None
@@ -402,12 +407,19 @@ class DillavouUpdateCorrector:
             self._centers.append(parameter.clone())
             self._scales.append(parameter_scale.clone())
             feature_scales.append(torch.ones_like(parameter_scale))
-        self.debiaser = LocalAffineDebiaser(
-            clean,
-            feature_centers=[0.0] * len(clean),
-            feature_scales=feature_scales,
-            affine=self.mode == "innovation",
-        )
+        if self.mode == "innovation" and self.predictor_kind == "ols":
+            self.debiaser = LocalLeastSquaresDebiaser(
+                clean,
+                feature_centers=[0.0] * len(clean),
+                feature_scales=feature_scales,
+            )
+        else:
+            self.debiaser = LocalAffineDebiaser(
+                clean,
+                feature_centers=[0.0] * len(clean),
+                feature_scales=feature_scales,
+                affine=self.mode == "innovation",
+            )
 
     @torch.no_grad()
     def _measure(
diff --git a/sdil/two_state_debias.py b/sdil/two_state_debias.py
index b1dfdde..6273b2f 100644
--- a/sdil/two_state_debias.py
+++ b/sdil/two_state_debias.py
@@ -180,6 +180,143 @@ class LocalAffineDebiaser:
         return clone
 
 
+@dataclass
+class LocalLeastSquaresState:
+    intercept: Tensor
+    slope: Tensor
+    feature_center: Tensor
+    feature_scale: Tensor
+    mean_feature: Tensor
+    mean_measurement: Tensor
+    feature_sum_squares: Tensor
+    cross_sum: Tensor
+
+
+class LocalLeastSquaresDebiaser:
+    """Per-element online affine regression from neutral observations.
+
+    Every element keeps its own scalar sufficient statistics. No observation,
+    coefficient, or update is shared across elements, and no autograd graph is
+    created. This is the direct online analogue of fitting a local baseline
+    relation before taking its innovation.
+    """
+
+    def __init__(
+        self,
+        templates: Iterable[Tensor],
+        *,
+        feature_centers: Iterable[Tensor | float],
+        feature_scales: Iterable[Tensor | float],
+        variance_floor: float = 1e-12,
+    ) -> None:
+        templates = _detached(templates)
+        centers = list(feature_centers)
+        scales = list(feature_scales)
+        if not (len(templates) == len(centers) == len(scales)):
+            raise ValueError("template and feature metadata lengths disagree")
+        if variance_floor <= 0.0:
+            raise ValueError("variance floor must be positive")
+        self.variance_floor = variance_floor
+        self.states = []
+        with torch.no_grad():
+            for template, center, scale in zip(templates, centers, scales):
+                center_tensor = torch.as_tensor(
+                    center, dtype=template.dtype, device=template.device)
+                scale_tensor = torch.as_tensor(
+                    scale, dtype=template.dtype, device=template.device)
+                if torch.any(scale_tensor <= 0):
+                    raise ValueError("feature scales must be positive")
+                self.states.append(LocalLeastSquaresState(
+                    intercept=torch.zeros_like(template),
+                    slope=torch.zeros_like(template),
+                    feature_center=center_tensor.clone(),
+                    feature_scale=scale_tensor.clone(),
+                    mean_feature=torch.zeros_like(template),
+                    mean_measurement=torch.zeros_like(template),
+                    feature_sum_squares=torch.zeros_like(template),
+                    cross_sum=torch.zeros_like(template),
+                ))
+        self.neutral_observations = 0
+
+    @staticmethod
+    def _feature(
+        state: LocalLeastSquaresState, local_feature: Tensor
+    ) -> Tensor:
+        return (local_feature - state.feature_center) / state.feature_scale
+
+    @torch.no_grad()
+    def predict(self, local_features: Iterable[Tensor]) -> list[Tensor]:
+        features = _detached(local_features)
+        if len(features) != len(self.states):
+            raise ValueError("feature collection length changed")
+        return [
+            (
+                state.intercept
+                + state.slope * self._feature(state, feature)
+            ).clone()
+            for state, feature in zip(self.states, features)
+        ]
+
+    @torch.no_grad()
+    def update_neutral(
+        self,
+        local_features: Iterable[Tensor],
+        neutral_measurements: Iterable[Tensor],
+        learning_rate: float,
+    ) -> list[Tensor]:
+        if learning_rate != 1.0:
+            raise ValueError(
+                "online least squares requires predictor_rate=1")
+        features = _detached(local_features)
+        measurements = _detached(neutral_measurements)
+        if not (len(features) == len(measurements) == len(self.states)):
+            raise ValueError("neutral tuple lengths disagree")
+        count = self.neutral_observations + 1
+        residuals = []
+        for state, feature, measurement in zip(
+            self.states, features, measurements
+        ):
+            normalized = self._feature(state, feature)
+            residuals.append(
+                measurement
+                - state.intercept
+                - state.slope * normalized)
+            delta_feature = normalized - state.mean_feature
+            delta_measurement = measurement - state.mean_measurement
+            state.mean_feature.add_(delta_feature / count)
+            state.mean_measurement.add_(delta_measurement / count)
+            state.feature_sum_squares.add_(
+                delta_feature * (normalized - state.mean_feature))
+            state.cross_sum.add_(
+                delta_feature * (measurement - state.mean_measurement))
+            identifiable = state.feature_sum_squares > self.variance_floor
+            state.slope.copy_(torch.where(
+                identifiable,
+                state.cross_sum
+                / state.feature_sum_squares.clamp_min(self.variance_floor),
+                torch.zeros_like(state.slope),
+            ))
+            state.intercept.copy_(
+                state.mean_measurement - state.slope * state.mean_feature)
+        self.neutral_observations = count
+        return residuals
+
+    @torch.no_grad()
+    def residual(
+        self,
+        local_features: Iterable[Tensor],
+        teaching_measurements: Iterable[Tensor],
+    ) -> list[Tensor]:
+        measurements = _detached(teaching_measurements)
+        predictions = self.predict(local_features)
+        if len(measurements) != len(predictions):
+            raise ValueError("teaching tuple lengths disagree")
+        return [
+            measurement - prediction
+            for measurement, prediction in zip(measurements, predictions)
+        ]
+
+
 class BatchedLocalAffineDebiaser:
     """Per-cell LMS whose coefficients are shared across observations.