feat: add local online least-squares residual predictor

1 files changed, 137 insertions, 0 deletions

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.