1import os, sys, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
  8
  9SEED0 = 2145
 10EPOCHS = 18
 11NTR, NTE = 400, 200
 12N_CAND, ACTIVE = 8, 3
 13
 14class CandidateMLP(nn.Module):
 15    """Shared architecture: incumbent MLP plus a bank of gated feature blocks."""
 16    def __init__(self, input_dim, out_dim=1, active=(0,1,2)):
 17        super().__init__()
 18        self.base = nn.Sequential(nn.Linear(input_dim, 32), nn.Tanh(), nn.Linear(32, 16), nn.Tanh())
 19        self.base_head = nn.Linear(16, out_dim)
 20        self.cands = nn.ModuleList([nn.Sequential(nn.Linear(input_dim, 8), nn.Tanh(), nn.Linear(8, out_dim)) for _ in range(N_CAND)])
 21        self.active = tuple(int(x) for x in active)
 22    def forward(self, x):
 23        y = self.base_head(self.base(x))
 24        for j in self.active:
 25            y = y + self.cands[j](x)
 26        return y
 27    def incumbent(self, x):
 28        return self.base_head(self.base(x))
 29    def candidate_response(self, x, j):
 30        return self.cands[j](x)
 31
 32def seed_all(seed):
 33    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 34    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 35
 36def projected_scores(net, x):
 37    """Q_j=(I-P_B) C_j for scalar responses, Y=I; return gains and singular values."""
 38    with torch.no_grad():
 39        b = net.incumbent(x).detach().cpu().numpy().reshape(-1, 1)
 40        cs = [net.candidate_response(x, j).detach().cpu().numpy().reshape(-1, 1) for j in range(N_CAND)]
 41    # Stable Y-weighted projector via least squares, no explicit inverse assumption.
 42    scores = []
 43    qs = []
 44    for c in cs:
 45        coef, *_ = np.linalg.lstsq(b, c, rcond=1e-8)
 46        q = c - b @ coef
 47        sv = np.linalg.svd(q, compute_uv=False)
 48        scores.append(float(np.sum(q*q))); qs.append(q)
 49    return scores, qs
 50
 51def train_one(track, seed, lr, mode, return_net=False):
 52    seed_all(seed)
 53    d = get_dataset(track, seed, n_train=NTR, n_test=NTE)
 54    dim = int(np.prod(d['input_shape']))
 55    # Keep exact same initial bank for paired baseline/idea; only activation differs.
 56    probe = CandidateMLP(dim, d['out_dim'], active=())
 57    scores, qs = projected_scores(probe, d['xtr'][:min(128, NTR)])
 58    if mode == 'idea':
 59        ranked = sorted(range(N_CAND), key=lambda j: scores[j], reverse=True)
 60        # Riesz filter: reject nearly null and badly scaled responses; fallback preserves budget.
 61        accepted = [j for j in ranked if np.linalg.norm(qs[j]) >= 1e-4 and np.linalg.norm(qs[j]) <= 30][:ACTIVE]
 62        if len(accepted) < ACTIVE: accepted = ranked[:ACTIVE]
 63    else:
 64        accepted = list(range(ACTIVE))
 65    net = CandidateMLP(dim, d['out_dim'], active=accepted)
 66    # train_model is the canonical path for the trained system
 67    net, metric, hist = train_model(net, d, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 68    if net is None: return float('nan') if not return_net else (float('nan'), None, accepted, scores)
 69    if return_net: return float(metric), net, accepted, scores
 70    return float(metric)
 71
 72def make_baseline(track, cfg):
 73    return lambda seed: train_one(track, int(seed), float(cfg['lr']), 'baseline')
 74
 75def make_idea(track, cfg):
 76    return lambda seed: train_one(track, int(seed), float(cfg['lr']), 'idea')
 77
 78def signature(track, cfg, seeds=(0,1,2,3,4,5,6,7)):
 79    pred, obs = [], []
 80    for s in seeds:
 81        metric, net, selected, scores = train_one(track, s, cfg['lr'], 'idea', True)
 82        d = get_dataset(track, s, n_train=NTR, n_test=NTE)
 83        x = d['xte'].to(next(net.parameters()).device)
 84        with torch.no_grad():
 85            full = net(x).cpu().numpy().reshape(-1)
 86            inc = net.incumbent(x).cpu().numpy().reshape(-1)
 87        # predicted residual energy on calibration; observed response energy on held-out trained model
 88        pred.append(float(sum(scores[j] for j in selected)))
 89        obs.append(float(np.mean((full-inc)**2)))
 90    corr = float(np.corrcoef(pred, obs)[0,1]) if np.std(pred)>0 and np.std(obs)>0 else 0.0
 91    return {'predicted_projected_gain_mean': float(np.mean(pred)), 'observed_heldout_selected_response_mse_mean': float(np.mean(obs)), 'predicted_per_seed': pred, 'observed_per_seed': obs, 'correlation': corr, 'confirmed': bool(corr > 0.3)}
 92
 93def main():
 94    track, model = 'tabular', 'candidate_mlp_shared'
 95    # Union of all step sizes is shared by both sides; three settings satisfy idea nearby sweep.
 96    grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
 97    base = sweep_baseline(lambda cfg: make_baseline(track, cfg), grid)
 98    idea_trials = []
 99    for cfg in grid:
100        r = evaluate(make_idea(track, cfg))
101        idea_trials.append({'cfg': cfg, 'result': r})
102    best = min(idea_trials, key=lambda z: z['result']['mean'])
103    report = make_report(track, model, base, best['result'], extra=signature(track, best['cfg']))
104    report['idea_sweep'] = idea_trials
105    report['selection'] = {'candidates': N_CAND, 'active_blocks': ACTIVE, 'epochs': EPOCHS, 'n_train': NTR, 'n_test': NTE, 'track_rationale': 'tabular regression is the built-in optimizer/regularizer/architecture-compatible small MLP track; no PDE/control structure is required by this adapter-bank idea'}
106    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
107    print(json.dumps(report, indent=2))
108
109if __name__ == '__main__': main()