import os, sys, json, math, random import numpy as np import torch import torch.nn as nn import torch.nn.functional as F sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import train_model, evaluate, sweep_baseline, make_report from custom_orientation_lines import get_dataset as get_custom_dataset SEEDS = tuple(range(8)) NTRAIN, NTEST, EPOCHS = 400, 200, 8 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass def oriented_bank(): # Eight unoriented first-derivative filters. Energies are nonnegative. k = torch.zeros(8, 1, 5, 5) yy, xx = torch.meshgrid(torch.arange(-2, 3), torch.arange(-2, 3), indexing="ij") for b in range(8): a = math.pi * b / 8.0 # directional derivative normal to a, with a Gaussian envelope u = torch.cos(torch.tensor(a)); v = torch.sin(torch.tensor(a)) g = torch.exp(-(xx.float()**2 + yy.float()**2) / 3.0) k[b] = (u * xx.float() + v * yy.float()) * g k[b] -= k[b].mean() k[b] /= k[b].abs().sum().clamp_min(1e-6) return k class OrientationCNN(nn.Module): def __init__(self, out_dim=10, mode="double"): super().__init__(); self.mode = mode self.register_buffer("bank", oriented_bank()) self.proj = nn.Conv2d(3, 32, 1) self.conv2 = nn.Conv2d(32, 64, 3, padding=1) self.conv3 = nn.Conv2d(64, 96, 3, padding=1) self.adapt = nn.AdaptiveAvgPool2d((4, 4)) self.fc1 = nn.Linear(96 * 4 * 4, 128); self.fc2 = nn.Linear(128, out_dim) self.last_pool = None def pool(self, x): gray = (0.299*x[:,0:1] + 0.587*x[:,1:2] + 0.114*x[:,2:3]) e = F.conv2d(gray, self.bank, padding=2).abs() if self.mode == "mean": p = e.mean(1, keepdim=True).expand(-1, 3, -1, -1) else: B = e.shape[1] theta = torch.arange(B, device=x.device, dtype=x.dtype) * math.pi / B zr = (e * torch.cos(2*theta)[None,:,None,None]).sum(1, keepdim=True) zi = (e * torch.sin(2*theta)[None,:,None,None]).sum(1, keepdim=True) r = e.sum(1, keepdim=True) q = torch.sqrt(zr.square() + zi.square() + 1e-12) / (r + 1e-6) p = torch.cat((zr/(r+1e-6), zi/(r+1e-6), q), 1) self.last_pool = p return p def forward(self, x): h = F.relu(self.proj(self.pool(x))); h = F.max_pool2d(h, 2) h = F.relu(self.conv2(h)); h = F.max_pool2d(h, 2) h = F.relu(self.conv3(h)); h = F.max_pool2d(h, 2) h = self.adapt(h) h = F.relu(self.fc1(h.flatten(1))) return self.fc2(h) def run_system(mode, cfg, seed, return_model=False): seed_all(seed) raw = get_custom_dataset(seed=seed, n_train=NTRAIN, n_test=NTEST) d = {"track": "oriented_line_quotient", "task": raw["task"], "metric": raw["metric"], "xtr": torch.from_numpy(raw["xtr"]), "ytr": torch.from_numpy(raw["ytr"]), "xte": torch.from_numpy(raw["xte"]), "yte": torch.from_numpy(raw["yte"]), "input_shape": raw["xtr"].shape[1:], "out_dim": raw["out_dim"]} net = OrientationCNN(d["out_dim"], mode=mode) trained, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"], weight_decay=cfg["weight_decay"], batch=128, log=lambda *_: None) if trained is None: raise RuntimeError("benchmark training failed") return (float(metric), trained, d) if return_model else float(metric) def main(): # Equal union: every idea lr/weight-decay is also evaluated for baseline. grid = [{"lr": lr, "weight_decay": wd} for lr in (1e-3, 3e-3, 6e-3) for wd in (0.0, 1e-4)] base = sweep_baseline(lambda c: lambda s: run_system("mean", c, s), grid) idea_cfgs = grid # same-sized, same-space idea sweep; full paired seeds idea_trials = [] for cfg in idea_cfgs: r = evaluate(lambda s, c=cfg: run_system("double", c, s), seeds=SEEDS) idea_trials.append({"cfg": cfg, **r}) best = min(idea_trials, key=lambda r: r["mean"]) # Re-test trained models at scale for a behavior-derived mechanism signature. metric, model, d = run_system("double", best["cfg"], 0, return_model=True) model = model.cpu() device = torch.device("cpu") x = d["xte"][:64].to(device) with torch.no_grad(): p = model.pool(x); p180 = model.pool(torch.rot90(x, 2, dims=(-2,-1))) # π rotation should preserve doubled-angle channels and q; measured, not analytic. reversal_err = float((p - p180).abs().mean()) x90 = torch.rot90(x, 1, dims=(-2,-1)) p90 = model.pool(x90) # A 90-degree image rotation predicts a doubled-angle sign flip for zR,zI. rot90_err = float(torch.cat((p90[:,0:1] + p[:,0:1], p90[:,1:2] + p[:,1:2]), 1).abs().mean()) q_diff = float((p[:,2:3] - p180[:,2:3]).abs().mean()) signature = { "predicted": {"pi_rotation_mean_pool_error": 0.0, "q_pi_rotation_error": 0.0, "90deg_doubled_angle_relation": "zR,zI should sign-flip"}, "observed": {"pi_rotation_mean_pool_error": reversal_err, "q_pi_rotation_error": q_diff, "90deg_relation_abs_error": rot90_err}, "tolerance": {"pi_rotation_mean_pool_error": 0.03, "q_pi_rotation_error": 0.03, "90deg_relation_abs_error": 0.10}, "confirmed": bool(reversal_err < 0.03 and q_diff < 0.03 and rot90_err < 0.10), "measurement": "trained idea model on held-out CIFAR test images" } # make_report expects the best idea evaluation and baseline's full paired result. rep = make_report("oriented_line_quotient", "cnn_small", base, best, extra=signature) rep["custom_track"] = {"name": "oriented_line_quotient", "file": "custom_orientation_lines.py", "domain": "vision_orientation"} rep["baseline"]["architecture_note"] = "fixed oriented Sobel bank + scalar mean + shared CNN" rep["idea"]["architecture_note"] = "fixed oriented Sobel bank + doubled-angle zR,zI,q + shared CNN" 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] rep["protocol_note"] = "custom oriented-line subset, 400/200 samples, 8 epochs, 8 paired seeds; baseline sweep uses 4 seeds then full best config." with open("bench_report.json", "w") as f: json.dump(rep, f, indent=2) print(json.dumps(rep, indent=2)) if __name__ == "__main__": main()