1import os, sys, json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import train_model, evaluate, sweep_baseline, make_report
  9from custom_orientation_lines import get_dataset as get_custom_dataset
 10
 11SEEDS = tuple(range(8))
 12NTRAIN, NTEST, EPOCHS = 400, 200, 8
 13
 14
 15def seed_all(seed):
 16    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 17    if torch.cuda.is_available():
 18        try: torch.cuda.manual_seed_all(seed)
 19        except Exception: pass
 20
 21
 22def oriented_bank():
 23    # Eight unoriented first-derivative filters. Energies are nonnegative.
 24    k = torch.zeros(8, 1, 5, 5)
 25    yy, xx = torch.meshgrid(torch.arange(-2, 3), torch.arange(-2, 3), indexing="ij")
 26    for b in range(8):
 27        a = math.pi * b / 8.0
 28        # directional derivative normal to a, with a Gaussian envelope
 29        u = torch.cos(torch.tensor(a)); v = torch.sin(torch.tensor(a))
 30        g = torch.exp(-(xx.float()**2 + yy.float()**2) / 3.0)
 31        k[b] = (u * xx.float() + v * yy.float()) * g
 32        k[b] -= k[b].mean()
 33        k[b] /= k[b].abs().sum().clamp_min(1e-6)
 34    return k
 35
 36
 37class OrientationCNN(nn.Module):
 38    def __init__(self, out_dim=10, mode="double"):
 39        super().__init__(); self.mode = mode
 40        self.register_buffer("bank", oriented_bank())
 41        self.proj = nn.Conv2d(3, 32, 1)
 42        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
 43        self.conv3 = nn.Conv2d(64, 96, 3, padding=1)
 44        self.adapt = nn.AdaptiveAvgPool2d((4, 4))
 45        self.fc1 = nn.Linear(96 * 4 * 4, 128); self.fc2 = nn.Linear(128, out_dim)
 46        self.last_pool = None
 47
 48    def pool(self, x):
 49        gray = (0.299*x[:,0:1] + 0.587*x[:,1:2] + 0.114*x[:,2:3])
 50        e = F.conv2d(gray, self.bank, padding=2).abs()
 51        if self.mode == "mean":
 52            p = e.mean(1, keepdim=True).expand(-1, 3, -1, -1)
 53        else:
 54            B = e.shape[1]
 55            theta = torch.arange(B, device=x.device, dtype=x.dtype) * math.pi / B
 56            zr = (e * torch.cos(2*theta)[None,:,None,None]).sum(1, keepdim=True)
 57            zi = (e * torch.sin(2*theta)[None,:,None,None]).sum(1, keepdim=True)
 58            r = e.sum(1, keepdim=True)
 59            q = torch.sqrt(zr.square() + zi.square() + 1e-12) / (r + 1e-6)
 60            p = torch.cat((zr/(r+1e-6), zi/(r+1e-6), q), 1)
 61        self.last_pool = p
 62        return p
 63
 64    def forward(self, x):
 65        h = F.relu(self.proj(self.pool(x))); h = F.max_pool2d(h, 2)
 66        h = F.relu(self.conv2(h)); h = F.max_pool2d(h, 2)
 67        h = F.relu(self.conv3(h)); h = F.max_pool2d(h, 2)
 68        h = self.adapt(h)
 69        h = F.relu(self.fc1(h.flatten(1)))
 70        return self.fc2(h)
 71
 72
 73def run_system(mode, cfg, seed, return_model=False):
 74    seed_all(seed)
 75    raw = get_custom_dataset(seed=seed, n_train=NTRAIN, n_test=NTEST)
 76    d = {"track": "oriented_line_quotient", "task": raw["task"], "metric": raw["metric"],
 77         "xtr": torch.from_numpy(raw["xtr"]), "ytr": torch.from_numpy(raw["ytr"]),
 78         "xte": torch.from_numpy(raw["xte"]), "yte": torch.from_numpy(raw["yte"]),
 79         "input_shape": raw["xtr"].shape[1:], "out_dim": raw["out_dim"]}
 80    net = OrientationCNN(d["out_dim"], mode=mode)
 81    trained, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"],
 82                                     weight_decay=cfg["weight_decay"], batch=128,
 83                                     log=lambda *_: None)
 84    if trained is None: raise RuntimeError("benchmark training failed")
 85    return (float(metric), trained, d) if return_model else float(metric)
 86
 87
 88def main():
 89    # Equal union: every idea lr/weight-decay is also evaluated for baseline.
 90    grid = [{"lr": lr, "weight_decay": wd}
 91            for lr in (1e-3, 3e-3, 6e-3) for wd in (0.0, 1e-4)]
 92    base = sweep_baseline(lambda c: lambda s: run_system("mean", c, s), grid)
 93    idea_cfgs = grid  # same-sized, same-space idea sweep; full paired seeds
 94    idea_trials = []
 95    for cfg in idea_cfgs:
 96        r = evaluate(lambda s, c=cfg: run_system("double", c, s), seeds=SEEDS)
 97        idea_trials.append({"cfg": cfg, **r})
 98    best = min(idea_trials, key=lambda r: r["mean"])
 99    # Re-test trained models at scale for a behavior-derived mechanism signature.
