1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7import sys
  8sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  9from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 10
 11SEED = 2852
 12OUT = Path('bench_report.json')
 13
 14class KineticMix(nn.Module):
 15    """Discrete phase-space operator on an 8-token dynamics window.
 16    Token index is x; feature channels are [theta, omega, u].  Diffusion is
 17    depthwise over x and transport shifts each channel along x by h*x_coord.
 18    """
 19    def __init__(self, h=0.08, gate_bias=1.5):
 20        super().__init__()
 21        self.h = float(h)
 22        self.gate = nn.Conv1d(3, 3, 1)
 23        nn.init.zeros_(self.gate.weight)
 24        nn.init.constant_(self.gate.bias, gate_bias)
 25        self.mix = nn.Conv1d(3, 3, 1)
 26
 27    def forward(self, z):  # [B,T,3]
 28        b, t, c = z.shape
 29        q = z.transpose(1, 2)
 30        # normalized x coordinate, matching a characteristic velocity field
 31        xcoord = torch.linspace(-1., 1., t, device=z.device, dtype=z.dtype)
 32        sigma = max(math.sqrt(2.0 * self.h), 1e-3)
 33        radius = max(1, int(math.ceil(3.0 * sigma * t / 2.0)))
 34        offs = torch.arange(-radius, radius + 1, device=z.device)
 35        w = torch.exp(-0.5 * (offs / (sigma * t / 2.0)) ** 2)
 36        w = w / w.sum()
 37        diff = torch.zeros_like(q)
 38        for oi, wi in zip(offs.tolist(), w):
 39            diff = diff + wi * torch.roll(q, int(oi), dims=2)
 40        # periodic linear interpolation at token coordinate x + h*xcoord
 41        pos = torch.arange(t, device=z.device, dtype=z.dtype)[None, :] + self.h * xcoord[None, :] * t / 2.0
 42        pos = pos.remainder(t)
 43        j0 = pos.floor().long().squeeze(0); a = (pos - j0).squeeze(0)
 44        j1 = (j0 + 1).remainder(t)
 45        shifted = diff[:, :, j0] * (1-a)[None, None, :] + diff[:, :, j1] * a[None, None, :]
 46        g = torch.sigmoid(self.gate(q))
 47        return self.mix(q + g * (shifted - q)).transpose(1, 2)
 48
 49class DynamicsNet(nn.Module):
 50    def __init__(self, idea=False, h=.08):
 51        super().__init__()
 52        self.idea = idea
 53        self.kin = KineticMix(h) if idea else nn.Identity()
 54        self.rnn = nn.GRU(3, 32, batch_first=True)
 55        self.head = nn.Linear(32, 1)
 56    def forward(self, x):
 57        z = x.view(x.shape[0], 8, 3)
 58        z = self.kin(z)
 59        _, h = self.rnn(z)
 60        return self.head(h[-1])
 61
 62def seed_all(seed):
 63    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 64    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 65
 66def run_one(seed, idea, lr, h=.08, epochs=14):
 67    seed_all(seed)
 68    ds = get_dataset('dynamics', seed, n_train=400, n_test=400)
 69    # canonical train_model supplies CUDA -> CPU fallback
 70    model = DynamicsNet(idea=idea, h=h)
 71    _, metric, _ = train_model(model, ds, epochs=epochs, lr=lr, batch=64, log=lambda *_: None)
 72    return float(metric)
 73
 74def main():
 75    # Union parity: every lr tested for idea is in baseline sweep.
 76    lrs = [0.0015, 0.003, 0.006]
 77    baseline_grid = [{'lr': x, 'h': None} for x in lrs]
 78    base = sweep_baseline(
 79        lambda cfg: (lambda seed: run_one(seed, False, cfg['lr'], epochs=14)),
 80        baseline_grid)
 81    best_lr = float(base['best_cfg']['lr'])
 82    idea_grid = [{'lr': x, 'h': .08} for x in lrs]
 83    idea_sweep = []
 84    for cfg in idea_grid:
 85        r = evaluate(lambda seed, cfg=cfg: run_one(seed, True, cfg['lr'], cfg['h'], epochs=14), seeds=range(4))
 86        idea_sweep.append({'cfg': cfg, 'mean': r['mean']})
 87    best_idea_cfg = min(idea_sweep, key=lambda x: x['mean'])['cfg']
 88    idea = evaluate(lambda seed: run_one(seed, True, best_idea_cfg['lr'], best_idea_cfg['h'], epochs=14), seeds=range(8))
 89
 90    # Re-test the trained-model mechanism at NN scale: observed Fourier
 91    # attenuation and characteristic centroid displacement of KineticMix.
 92    seed_all(17)
 93    layer = KineticMix(.08).eval()
 94    t = 32; xx = torch.linspace(-1, 1, t)[None, :, None]
 95    k = 3.0; inp = torch.cos(k * xx).expand(1, t, 3).clone()
 96    with torch.no_grad(): out = layer(inp)
 97    # isolate the fixed operator's diffusive response using gate/mix neutralized
 98    raw = layer.gate; mix = layer.mix
 99    with torch.no_grad():
100        layer.gate.weight.zero_(); layer.gate.bias.fill_(30); layer.mix.weight.copy_(torch.eye(3).unsqueeze(-1)); layer.mix.bias.zero_()
101        op = layer(inp)
102    observed_att = float(op[:, :, 0].abs().mean() / inp[:, :, 0].abs().mean())
103    predicted_att = float(math.exp(-.08 * k*k))
104    # transport prediction: use a localized impulse and centroid shift
105    impulse = torch.zeros(1,t,3); impulse[0,t//2,0] = 1
106    with torch.no_grad(): shifted = layer(impulse)
107    mass = shifted[0,:,0].abs(); centroid = float((mass*torch.arange(t)).sum()/(mass.sum()+1e-8))
108    displacement = centroid - t//2
109    signature = {
110        'prediction': 'x Fourier attenuation approximately exp(-h*k^2) and transport displacement proportional to h*x',
111        'predicted_fourier_attenuation': predicted_att,
112        'observed_fourier_attenuation': observed_att,
113        'attenuation_relative_error': abs(observed_att-predicted_att)/(abs(predicted_att)+1e-8),
114        'predicted_transport_displacement': .08 * (t/2),
115        'observed_transport_displacement': displacement,
116        'confirmed': bool(abs(observed_att-predicted_att)/(abs(predicted_att)+1e-8) < .25 and np.isfinite(displacement))
117    }
118    rep = make_report('dynamics', 'rnn_small_kinetic_vs_raw', base, idea,
119                      {'signature': signature, 'idea_sweep': idea_sweep,
120                       'best_idea_cfg': best_idea_cfg, 'parameter_parity': {
121                           'baseline': sum(p.numel() for p in DynamicsNet(False).parameters()),
122                           'idea': sum(p.numel() for p in DynamicsNet(True).parameters())}})
123    rep['baseline']['union_lr_grid'] = lrs
124    rep['custom_track'] = None
125    OUT.write_text(json.dumps(rep, indent=2))
126    print(json.dumps(rep, indent=2))
127
128if __name__ == '__main__':
129    try:
130        main()
131    except RuntimeError as e:
132        print('runtime failure:', repr(e)); raise