1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import train_model, evaluate, sweep_baseline, make_report
  9
 10META = {
 11    'name': 'singular_radial_poisson',
 12    'domain': 'pde',
 13    '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.'
 14}
 15LAMBDA = 0.25
 16Q = 2.0
 17
 18
 19def get_dataset(seed, n_train=400, n_test=400):
 20    rng = np.random.RandomState(seed)
 21    def sample(n):
 22        s = np.maximum(rng.rand(n), 1e-5).astype(np.float32)
 23        a = (2.0 * np.pi * rng.rand(n)).astype(np.float32)
 24        # First coordinate is computational radius; the other two encode angle.
 25        x = np.stack([s, np.cos(a), np.sin(a)], axis=1).astype(np.float32)
 26        y = (s ** LAMBDA).astype(np.float32)[:, None]
 27        return x, y
 28    xtr, ytr = sample(n_train)
 29    xte, yte = sample(n_test)
 30    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
 31            'task': 'regression', 'metric': 'mse', 'out_dim': 1,
 32            'input_shape': (3,)}
 33
 34
 35class CoreMLP(nn.Module):
 36    """Shared field architecture for both systems."""
 37    def __init__(self):
 38        super().__init__()
 39        self.net = nn.Sequential(
 40            nn.Linear(3, 64), nn.Tanh(),
 41            nn.Linear(64, 64), nn.Tanh(),
 42            nn.Linear(64, 1))
 43
 44    def forward(self, x):
 45        return self.net(x)
 46
 47
 48class PositiveRadialWarp(nn.Module):
 49    """rho(s)=eps+s^(q-1) exp(h(s)), normalized by cumulative trapezoids."""
 50    def __init__(self, q=Q, ngrid=129):
 51        super().__init__()
 52        self.q = float(q)
 53        self.grid = torch.linspace(0.0, 1.0, ngrid).view(-1, 1)
 54        self.h = nn.Sequential(nn.Linear(1, 12), nn.Tanh(), nn.Linear(12, 1))
 55
 56    def forward(self, s):
 57        g = self.grid.to(device=s.device, dtype=s.dtype)
 58        hg = torch.clamp(self.h(g), -2.0, 2.0)
 59        rho = 1e-4 + torch.clamp(g, min=1e-6).pow(self.q - 1.0) * torch.exp(hg)
 60        ds = g[1:] - g[:-1]
 61        inc = 0.5 * (rho[1:] + rho[:-1]) * ds
 62        R = torch.cat([torch.zeros(1, 1, device=s.device, dtype=s.dtype), torch.cumsum(inc, 0)], 0)
 63        R = R / R[-1].clamp_min(1e-12)
 64        z = s.clamp(0.0, 1.0) * (len(g) - 1)
 65        i = z.long().clamp(0, len(g)-2)
 66        t = z - i
 67        return R[i, 0] * (1.0-t) + R[i+1, 0] * t
 68
 69
 70class WarpedSystem(nn.Module):
 71    def __init__(self):
 72        super().__init__()
 73        self.field = CoreMLP()
 74        self.warp = PositiveRadialWarp()
 75
 76    def forward(self, x):
 77        z = x.clone()
 78        z[:, 0:1] = self.warp(x[:, 0:1])
 79        return self.field(z)
 80
 81
 82def make_system(kind, seed):
 83    torch.manual_seed(int(seed))
 84    np.random.seed(int(seed))
 85    return CoreMLP() if kind == 'baseline' else WarpedSystem()
 86
 87
 88def train_one(kind, seed, cfg, capture=False):
 89    ds = get_dataset(seed, 400, 400)
 90    ds = {k: (torch.from_numpy(v) if isinstance(v, np.ndarray) else v) for k, v in ds.items()}
 91    model = make_system(kind, seed)
 92    model, metric, _ = train_model(model, ds, epochs=int(cfg['epochs']), lr=float(cfg['lr']),
 93                                   batch=128, weight_decay=float(cfg.get('weight_decay', 0.0)),
 94                                   log=lambda *_: None)
 95    if model is None or metric is None:
 96        return float('nan'), None
 97    sig = None
 98    if capture:
 99        # Re-test the mathematical prediction on this trained model: estimate the
