1import json, random, sys
  2import numpy as np
  3import torch
  4from torch import nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import train_model, sweep_baseline, evaluate, make_report
  7import tangent_surface_track as track
  8
  9SEEDS = tuple(range(8))
 10EPOCHS = 14
 11BATCH = 64
 12WEIGHT_DECAY = 1e-5
 13
 14
 15def seed_all(s):
 16    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 17    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 18
 19
 20def project_input(x):
 21    p, v = x[..., :3], x[..., 4:]
 22    d = torch.cdist(p, p)
 23    k = min(8, p.shape[1]-1)
 24    idx = d.topk(k+1, largest=False).indices[..., 1:]
 25    src = p.unsqueeze(1).expand(-1, p.shape[1], -1, -1)
 26    nb = torch.gather(src, 2, idx.unsqueeze(-1).expand(-1, -1, -1, 3))
 27    c = nb - p.unsqueeze(2)
 28    cov = torch.einsum('bmki,bmkj->bmij', c, c) / float(k)
 29    _, e = torch.linalg.eigh(cov)
 30    n = e[..., 0]
 31    n = n / (n.norm(dim=-1, keepdim=True) + 1e-8)
 32    eye = torch.eye(3, dtype=x.dtype, device=x.device).view(1,1,3,3)
 33    P = eye - n.unsqueeze(-1) * n.unsqueeze(-2)
 34    vp = torch.einsum('bmij,bmj->bmi', P, v)
 35    r = p - p[:, :1]
 36    rq = torch.einsum('bij,bmj->bmi', P[:, 0], r)
 37    out = torch.cat([rq, x[..., 3:4], vp], dim=-1)
 38    return out.reshape(x.shape[0], -1), n, vp
 39
 40
 41class SharedNet(nn.Module):
 42    def __init__(self, intrinsic=False, width=64):
 43        super().__init__(); self.intrinsic = intrinsic
 44        dim = 20 * (7 if not intrinsic else 7)
 45        self.net = nn.Sequential(nn.Linear(dim, width), nn.ReLU(),
 46                                 nn.Linear(width, width), nn.ReLU(),
 47                                 nn.Linear(width, 1))
 48
 49    def forward(self, x):
 50        if self.intrinsic:
 51            z, n, vp = project_input(x)
 52        else:
 53            z = x.reshape(x.shape[0], -1); n = None; vp = None
 54        return self.net(z), n, vp
 55
 56
 57def run_one(seed, intrinsic, lr):
 58    seed_all(seed)
 59    raw = track.get_dataset(seed, 400, 100)
 60    ds = {k: torch.as_tensor(raw[k], dtype=torch.float32) for k in ('xtr','ytr','xte','yte')}
 61    ds.update(task='regression', metric='mse')
 62    model = SharedNet(intrinsic=intrinsic)
 63    class Wrapper(nn.Module):
 64        def __init__(self, m): super().__init__(); self.m=m
 65        def forward(self, x): return self.m(x)[0]
 66    net, metric, hist = train_model(Wrapper(model), ds, epochs=EPOCHS, lr=lr,
 67                                    batch=BATCH, weight_decay=WEIGHT_DECAY, log=lambda *_: None)
 68    if net is None: return float('nan'), None
 69    return float(metric), model
 70
 71
 72def factory(intrinsic, lr):
 73    return lambda seed: run_one(seed, intrinsic, lr)[0]
 74
 75
 76def main():
 77    # Both methods are evaluated at every learning rate in the shared union.
 78    grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
 79    base = sweep_baseline(lambda c: factory(False, c['lr']), grid, seeds=(0,1,2,3))
 80    # Ensure baseline is evaluated on all candidate settings as required.
 81    baseline_candidates = []
 82    for c in grid:
 83        baseline_candidates.append({'cfg': c, 'full': evaluate(factory(False, c['lr']), SEEDS)})
 84    best_cfg = min(baseline_candidates, key=lambda z: z['full']['mean'])['cfg']
 85    base['union_full'] = baseline_candidates
 86    idea_candidates = []
 87    for c in grid:
 88        idea_candidates.append({'cfg': c, 'full': evaluate(factory(True, c['lr']), SEEDS)})
 89    idea = min(idea_candidates, key=lambda z: z['full']['mean'])['full']
 90    best_idea_cfg = min(idea_candidates, key=lambda z: z['full']['mean'])['cfg']
 91
 92    # Re-test the proposed mechanism on independently trained benchmark models.
 93    raw_leak, proj_leak, rot_change = [], [], []
 94    for s in SEEDS:
 95        metric, model = run_one(s, True, best_idea_cfg['lr'])
 96        if model is None: continue
 97        raw = track.get_dataset(s+9000, 400, 20)
 98        x = torch.as_tensor(raw['xte'], dtype=torch.float32)
 99        device = next(model.parameters()).device
100        x = x.to(device)
101        with torch.no_grad():
102            _, n, vp = model(x)
103            rawv = x[...,4:]
104            raw_leak.append(float(torch.abs((n*rawv).sum(-1)).mean()))
105            proj_leak.append(float(torch.abs((n*vp).sum(-1)).max()))
106            R = torch.tensor([[0.,-1.,0.],[1.,0.,0.],[0.,0.,1.]], device=device)
107            xr=x.clone(); xr[...,:3]=x[...,:3]@R.T; xr[...,4:]=x[...,4:]@R.T
108            pred1=model(x)[0]; pred2=model(xr)[0]
109            rot_change.append(float(torch.abs(pred1-pred2).mean()/(torch.abs(pred1).mean()+1e-8)))
110    signature = {'raw_mean_normal_component': float(np.mean(raw_leak)),
111                 'projected_max_normal_component': float(np.max(proj_leak)),
112                 'rotation_relative_prediction_change': float(np.mean(rot_change)),
113                 'predicted_projected_leakage': '<1e-5',
114                 'confirmed': bool(np.max(proj_leak) < 1e-5)}
115    report = make_report('tangent_surface_vector_regression', 'shared_mlp',
116                         {'best_cfg': best_cfg, 'sweep': base['sweep'],
117                          'union_full': baseline_candidates, 'full': base['full']},
118                         idea, {'mechanism_signature': signature,
119                                'custom_track': {'name': track.META['name'],
120                                                 'file': 'tangent_surface_track.py',
121                                                 'domain': track.META['domain']},
122                                'idea_sweep': idea_candidates,
123                                'best_idea_cfg': best_idea_cfg})
124    print(json.dumps(report, indent=2))
125
126if __name__ == '__main__': main()