feat: add magnitude-matched raw apical control

1 files changed, 21 insertions, 1 deletions

diff --git a/experiments/smoke.py b/experiments/smoke.py
index 575f0bd..0bdb0ca 100644
--- a/experiments/smoke.py
+++ b/experiments/smoke.py
@@ -8,7 +8,7 @@ import torch
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from sdil.core import (SDILNet, SDILConfig, sdil_step,
                        node_perturbation_targets, simultaneous_node_perturbation_targets,
-                       neutral_p_update)
+                       neutral_p_update, teaching_signal)
 from sdil.baselines import dfa_config
 from sdil import probes
 from sdil.data import get_dataset, onehot
@@ -164,6 +164,26 @@ def main():
     assert al_raw["cos_r_negg"] == al_raw["cos_apical_negg"], (
         "reported teaching alignment must honor use_residual=False")
 
+    # A magnitude-matched raw signal is still pointed along raw apical activity,
+    # but has exactly the innovation's norm for every sample. This isolates the
+    # directional effect of subtracting the soma-predictable component.
+    cfg_matched = SDILConfig(use_residual=False, learn_A=False, learn_P=True,
+                             raw_scale_control="match_innovation_norm")
+    with torch.no_grad():
+        fwd_n = net_n.forward(x)
+        error_n = net_n.output_error(fwd_n["h"][-1], yoh)
+        for l in range(net_n.L - 1):
+            matched, raw, innovation = teaching_signal(
+                net_n, l, error_n, fwd_n["h"][l + 1], cfg_matched)
+            assert torch.allclose(matched.norm(dim=1), innovation.norm(dim=1),
+                                  atol=1e-6, rtol=1e-5)
+            assert row_cos(matched, raw) > 0.99999
+    al_matched = probes.alignment_report(net_n, x, y, yoh, cfg_matched)
+    for matched_cos, raw_cos in zip(al_matched["cos_r_negg"],
+                                    al_matched["cos_apical_negg"]):
+        assert abs(matched_cos - raw_cos) < 1e-6
+    print("  matched raw: direction preserved; per-sample innovation norm matched")
+
     # ---- CHECK 5: single-step loss-decrease ratio vs exact GD ----
     print("\nCHECK5 single-step loss-decrease ratio vs exact GD:")
     ldr = probes.loss_decrease_ratio(net, x, y, yoh, cfg, step=100)