1import json, random, sys
  2import numpy as np
  3import torch
  4
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report
  7
  8TRACK = 'poisson_dirichlet'
  9MODEL = 'mlp_tiny'
 10SEEDS = tuple(range(8))
 11EPOCHS = 20
 12BATCH = 128
 13# Shared union: every idea LR is also tested by baseline.
 14GRID = [
 15    {'lr': 1e-3, 'weight_decay': 0.0},
 16    {'lr': 3e-3, 'weight_decay': 0.0},
 17    {'lr': 1e-2, 'weight_decay': 0.0},
 18]
 19
 20
 21def seed_all(seed):
 22    random.seed(seed)
 23    np.random.seed(seed)
 24    torch.manual_seed(seed)
 25    if torch.cuda.is_available():
 26        torch.cuda.manual_seed_all(seed)
 27
 28
 29
 30def basis_matrix():
 31    n_grid, n_modes = 32, 12
 32    x = np.arange(1, n_grid + 1, dtype=np.float32) / (n_grid + 1)
 33    return (np.sqrt(2.0) * np.sin(np.pi * np.outer(
 34        x, np.arange(1, n_modes + 1)))).astype(np.float32)
 35
 36def get_ds(seed):
 37    d = get_dataset(TRACK, seed=int(seed), n_train=400, n_test=400)
 38    # bench flattens custom-track y arrays; restore one field per sample.
 39    nout = int(d['out_dim'])
 40    for k in ('xtr', 'xte'):
 41        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 42    for k, n in (('ytr', len(d['xtr'])), ('yte', len(d['xte']))):
 43        a = np.asarray(d[k]).reshape(n, nout)
 44        d[k] = torch.as_tensor(a, dtype=torch.float32)
 45    return d
 46
 47
 48class NullspaceModel(torch.nn.Module):
 49    def __init__(self, core, basis):
 50        super().__init__()
 51        self.core = core
 52        self.register_buffer('V', torch.as_tensor(basis, dtype=torch.float32))
 53
 54    def forward(self, x):
 55        # core predicts z; decoder always lies in ker(C).
 56        interior = self.core(x) @ self.V.T
 57        z = torch.zeros((interior.shape[0], 1), device=interior.device, dtype=interior.dtype)
 58        return torch.cat((z, interior, z), dim=1)
 59
 60
 61def fit(seed, cfg, constrained, return_model=False):
 62    seed_all(seed)
 63    d = get_ds(seed)
 64    # Baseline predicts the full field, idea predicts the 12 nullspace coordinates.
 65    out_dim = int(basis_matrix().shape[1]) if constrained else int(d['out_dim'])
 66    core = make_model(MODEL, tuple(d['xtr'].shape[1:]), out_dim)
 67    net = NullspaceModel(core, basis_matrix()) if constrained else core
 68    net, metric, history = train_model(
 69        net, d, epochs=EPOCHS, lr=float(cfg['lr']), batch=BATCH,
 70        weight_decay=float(cfg['weight_decay']), log=lambda *_: None)
 71    if net is None or metric is None:
 72        return (float('inf'), None, d)
 73    return (float(metric), net if return_model else None, d)
 74
 75
 76def make_fn(constrained):
 77    def fn(cfg):
 78        return lambda seed: fit(seed, cfg, constrained)[0]
 79    return fn
 80
 81
 82def run():
 83    # Baseline sweep uses four seeds as specified, then re-evaluates best on eight.
 84    base = sweep_baseline(make_fn(False), GRID, seeds=(0, 1, 2, 3))
 85    # Explicitly run the idea at all three shared settings; report best using the
 86    # same four-seed selection budget, then evaluate its selected config on 8.
 87    idea_trials = []
 88    best_cfg, best_mean = None, float('inf')
 89    for cfg in GRID:
 90        r = evaluate(make_fn(True)(cfg), seeds=(0, 1, 2, 3))
 91        idea_trials.append({'cfg': cfg, 'mean': r['mean']})
 92        if r['mean'] < best_mean:
 93            best_cfg, best_mean = cfg, r['mean']
 94    idea_full = evaluate(make_fn(True)(best_cfg), seeds=SEEDS)
 95    base['idea_union_trials'] = idea_trials
 96
 97    # Signature is measured from trained systems: boundary residual and decoded
 98    # nullspace violation on actual held-out inputs. C selects endpoints.
 99    probe_seed = 0
100    bmetric, bnet, bd = fit(probe_seed, base['best_cfg'], False, True)
101    imetric, inet, idata = fit(probe_seed, best_cfg, True, True)
102    C = torch.zeros((2, 34), dtype=torch.float32)
103    C[0, 0] = 1.0; C[1, -1] = 1.0
104    with torch.no_grad():
105        db = next(bnet.parameters()).device
106        di = next(inet.parameters()).device
107        ub = bnet(idata['xte'].to(db)).detach().cpu()
108        ui = inet(idata['xte'].to(di)).detach().cpu()
109    rb = torch.linalg.norm(ub @ C.T, dim=1)
110    ri = torch.linalg.norm(ui @ C.T, dim=1)
111    signature = {
112        'prediction': 'nullspace decoder should make C u exactly zero while ambient output violates endpoints',
113        'predicted_boundary_residual_idea': 0.0,
114        'observed_boundary_residual_idea_max': float(ri.max()),
115        'observed_boundary_residual_baseline_max': float(rb.max()),
116        'observed_boundary_residual_baseline_mean': float(rb.mean()),
117        'observed_boundary_residual_idea_mean': float(ri.mean()),
118        'ratio_baseline_to_idea_max': float((rb.max() + 1e-12) / (ri.max() + 1e-12)),
119        'confirmed': bool(float(ri.max()) < 1e-5 and float(rb.max()) > 1e-3),
120        'trained_models': True,
121    }
122    report = make_report(TRACK, MODEL, base, idea_full, {
123        'selection': 'best of shared 3-point LR/weight-decay grid on seeds 0-3',
124        'idea_best_cfg': best_cfg,
125        **signature,
126    })
127    report['custom_track'] = {
128        'name': 'poisson_dirichlet',
129        'file': '/home/maxwelhelp/all/math2nn/bench/custom_tracks/poisson_dirichlet.py',
130        'domain': 'pde',
131        'reason': 'built-in custom PDE track has fixed-grid field outputs and homogeneous Dirichlet constraints',
132    }
133    report['protocol'] = {'epochs': EPOCHS, 'batch': BATCH, 'seeds': list(SEEDS), 'grid': GRID}
134    with open('bench_report.json', 'w') as f:
135        json.dump(report, f, indent=2)
136    print(json.dumps(report, indent=2))
137
138
139if __name__ == '__main__':
140    run()