feat: train the reconstructed physical grid

1 files changed, 31 insertions, 0 deletions

diff --git a/experiments/physical_grid_smoke.py b/experiments/physical_grid_smoke.py
index a84ed84..d5ababd 100644
--- a/experiments/physical_grid_smoke.py
+++ b/experiments/physical_grid_smoke.py
@@ -15,11 +15,14 @@ from sdil.physical_grid import (  # noqa: E402
     EdgePolynomialPredictor,
     GridCircuit,
     GridSquareLawImperfection,
+    RingClassificationDataset,
     _residual_and_jacobian,
     edge_voltage_drops,
     fit_edge_predictor,
     output_difference,
     solve_grid_state,
+    evaluate_grid_classifier,
+    train_grid_classifier,
 )
 
 
@@ -106,6 +109,32 @@ def main() -> None:
     ])))
     assert quadratic_rmse < 1e-7
     assert constant_rmse > 1e-3
+
+    angles = np.arange(8) * (2.0 * np.pi / 8.0)
+    midpoint = 0.5 * (circuit.low_voltage + circuit.high_voltage)
+    ring = RingClassificationDataset(
+        inputs_v=np.column_stack((
+            midpoint + 0.18 * np.cos(angles),
+            midpoint - 0.18 * np.sin(angles),
+        )),
+        labels_v=np.asarray([-0.018] * 4 + [0.018] * 4),
+    )
+    initial_gates = np.random.default_rng(7).normal(
+        2.33, 0.02, circuit.edge_count)
+    initial_metrics, _ = evaluate_grid_classifier(
+        circuit, initial_gates, ring)
+    trained = train_grid_classifier(
+        circuit,
+        initial_gates,
+        ring,
+        ideal,
+        method="clean",
+        epochs=30,
+        standard_learning_time_seconds=1e-3,
+        record_every=10,
+    )
+    assert trained["hinge_loss_v2"] < initial_metrics["hinge_loss_v2"]
+    assert trained["local_updates"] > 0
     print({
         "nodes": circuit.node_count,
         "edges": circuit.edge_count,
@@ -113,6 +142,8 @@ def main() -> None:
         "target_nodes": circuit.target_nodes,
         "constant_neutral_rmse_v_per_s": constant_rmse,
         "quadratic_neutral_rmse_v_per_s": quadratic_rmse,
+        "clean_hinge_before": initial_metrics["hinge_loss_v2"],
+        "clean_hinge_after": trained["hinge_loss_v2"],
         "autodiff_used": False,
     })