feat: match physical classification early stopping

1 files changed, 16 insertions, 1 deletions

diff --git a/sdil/physical_grid.py b/sdil/physical_grid.py
index bb6b4a4..2ef95e4 100644
--- a/sdil/physical_grid.py
+++ b/sdil/physical_grid.py
@@ -379,6 +379,7 @@ def train_grid_classifier(
     overclamp_target_magnitude_v: float | None = None,
     overclamp_time_seconds_per_v: float = 0.05,
     record_every: int = 50,
+    early_stop_perfect_checkpoints: int | None = None,
 ) -> dict:
     """Train the physical grid with explicit local voltage-square updates."""
     allowed = {
@@ -418,6 +419,8 @@ def train_grid_classifier(
     max_clamp_displacement = 0.0
     local_updates = 0
     clipped_updates = 0
+    perfect_checkpoints = 0
+    completed_epochs = 0
 
     for epoch in range(epochs):
         for sample, (inputs, label) in enumerate(zip(
@@ -494,11 +497,23 @@ def train_grid_classifier(
             metrics, free_cache = evaluate_grid_classifier(
                 circuit, gates, dataset, initial_states=free_cache)
             trace.append({"epoch": epoch, **metrics})
+            if metrics["classification_error"] == 0.0:
+                perfect_checkpoints += 1
+            else:
+                perfect_checkpoints = 0
+            if (
+                early_stop_perfect_checkpoints is not None
+                and perfect_checkpoints >= early_stop_perfect_checkpoints
+            ):
+                completed_epochs = epoch + 1
+                break
+        completed_epochs = epoch + 1
 
     final = trace[-1]
     return {
         "method": method,
-        "epochs": epochs,
+        "requested_epochs": epochs,
+        "completed_epochs": completed_epochs,
         "initial_gates_v": np.asarray(initial_gates).tolist(),
         "final_gates_v": gates.tolist(),
         "classification_error": final["classification_error"],