1import sys, json, math, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import make_model, train_model, make_report, sweep_baseline
  7
  8META = {
  9    'name': 'resolution_operator',
 10    'domain': 'pde',
 11    'description': 'Parametric 1D elliptic-style operator mapping coefficient parameters to solution observations on an explicit spatial output grid.'
 12}
 13
 14
 15def operator(z, grid):
 16    z = np.asarray(z); grid = np.asarray(grid)[None, :]
 17    y = np.zeros((len(z), grid.shape[1]), dtype=np.float32)
 18    for k in range(1, 9):
 19        c = (np.sin((k + .2) * z[:, 0]) + .7*np.cos((k + 1.1)*z[:, 1])
 20             + .3*z[:, 2]*z[:, 3]) / (k ** 1.12)
 21        y += c[:, None] * np.sin(2*np.pi*k*grid)
 22    return y + (.3*z[:, 0]*z[:, 1])[:, None]*np.cos(2*np.pi*grid)
 23
 24
 25def get_dataset(seed, n_train=400, n_test=400, m=32):
 26    rng = np.random.default_rng(int(seed))
 27    n = int(n_train) + int(n_test)
 28    z = rng.uniform(-1, 1, (n, 6)).astype('float32')
 29    grid = np.linspace(0, 1, int(m), dtype='float32')
 30    y = operator(z, grid) + rng.normal(0, .01, (n, int(m))).astype('float32')
 31    return {'xtr': z[:n_train], 'ytr': y[:n_train], 'xte': z[n_train:],
 32            'yte': y[n_train:], 'task': 'regression', 'metric': 'mse',
 33            'out_dim': int(m), 'input_shape': (6,)}
 34
 35
 36def fill_distance(points, grid):
 37    p = np.asarray(points)[:, None]
 38    q = np.asarray(grid)[None, :]
 39    return float(np.max(np.min(np.abs(q-p), axis=0)))
 40
 41
 42def math_check():
 43    ms = np.array([4, 8, 16, 32, 64])
 44    hs = np.array([fill_distance(np.linspace(0, 1, m), np.linspace(0, 1, 20001)) for m in ms])
 45    theory = 1/(2*(ms-1))
 46    slope = float(np.polyfit(np.log(ms), np.log(hs), 1)[0])
 47    # Oracle decoder: linear interpolation of exact observations on a fine reference grid.
 48    z = np.array([[.2, -.4, .3, .7, -.1, .5]], dtype=np.float32)
 49    fine = np.linspace(0, 1, 513)
 50    truth = operator(z, fine)[0]
 51    oracle = {}
 52    for m in (16, 32, 64):
 53        coarse = operator(z, np.linspace(0, 1, m))[0]
 54        rec = np.interp(fine, np.linspace(0, 1, m), coarse)
 55        oracle[str(m)] = float(np.mean((rec-truth)**2))
 56    return {'uniform_fill_distance': hs.tolist(), 'theory': theory.tolist(),
 57            'loglog_slope': slope, 'max_theory_error': float(np.max(abs(hs-theory))),
 58            'oracle_mse': oracle, 'slope_close_to_minus_one': bool(abs(slope+1) < .15)}
 59
 60
 61def train_one(seed, m, n_pairs, lr, epochs=12, return_model=False):
 62    d = get_dataset(seed, n_pairs, 200, m)
 63    for key in ('xtr', 'ytr', 'xte', 'yte'):
 64        d[key] = torch.as_tensor(d[key], dtype=torch.float32)
 65    torch.manual_seed(10000 + int(seed)); np.random.seed(10000 + int(seed))
 66    net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
 67    t0 = time.perf_counter()
 68    trained, metric, hist = train_model(net, d, epochs=epochs, lr=float(lr), batch=128, log=lambda *_: None)
 69    out = {'metric': float(metric), 'seconds': time.perf_counter()-t0,
 70           'last_train_mse': float(hist[-1]) if hist else None,
 71           'n_pairs': int(n_pairs), 'm': int(m), 'lr': float(lr)}
 72    if return_model and trained is not None:
 73        trained.eval()
 74        with torch.no_grad():
 75            dev = next(trained.parameters()).device
 76            out['pred'] = trained(d['xte'].to(dev)).detach().cpu().numpy()
 77        out['truth'] = d['yte'].numpy()
 78    return out
 79
 80
 81def main():
 82    check = math_check()
 83    assert check['slope_close_to_minus_one'] and check['max_theory_error'] < 1e-8
 84    m = 64; epochs = 12
 85    lrs = [.001, .003, .01]
 86    # Full union parity: baseline sweeps every LR and every N used by kappa=.5,1,1.5.
 87    budgets = [int(math.ceil(m**k)) for k in (.5, 1.0, 1.5)]
 88    base_grid = [{'lr': lr, 'n_pairs': n, 'm': m, 'epochs': epochs} for lr in lrs for n in budgets]
 89    def baseline_factory(cfg):
 90        return lambda seed: train_one(seed, cfg['m'], cfg['n_pairs'], cfg['lr'], cfg['epochs'])['metric']
 91    sweep = sweep_baseline(baseline_factory, base_grid)
 92    best_cfg = sweep['best_cfg']
 93    # Idea sweep: same three LR values and the a-priori kappa values.
 94    idea_grid = [{'lr': lr, 'kappa': k, 'n_pairs': int(math.ceil(m**k)), 'm': m, 'epochs': epochs}
 95                 for lr in lrs for k in (.5, 1.0, 1.5)]
 96    idea_sweep = []
 97    for cfg in idea_grid:
 98        vals = [train_one(s, m, cfg['n_pairs'], cfg['lr'], epochs)['metric'] for s in range(4)]
 99        idea_sweep.append({'cfg': cfg, 'mean': float(np.mean(vals))})
