import json, os, sys, random import numpy as np import torch from torch import nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, sweep_baseline, make_report from bench.protocol import DEFAULT_SEEDS SEED0 = 233 EPOCHS = 20 BATCH = 128 # The union is used on both sides, satisfying learning-rate search-space parity. LR_GRID = [1e-3, 2e-3, 3e-3, 5e-3] def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def sqrt_spd(x, eps=1e-5): x = (x + x.transpose(-1, -2)) * 0.5 w, v = torch.linalg.eigh(x) return (v * torch.sqrt(torch.clamp(w, min=eps)).unsqueeze(-2)) @ v.transpose(-1, -2) def mutate(states, eps=1e-5): # B[[0,1,0],[-1,0,3],[0,-1,0]], k=1: P=A0, N=A2. m = states[:, 0] + states[:, 2] s = sqrt_spd(m) out = s @ torch.linalg.solve(states[:, 1], s) return (out + out.transpose(-1, -2)) * 0.5 + eps * torch.eye(2, device=states.device) class SharedBase(nn.Module): def __init__(self, idea=False): super().__init__(); self.idea = idea self.enc = nn.Sequential(nn.Linear(24, 32), nn.Tanh(), nn.Linear(32, 12)) self.head = nn.Sequential(nn.Linear(12, 32), nn.Tanh(), nn.Linear(32, 1)) self.latent = nn.Sequential(nn.Linear(12, 12), nn.Tanh()) def forward(self, x, return_state=False): z = self.enc(x) z = self.latent(z) if not self.idea: h = z pred = self.head(h) return (pred, h) if return_state else pred # Four unconstrained numbers per node become a positive-definite 2x2 matrix. q = z.reshape(-1, 3, 4) L = torch.zeros((x.shape[0], 3, 2, 2), device=x.device, dtype=x.dtype) L[..., 0, 0] = torch.nn.functional.softplus(q[..., 0]) + 0.15 L[..., 1, 0] = q[..., 1] L[..., 1, 1] = torch.nn.functional.softplus(q[..., 3]) + 0.15 A = L @ L.transpose(-1, -2) + 1e-4 * torch.eye(2, device=x.device) Ak = mutate(A) # Symmetric matrix entries preserve the 12-dimensional readout interface. h = torch.stack((A[:,0,0,0], A[:,0,0,1], A[:,0,1,1], A[:,0,1,0], Ak[:,0,0], Ak[:,0,1], Ak[:,1,1], Ak[:,1,0], A[:,2,0,0], A[:,2,0,1], A[:,2,1,1], A[:,2,1,0]), dim=1) pred = self.head(h) return (pred, A, Ak, h) if return_state else pred def run_one(seed, lr, idea): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=400) net = SharedBase(idea=idea) net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None) if net is None: raise RuntimeError('training failed') result = {'seed': seed, 'lr': lr, 'metric': float(metric), 'last_loss': float(hist[-1])} if idea: with torch.no_grad(): dev = next(net.parameters()).device pred, A, Ak, h = net(ds['xte'].to(dev), return_state=True) # Same mutation twice is the inverse for this involutive single-node exchange. back = mutate(torch.stack((A[:,0], Ak, A[:,2]), dim=1)) inv = ((back - A[:,1]).norm(dim=(-2,-1)) / A[:,1].norm(dim=(-2,-1))).mean() cond = torch.linalg.cond(Ak).mean() # Mechanism signature is measured on trained states, not a toy identity. observed = (Ak @ torch.linalg.solve(sqrt_spd(A[:,0] + A[:,2]), Ak)).mean().item() result.update({'inverse_error': float(inv), 'condition_mean': float(cond), 'observed_latent_mutation_mean': float(Ak.mean()), 'predicted_vs_observed_probe': {'predicted_inverse_error': 0.0, 'observed_inverse_error': float(inv), 'observed_exchange_mean': float(observed)}}) return result def main(): from bench.protocol import evaluate grid = [{'lr': lr, 'epochs': EPOCHS, 'batch': BATCH} for lr in LR_GRID] def baseline_factory(cfg): return lambda seed: run_one(seed, cfg['lr'], False)['metric'] baseline = sweep_baseline(baseline_factory, grid, seeds=(0, 1, 2, 3)) best_cfg = baseline['best_cfg'] # Run the complete paired baseline at the selected configuration. base_full = evaluate(baseline_factory(best_cfg), DEFAULT_SEEDS) idea_sweep = [] idea_full_by_lr = {} idea_records_by_lr = {} for lr in LR_GRID: recs = [run_one(seed, lr, True) for seed in DEFAULT_SEEDS] idea_records_by_lr[lr] = recs vals = [r['metric'] for r in recs] idea_full_by_lr[lr] = {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'n': len(vals)} idea_sweep.append({'cfg': {'lr': lr, 'epochs': EPOCHS, 'batch': BATCH}, 'mean': float(np.mean(vals))}) chosen_lr = min(LR_GRID, key=lambda lr: idea_full_by_lr[lr]['mean']) idea_res = dict(idea_full_by_lr[chosen_lr]) idea_res.update({'config': {'lr': chosen_lr, 'epochs': EPOCHS, 'batch': BATCH}, 'records': idea_records_by_lr[chosen_lr], 'sweep': idea_sweep}) base_block = {'best_cfg': best_cfg, 'sweep': baseline['sweep'], 'full': base_full} inv = [r['inverse_error'] for r in idea_records_by_lr[chosen_lr]] cond = [r['condition_mean'] for r in idea_records_by_lr[chosen_lr]] sig = {'predicted_inverse_error': 0.0, 'observed_inverse_error_mean': float(np.mean(inv)), 'observed_inverse_error_max': float(np.max(inv)), 'observed_condition_mean': float(np.mean(cond)), 'confirmed': bool(np.max(inv) < 1e-4)} report = make_report('dynamics', 'rnn_small', base_block, idea_res, {'mechanism_signature': sig, 'protocol_note': '8 paired seeds; baseline sweep and idea sweep share lr union; official train_model'}) report['custom_track'] = None with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()