100    metric, model, d = run_system("double", best["cfg"], 0, return_model=True)
101    model = model.cpu()
102    device = torch.device("cpu")
103    x = d["xte"][:64].to(device)
104    with torch.no_grad():
105        p = model.pool(x); p180 = model.pool(torch.rot90(x, 2, dims=(-2,-1)))
106        # π rotation should preserve doubled-angle channels and q; measured, not analytic.
107        reversal_err = float((p - p180).abs().mean())
108        x90 = torch.rot90(x, 1, dims=(-2,-1))
109        p90 = model.pool(x90)
110        # A 90-degree image rotation predicts a doubled-angle sign flip for zR,zI.
111        rot90_err = float(torch.cat((p90[:,0:1] + p[:,0:1], p90[:,1:2] + p[:,1:2]), 1).abs().mean())
112        q_diff = float((p[:,2:3] - p180[:,2:3]).abs().mean())
113    signature = {
114        "predicted": {"pi_rotation_mean_pool_error": 0.0, "q_pi_rotation_error": 0.0,
115                       "90deg_doubled_angle_relation": "zR,zI should sign-flip"},
116        "observed": {"pi_rotation_mean_pool_error": reversal_err,
117                      "q_pi_rotation_error": q_diff, "90deg_relation_abs_error": rot90_err},
118        "tolerance": {"pi_rotation_mean_pool_error": 0.03, "q_pi_rotation_error": 0.03,
119                       "90deg_relation_abs_error": 0.10},
120        "confirmed": bool(reversal_err < 0.03 and q_diff < 0.03 and rot90_err < 0.10),
121        "measurement": "trained idea model on held-out CIFAR test images"
122    }
123    # make_report expects the best idea evaluation and baseline's full paired result.
124    rep = make_report("oriented_line_quotient", "cnn_small", base, best, extra=signature)
125    rep["custom_track"] = {"name": "oriented_line_quotient", "file": "custom_orientation_lines.py", "domain": "vision_orientation"}
126    rep["baseline"]["architecture_note"] = "fixed oriented Sobel bank + scalar mean + shared CNN"
127    rep["idea"]["architecture_note"] = "fixed oriented Sobel bank + doubled-angle zR,zI,q + shared CNN"
128    rep["idea"]["trials"] = [{k: v for k, v in r.items() if k in ("cfg", "mean", "std", "per_seed", "n")} for r in idea_trials]
129    rep["protocol_note"] = "custom oriented-line subset, 400/200 samples, 8 epochs, 8 paired seeds; baseline sweep uses 4 seeds then full best config."
130    with open("bench_report.json", "w") as f: json.dump(rep, f, indent=2)
131    print(json.dumps(rep, indent=2))
132
133if __name__ == "__main__": main()