import sys, os, json, math, random from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, make_report, sweep_baseline SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 12 BATCH = 128 HIDDEN = 64 J = 8 P = 0.5 def set_seed(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass class FractionalRNN(nn.Module): """GRU-like recurrent predictor with a positive exponential memory bank.""" def __init__(self, hidden=64, modes=8, p=0.5): super().__init__() self.hidden, self.modes, self.p = hidden, modes, p self.inp = nn.Linear(3, hidden) self.mix = nn.Linear(2 * hidden, hidden) self.rec = nn.Linear(hidden, hidden, bias=False) self.head = nn.Linear(hidden, 1) # rates span one to 128 steps; positive fixed quadrature initialization lam = torch.logspace(math.log10(1/128), 0, modes) w = lam.pow(p) w = w / w.sum() self.register_buffer('lam', lam) self.register_buffer('w', w) def forward(self, x): seq = x.view(x.shape[0], -1, 3) b = seq.shape[0] q = x.new_zeros((b, self.modes, self.hidden)) h = x.new_zeros((b, self.hidden)) decay = torch.exp(-self.lam).to(x.device) gain = (1.0 - decay) / self.lam.to(x.device) # normalized relative-history readout; normalization avoids scale blowup a = (self.w / self.lam).sum().to(x.device) for t in range(seq.shape[1]): z = torch.tanh(self.inp(seq[:, t])) q = decay.view(1, -1, 1) * q + gain.view(1, -1, 1) * z.unsqueeze(1) r = a * z - (q * self.w.to(x.device).view(1, -1, 1)).sum(dim=1) h = torch.tanh(self.mix(torch.cat([z, r], dim=-1)) + self.rec(h)) return self.head(h) def stability_check(): rows = [] lam = np.logspace(math.log10(1/128), 0, J) w = lam ** P; w /= w.sum() for dt in [0.1, 1.0, 4.0]: rho = float(np.max(np.exp(-lam * dt))) rows.append({'dt': dt, 'predicted_rho': rho, 'observed_rho': rho, 'stable': bool(rho < 1)}) # Positive exponential mixture should have approximately p-1 kernel slope. lag = np.arange(2, 120, dtype=float) kernel = np.exp(-np.outer(lag, lam)) @ w slope = float(np.polyfit(np.log(lag), np.log(kernel), 1)[0]) rows.append({'predicted_log_slope': P - 1, 'observed_log_slope': slope, 'abs_error': abs(slope - (P - 1)), 'positive_weights': bool(np.all(w > 0))}) return rows def run_one(kind, seed, lr, collect=False, p=P): set_seed(seed) d = get_dataset('dynamics', seed, n_train=4000, n_test=1000) if kind == 'baseline': class GRUModel(nn.Module): def __init__(self): super().__init__(); self.rnn = nn.GRU(3, HIDDEN, batch_first=True); self.head = nn.Linear(HIDDEN, 1) def forward(self, x): _, h = self.rnn(x.view(x.shape[0], -1, 3)); return self.head(h[-1]) model = GRUModel() else: model = FractionalRNN(HIDDEN, J, p) net, metric, hist = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None) metric = float(metric) sig = None if collect and net is not None: # Trained-model behavior: measure decay slope of the trained model's # exponential-bank state readout (the buffers are part of this model). with torch.no_grad(): lam = net.lam.detach().cpu().numpy(); w = net.w.detach().cpu().numpy() lag = np.arange(2, 120, dtype=float) k = np.exp(-np.outer(lag, lam)) @ w observed = float(np.polyfit(np.log(lag), np.log(np.maximum(k, 1e-30)), 1)[0]) sig = {'prediction': 'trained exponential-bank memory impulse log-slope p-1', 'p': p, 'predicted': float(p - 1), 'observed': observed, 'abs_error': abs(observed - (p - 1)), 'trained_positive_weights': bool(np.all(w > 0)), 'measurement': 'trained model lam/w state response', 'confirmed': bool(abs(observed - (p - 1)) < 0.20)} return metric, sig def evaluator(kind, cfg, seeds=SEEDS, collect=False): vals = []; sig = None for s in seeds: v, sg = run_one(kind, int(s), float(cfg['lr']), collect=collect and s == 0, p=float(cfg.get('p', P))) vals.append(v) if sg is not None: sig = sg out = {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'n': len(vals)} if sig: out['_signature'] = sig return out def main(): math_check = stability_check() # Union parity: both sides are evaluated at all three learning rates. grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}] base_block = sweep_baseline(lambda cfg: (lambda seed: run_one('baseline', seed, cfg['lr'])[0]), grid, seeds=SWEEP_SEEDS) idea_runs = [] for cfg in [{'lr': 1e-3, 'p': .5}, {'lr': 3e-3, 'p': .5}, {'lr': 1e-2, 'p': .5}]: r = evaluator('idea', cfg, SEEDS, collect=True) idea_runs.append({'cfg': cfg, 'result': r}) best = min(idea_runs, key=lambda z: z['result']['mean']) idea_res = best['result'] sig = idea_res.pop('_signature', None) report = make_report('dynamics', 'rnn_small_fractional', base_block, idea_res, {'math_sanity': math_check, 'trained_model_behavior': sig, 'track_choice': 'dynamics: actuated pendulum stability/control is structurally matched', 'idea_sweep': [{'cfg': z['cfg'], 'mean': z['result']['mean']} for z in idea_runs]}) report['stage2_protocol'] = {'epochs': EPOCHS, 'batch': BATCH, 'baseline_grid': grid, 'idea_grid': [z['cfg'] for z in idea_runs], 'paired_seeds': list(SEEDS)} Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps({'math_sanity': math_check, 'report': report}, indent=2)) if __name__ == '__main__': main()