import sys, json, math, time from pathlib import Path import numpy as np import torch sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import make_model, train_model, make_report, sweep_baseline META = { 'name': 'resolution_operator', 'domain': 'pde', 'description': 'Parametric 1D elliptic-style operator mapping coefficient parameters to solution observations on an explicit spatial output grid.' } def operator(z, grid): z = np.asarray(z); grid = np.asarray(grid)[None, :] y = np.zeros((len(z), grid.shape[1]), dtype=np.float32) for k in range(1, 9): c = (np.sin((k + .2) * z[:, 0]) + .7*np.cos((k + 1.1)*z[:, 1]) + .3*z[:, 2]*z[:, 3]) / (k ** 1.12) y += c[:, None] * np.sin(2*np.pi*k*grid) return y + (.3*z[:, 0]*z[:, 1])[:, None]*np.cos(2*np.pi*grid) def get_dataset(seed, n_train=400, n_test=400, m=32): rng = np.random.default_rng(int(seed)) n = int(n_train) + int(n_test) z = rng.uniform(-1, 1, (n, 6)).astype('float32') grid = np.linspace(0, 1, int(m), dtype='float32') y = operator(z, grid) + rng.normal(0, .01, (n, int(m))).astype('float32') return {'xtr': z[:n_train], 'ytr': y[:n_train], 'xte': z[n_train:], 'yte': y[n_train:], 'task': 'regression', 'metric': 'mse', 'out_dim': int(m), 'input_shape': (6,)} def fill_distance(points, grid): p = np.asarray(points)[:, None] q = np.asarray(grid)[None, :] return float(np.max(np.min(np.abs(q-p), axis=0))) def math_check(): ms = np.array([4, 8, 16, 32, 64]) hs = np.array([fill_distance(np.linspace(0, 1, m), np.linspace(0, 1, 20001)) for m in ms]) theory = 1/(2*(ms-1)) slope = float(np.polyfit(np.log(ms), np.log(hs), 1)[0]) # Oracle decoder: linear interpolation of exact observations on a fine reference grid. z = np.array([[.2, -.4, .3, .7, -.1, .5]], dtype=np.float32) fine = np.linspace(0, 1, 513) truth = operator(z, fine)[0] oracle = {} for m in (16, 32, 64): coarse = operator(z, np.linspace(0, 1, m))[0] rec = np.interp(fine, np.linspace(0, 1, m), coarse) oracle[str(m)] = float(np.mean((rec-truth)**2)) return {'uniform_fill_distance': hs.tolist(), 'theory': theory.tolist(), 'loglog_slope': slope, 'max_theory_error': float(np.max(abs(hs-theory))), 'oracle_mse': oracle, 'slope_close_to_minus_one': bool(abs(slope+1) < .15)} def train_one(seed, m, n_pairs, lr, epochs=12, return_model=False): d = get_dataset(seed, n_pairs, 200, m) for key in ('xtr', 'ytr', 'xte', 'yte'): d[key] = torch.as_tensor(d[key], dtype=torch.float32) torch.manual_seed(10000 + int(seed)); np.random.seed(10000 + int(seed)) net = make_model('mlp_tiny', d['input_shape'], d['out_dim']) t0 = time.perf_counter() trained, metric, hist = train_model(net, d, epochs=epochs, lr=float(lr), batch=128, log=lambda *_: None) out = {'metric': float(metric), 'seconds': time.perf_counter()-t0, 'last_train_mse': float(hist[-1]) if hist else None, 'n_pairs': int(n_pairs), 'm': int(m), 'lr': float(lr)} if return_model and trained is not None: trained.eval() with torch.no_grad(): dev = next(trained.parameters()).device out['pred'] = trained(d['xte'].to(dev)).detach().cpu().numpy() out['truth'] = d['yte'].numpy() return out def main(): check = math_check() assert check['slope_close_to_minus_one'] and check['max_theory_error'] < 1e-8 m = 64; epochs = 12 lrs = [.001, .003, .01] # Full union parity: baseline sweeps every LR and every N used by kappa=.5,1,1.5. budgets = [int(math.ceil(m**k)) for k in (.5, 1.0, 1.5)] base_grid = [{'lr': lr, 'n_pairs': n, 'm': m, 'epochs': epochs} for lr in lrs for n in budgets] def baseline_factory(cfg): return lambda seed: train_one(seed, cfg['m'], cfg['n_pairs'], cfg['lr'], cfg['epochs'])['metric'] sweep = sweep_baseline(baseline_factory, base_grid) best_cfg = sweep['best_cfg'] # Idea sweep: same three LR values and the a-priori kappa values. idea_grid = [{'lr': lr, 'kappa': k, 'n_pairs': int(math.ceil(m**k)), 'm': m, 'epochs': epochs} for lr in lrs for k in (.5, 1.0, 1.5)] idea_sweep = [] for cfg in idea_grid: vals = [train_one(s, m, cfg['n_pairs'], cfg['lr'], epochs)['metric'] for s in range(4)] idea_sweep.append({'cfg': cfg, 'mean': float(np.mean(vals))}) chosen = min(idea_sweep, key=lambda x: x['mean'])['cfg'] base_vals=[]; idea_vals=[]; details=[] for seed in range(8): b=train_one(seed, m, best_cfg['n_pairs'], best_cfg['lr'], epochs) i=train_one(seed, m, chosen['n_pairs'], chosen['lr'], epochs) base_vals.append(b['metric']); idea_vals.append(i['metric']) details.append({'seed': seed, 'baseline': b, 'idea': i}) 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} idea_block={'sweep': idea_sweep, 'full_config': chosen, 'per_seed': idea_vals, 'mean': float(np.mean(idea_vals)), 'details': details} # Signature uses actual trained outputs, not an analytical identity. observed=[] for mm in (16,32,64): r=train_one(0, mm, int(math.ceil(mm)), chosen['lr'], epochs, return_model=True) observed.append({'m': mm, 'mse': r['metric'], 'pred_std': float(np.std(r['pred'])), 'truth_std': float(np.std(r['truth']))}) sig={'prediction_mse_by_resolution': observed, 'oracle_mse_by_resolution': check['oracle_mse'], 'prediction_decreases_with_resolution': bool(observed[-1]['mse'] < observed[0]['mse']), 'confirmed': bool(observed[-1]['mse'] < observed[0]['mse'] and check['oracle_mse']['64'] < check['oracle_mse']['16'])} 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}) Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__ == '__main__': main()