100    chosen = min(idea_sweep, key=lambda x: x['mean'])['cfg']
101    base_vals=[]; idea_vals=[]; details=[]
102    for seed in range(8):
103        b=train_one(seed, m, best_cfg['n_pairs'], best_cfg['lr'], epochs)
104        i=train_one(seed, m, chosen['n_pairs'], chosen['lr'], epochs)
105        base_vals.append(b['metric']); idea_vals.append(i['metric'])
106        details.append({'seed': seed, 'baseline': b, 'idea': i})
107    base_block={'sweep': sweep, 'full': {'mean': float(np.mean(base_vals)), 'std': float(np.std(base_vals)), 'per_seed': base_vals, 'n': 8}, 'best_config': best_cfg}
108    idea_block={'sweep': idea_sweep, 'full_config': chosen, 'per_seed': idea_vals, 'mean': float(np.mean(idea_vals)), 'details': details}
109    # Signature uses actual trained outputs, not an analytical identity.
110    observed=[]
111    for mm in (16,32,64):
112        r=train_one(0, mm, int(math.ceil(mm)), chosen['lr'], epochs, return_model=True)
113        observed.append({'m': mm, 'mse': r['metric'], 'pred_std': float(np.std(r['pred'])), 'truth_std': float(np.std(r['truth']))})
114    sig={'prediction_mse_by_resolution': observed, 'oracle_mse_by_resolution': check['oracle_mse'],
115         'prediction_decreases_with_resolution': bool(observed[-1]['mse'] < observed[0]['mse']),
116         'confirmed': bool(observed[-1]['mse'] < observed[0]['mse'] and check['oracle_mse']['64'] < check['oracle_mse']['16'])}
117    report=make_report('resolution_operator','mlp_tiny',base_block,idea_block,{'math_check':check,'idea_sweep':idea_sweep,'custom_track':{'name':'resolution_operator','file':'resolution_aware_bench.py','domain':'pde'},'mechanism_signature':sig})
118    Path('bench_report.json').write_text(json.dumps(report,indent=2))
119    print(json.dumps(report,indent=2))
120
121if __name__ == '__main__': main()