100        # learned response exponent against the observed transformed coordinate.
101        model.eval()
102        s = torch.logspace(-4, 0, 256).view(-1, 1)
103        ang = torch.zeros_like(s)
104        x = torch.cat([s, torch.ones_like(s), ang], 1)
105        dev = next(model.parameters()).device
106        x, s = x.to(dev), s.to(dev)
107        with torch.no_grad():
108            pred = model(x).abs().flatten().detach().cpu().numpy()
109            if kind == 'idea':
110                obs_coord = model.warp(s).detach().flatten().cpu().numpy()
111            else:
112                obs_coord = s.flatten().detach().cpu().numpy()
113        mask = (pred > 1e-5) & np.isfinite(pred) & (obs_coord > 1e-5)
114        slope = float(np.polyfit(np.log(obs_coord[mask]), np.log(pred[mask]), 1)[0])
115        expected = 1.0  # output is u(r)=r^lambda; learned field should be linear in its input radius
116        sig = {'predicted_transformed_field_exponent': expected,
117               'observed_trained_model_exponent': slope,
118               'warp_q': Q, 'lambda': LAMBDA,
119               'confirmed': bool(abs(slope-expected) < 0.35)}
120    return float(metric), sig
121
122
123# Baseline sweep includes every learning rate used by the idea sweep (parity).
124GRID = [
125    {'lr': 1e-3, 'epochs': 30},
126    {'lr': 2e-3, 'epochs': 30},
127    {'lr': 3e-3, 'epochs': 30},
128]
129
130
131def main():
132    # The canonical sweep chooses the standard raw-coordinate baseline.
133    base = sweep_baseline(lambda cfg: lambda seed: train_one('baseline', seed, cfg)[0], GRID)
134    # Idea uses the same three configurations, then is evaluated on all eight seeds.
135    idea_by_cfg = []
136    for cfg in GRID:
137        r = evaluate(lambda seed, cfg=cfg: train_one('idea', seed, cfg)[0])
138        idea_by_cfg.append({'cfg': cfg, 'result': r})
139    best = min(idea_by_cfg, key=lambda z: z['result']['mean'])
140    idea = best['result']
141
142    # Trained-model behaviour signature, separately measured on all paired seeds.
143    bsigs, isigs = [], []
144    for seed in range(8):
145        _, bs = train_one('baseline', seed, base['best_cfg'], capture=True)
146        _, ins = train_one('idea', seed, best['cfg'], capture=True)
147        bsigs.append(bs); isigs.append(ins)
148    observed = [x['observed_trained_model_exponent'] for x in isigs if x]
149    sig = {
150        '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',
151        'baseline_trained_exponents': [x['observed_trained_model_exponent'] for x in bsigs],
152        'idea_trained_exponents': observed,
153        'predicted': 1.0,
154        'observed_mean': float(np.mean(observed)),
155        'confirmed': bool(observed and abs(float(np.mean(observed))-1.0) < 0.35),
156        'note': 'Signature is computed from predictions of independently trained benchmark models.'
157    }
158    report = make_report('singular_radial_poisson', 'mlp_med', base, idea, {
159        'custom_track': {'name': META['name'], 'file': 'bench_warp.py', 'domain': META['domain']},
160        **sig,
161        'idea_sweep': [{'cfg': z['cfg'], 'mean': z['result']['mean'], 'std': z['result']['std']} for z in idea_by_cfg]
162    })
163    Path('bench_report.json').write_text(json.dumps(report, indent=2))
164    print(json.dumps(report, indent=2))
165
166if __name__ == '__main__':
167    main()