feat: add correlated autozero sampling

1 files changed, 123 insertions, 4 deletions

diff --git a/sdil/physical_grid.py b/sdil/physical_grid.py
index 1d46714..51cd1eb 100644
--- a/sdil/physical_grid.py
+++ b/sdil/physical_grid.py
@@ -287,6 +287,88 @@ class AutozeroSampleHold:
         return self.sample_gain * neutral + offset + noise
 
 
+@dataclass(frozen=True)
+class CorrelatedDoubleSampleHold:
+    """Matched local sampling of neutral and active learning outputs.
+
+    A common pedestal is added to both samples and therefore cancels in their
+    difference.  The mismatch fields model the remaining difference between
+    the two sampling phases rather than exposing correction parameters to the
+    learning rule.
+    """
+
+    common_sample_gain: float = 1.0
+    sample_gain_mismatch: float | Array = 0.0
+    common_pedestal_offset_v_per_s: float | Array = 0.0
+    pedestal_mismatch_v_per_s: float | Array = 0.0
+    sample_noise_standard_deviation_v_per_s: float = 0.0
+    refresh_interval_updates: int = 1
+
+    def __post_init__(self) -> None:
+        if (
+            not np.isfinite(self.common_sample_gain)
+            or self.common_sample_gain < 0.0
+        ):
+            raise ValueError("common sample gain must be finite and nonnegative")
+        if (
+            not np.isfinite(self.sample_noise_standard_deviation_v_per_s)
+            or self.sample_noise_standard_deviation_v_per_s < 0.0
+        ):
+            raise ValueError("sample noise must be finite and nonnegative")
+        if self.refresh_interval_updates < 1:
+            raise ValueError("refresh interval must be positive")
+
+    @staticmethod
+    def _edge_vector(value: float | Array, shape: tuple[int, ...]) -> Array:
+        array = np.asarray(value, dtype=float)
+        try:
+            return np.broadcast_to(array, shape)
+        except ValueError as error:
+            raise ValueError(
+                "correlated-sampling parameter cannot broadcast to edges"
+            ) from error
+
+    def sample_neutral(
+        self, neutral_output: Array, rng: np.random.Generator
+    ) -> Array:
+        neutral = np.asarray(neutral_output, dtype=float)
+        gain_mismatch = self._edge_vector(
+            self.sample_gain_mismatch, neutral.shape)
+        common_pedestal = self._edge_vector(
+            self.common_pedestal_offset_v_per_s, neutral.shape)
+        pedestal_mismatch = self._edge_vector(
+            self.pedestal_mismatch_v_per_s, neutral.shape)
+        noise = rng.normal(
+            0.0,
+            self.sample_noise_standard_deviation_v_per_s,
+            neutral.shape,
+        )
+        return (
+            (self.common_sample_gain - 0.5 * gain_mismatch) * neutral
+            + common_pedestal - 0.5 * pedestal_mismatch + noise
+        )
+
+    def sample_active(
+        self, active_output: Array, rng: np.random.Generator
+    ) -> Array:
+        active = np.asarray(active_output, dtype=float)
+        gain_mismatch = self._edge_vector(
+            self.sample_gain_mismatch, active.shape)
+        common_pedestal = self._edge_vector(
+            self.common_pedestal_offset_v_per_s, active.shape)
+        pedestal_mismatch = self._edge_vector(
+            self.pedestal_mismatch_v_per_s, active.shape)
+        noise = rng.normal(
+            0.0,
+            self.sample_noise_standard_deviation_v_per_s,
+            active.shape,
+        )
+        return (
+            (self.common_sample_gain + 0.5 * gain_mismatch) * active
+            + common_pedestal + 0.5 * pedestal_mismatch + noise
+        )
+
+
 @dataclass
 class EdgePolynomialPredictor:
     feature_center: Array
