import sys, json, math from pathlib import Path import numpy as np import torch from torch import nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import train_model, evaluate, sweep_baseline, make_report META = { 'name': 'singular_radial_poisson', 'domain': 'pde', 'description': 'Radial singular Poisson-style Dirichlet field regression in a polar corner chart; u(r)=r^lambda is learned from computational coordinates with or without a positive monotone radial warp.' } LAMBDA = 0.25 Q = 2.0 def get_dataset(seed, n_train=400, n_test=400): rng = np.random.RandomState(seed) def sample(n): s = np.maximum(rng.rand(n), 1e-5).astype(np.float32) a = (2.0 * np.pi * rng.rand(n)).astype(np.float32) # First coordinate is computational radius; the other two encode angle. x = np.stack([s, np.cos(a), np.sin(a)], axis=1).astype(np.float32) y = (s ** LAMBDA).astype(np.float32)[:, None] return x, y xtr, ytr = sample(n_train) xte, yte = sample(n_test) return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte, 'task': 'regression', 'metric': 'mse', 'out_dim': 1, 'input_shape': (3,)} class CoreMLP(nn.Module): """Shared field architecture for both systems.""" def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(3, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1)) def forward(self, x): return self.net(x) class PositiveRadialWarp(nn.Module): """rho(s)=eps+s^(q-1) exp(h(s)), normalized by cumulative trapezoids.""" def __init__(self, q=Q, ngrid=129): super().__init__() self.q = float(q) self.grid = torch.linspace(0.0, 1.0, ngrid).view(-1, 1) self.h = nn.Sequential(nn.Linear(1, 12), nn.Tanh(), nn.Linear(12, 1)) def forward(self, s): g = self.grid.to(device=s.device, dtype=s.dtype) hg = torch.clamp(self.h(g), -2.0, 2.0) rho = 1e-4 + torch.clamp(g, min=1e-6).pow(self.q - 1.0) * torch.exp(hg) ds = g[1:] - g[:-1] inc = 0.5 * (rho[1:] + rho[:-1]) * ds R = torch.cat([torch.zeros(1, 1, device=s.device, dtype=s.dtype), torch.cumsum(inc, 0)], 0) R = R / R[-1].clamp_min(1e-12) z = s.clamp(0.0, 1.0) * (len(g) - 1) i = z.long().clamp(0, len(g)-2) t = z - i return R[i, 0] * (1.0-t) + R[i+1, 0] * t class WarpedSystem(nn.Module): def __init__(self): super().__init__() self.field = CoreMLP() self.warp = PositiveRadialWarp() def forward(self, x): z = x.clone() z[:, 0:1] = self.warp(x[:, 0:1]) return self.field(z) def make_system(kind, seed): torch.manual_seed(int(seed)) np.random.seed(int(seed)) return CoreMLP() if kind == 'baseline' else WarpedSystem() def train_one(kind, seed, cfg, capture=False): ds = get_dataset(seed, 400, 400) ds = {k: (torch.from_numpy(v) if isinstance(v, np.ndarray) else v) for k, v in ds.items()} model = make_system(kind, seed) model, metric, _ = train_model(model, ds, epochs=int(cfg['epochs']), lr=float(cfg['lr']), batch=128, weight_decay=float(cfg.get('weight_decay', 0.0)), log=lambda *_: None) if model is None or metric is None: return float('nan'), None sig = None if capture: # Re-test the mathematical prediction on this trained model: estimate the # learned response exponent against the observed transformed coordinate. model.eval() s = torch.logspace(-4, 0, 256).view(-1, 1) ang = torch.zeros_like(s) x = torch.cat([s, torch.ones_like(s), ang], 1) dev = next(model.parameters()).device x, s = x.to(dev), s.to(dev) with torch.no_grad(): pred = model(x).abs().flatten().detach().cpu().numpy() if kind == 'idea': obs_coord = model.warp(s).detach().flatten().cpu().numpy() else: obs_coord = s.flatten().detach().cpu().numpy() mask = (pred > 1e-5) & np.isfinite(pred) & (obs_coord > 1e-5) slope = float(np.polyfit(np.log(obs_coord[mask]), np.log(pred[mask]), 1)[0]) expected = 1.0 # output is u(r)=r^lambda; learned field should be linear in its input radius sig = {'predicted_transformed_field_exponent': expected, 'observed_trained_model_exponent': slope, 'warp_q': Q, 'lambda': LAMBDA, 'confirmed': bool(abs(slope-expected) < 0.35)} return float(metric), sig # Baseline sweep includes every learning rate used by the idea sweep (parity). GRID = [ {'lr': 1e-3, 'epochs': 30}, {'lr': 2e-3, 'epochs': 30}, {'lr': 3e-3, 'epochs': 30}, ] def main(): # The canonical sweep chooses the standard raw-coordinate baseline. base = sweep_baseline(lambda cfg: lambda seed: train_one('baseline', seed, cfg)[0], GRID) # Idea uses the same three configurations, then is evaluated on all eight seeds. idea_by_cfg = [] for cfg in GRID: r = evaluate(lambda seed, cfg=cfg: train_one('idea', seed, cfg)[0]) idea_by_cfg.append({'cfg': cfg, 'result': r}) best = min(idea_by_cfg, key=lambda z: z['result']['mean']) idea = best['result'] # Trained-model behaviour signature, separately measured on all paired seeds. bsigs, isigs = [], [] for seed in range(8): _, bs = train_one('baseline', seed, base['best_cfg'], capture=True) _, ins = train_one('idea', seed, best['cfg'], capture=True) bsigs.append(bs); isigs.append(ins) observed = [x['observed_trained_model_exponent'] for x in isigs if x] sig = { 'prediction': 'after radial warp, u(r)=r^lambda should be approximately linear in warped coordinate raised to lambda; for the trained field, fitted output-vs-warped-radius exponent should be about 1', 'baseline_trained_exponents': [x['observed_trained_model_exponent'] for x in bsigs], 'idea_trained_exponents': observed, 'predicted': 1.0, 'observed_mean': float(np.mean(observed)), 'confirmed': bool(observed and abs(float(np.mean(observed))-1.0) < 0.35), 'note': 'Signature is computed from predictions of independently trained benchmark models.' } report = make_report('singular_radial_poisson', 'mlp_med', base, idea, { 'custom_track': {'name': META['name'], 'file': 'bench_warp.py', 'domain': META['domain']}, **sig, 'idea_sweep': [{'cfg': z['cfg'], 'mean': z['result']['mean'], 'std': z['result']['std']} for z in idea_by_cfg] }) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()