@@ -416,6 +498,7 @@ def train_grid_classifier(
     epochs: int,
     predictor: EdgePolynomialPredictor | None = None,
     autozero_sample_hold: AutozeroSampleHold | None = None,
+    correlated_sample_hold: CorrelatedDoubleSampleHold | None = None,
     autozero_seed: int = 0,
     standard_nudging: float = 128.0 / 129.0,
     standard_learning_time_seconds: float = 1.0e-3,
@@ -433,12 +516,14 @@ def train_grid_classifier(
         "sdil",
         "oracle_neutral",
         "autozero_sdil",
+        "cds_autozero_sdil",
         "overclamp_clean",
         "overclamp",
         "overclamp_constant",
         "overclamp_sdil",
         "overclamp_oracle_neutral",
         "overclamp_autozero_sdil",
+        "overclamp_cds_autozero_sdil",
     }
     if method not in allowed:
         raise ValueError(f"unrecognized method {method}")
@@ -451,8 +536,13 @@ def train_grid_classifier(
     autozero_methods = {
         "autozero_sdil", "overclamp_autozero_sdil"
     }
+    correlated_methods = {
+        "cds_autozero_sdil", "overclamp_cds_autozero_sdil"
+    }
     if method in autozero_methods and autozero_sample_hold is None:
         autozero_sample_hold = AutozeroSampleHold()
+    if method in correlated_methods and correlated_sample_hold is None:
+        correlated_sample_hold = CorrelatedDoubleSampleHold()
     gates = np.asarray(initial_gates, dtype=float).copy()
     if gates.shape != (circuit.edge_count,):
         raise ValueError("initial gate vector has the wrong shape")
@@ -471,10 +561,13 @@ def train_grid_classifier(
     local_updates = 0
     clipped_updates = 0
     autozero_samples = 0
+    autozero_active_samples = 0
     autozero_updates_since_sample = 0
     held_neutral_output: Array | None = None
     autozero_error_sum_squared = 0.0
     autozero_error_entries = 0
+    autozero_applied_error_sum_squared = 0.0
+    autozero_applied_error_entries = 0
     autozero_rng = np.random.default_rng(autozero_seed)
     perfect_checkpoints = 0
     completed_epochs = 0
@@ -496,16 +589,26 @@ def train_grid_classifier(
             if label * error <= 0.0:
                 continue
             free_drops = edge_voltage_drops(circuit, free_state)
-            if method in autozero_methods:
+            if method in autozero_methods | correlated_methods:
                 current_neutral_output = imperfection.observed_rate(
                     circuit.measured_learning_rate, free_drops, free_drops)
+                refresh_interval = (
+                    autozero_sample_hold.refresh_interval_updates
+                    if method in autozero_methods
+                    else correlated_sample_hold.refresh_interval_updates
+                )
                 if (
                     held_neutral_output is None
                     or autozero_updates_since_sample
-                    >= autozero_sample_hold.refresh_interval_updates
+                    >= refresh_interval
                 ):
-                    held_neutral_output = autozero_sample_hold.sample(
-                        current_neutral_output, autozero_rng)
+                    if method in autozero_methods:
+                        held_neutral_output = autozero_sample_hold.sample(
+                            current_neutral_output, autozero_rng)
+                    else:
+                        held_neutral_output = (
+                            correlated_sample_hold.sample_neutral(
+                                current_neutral_output, autozero_rng))
                     autozero_samples += 1
                     autozero_updates_since_sample = 0
                 autozero_error_sum_squared += float(np.sum(np.square(
@@ -550,6 +653,17 @@ def train_grid_classifier(
                 applied_rate = observed_rate - neutral_output
             elif method in autozero_methods:
                 applied_rate = observed_rate - held_neutral_output
+            elif method in correlated_methods:
+                sampled_active_rate = correlated_sample_hold.sample_active(
+                    observed_rate, autozero_rng)
+                autozero_active_samples += 1
+                applied_rate = sampled_active_rate - held_neutral_output
+                reference_rate = (
+                    correlated_sample_hold.common_sample_gain
+                    * (observed_rate - current_neutral_output))
+                autozero_applied_error_sum_squared += float(np.sum(np.square(
+                    applied_rate - reference_rate)))
+                autozero_applied_error_entries += circuit.edge_count
             elif method in {
                 "constant", "sdil", "overclamp_constant", "overclamp_sdil"
             }:
@@ -601,10 +715,15 @@ def train_grid_classifier(
         "max_abs_clamp_displacement_v": float(max_clamp_displacement),
         "clipped_updates": clipped_updates,
         "autozero_samples": autozero_samples,
+        "autozero_active_samples": autozero_active_samples,
         "autozero_sample_fraction_per_update": float(
             autozero_samples / local_updates if local_updates else 0.0),
         "autozero_baseline_rmse_v_per_s": float(np.sqrt(
             autozero_error_sum_squared / autozero_error_entries
         )) if autozero_error_entries else None,
+        "autozero_applied_rate_rmse_v_per_s": float(np.sqrt(
+            autozero_applied_error_sum_squared
+            / autozero_applied_error_entries
+        )) if autozero_applied_error_entries else None,
         "trace": trace,
     }