1import json, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6from sklearn.datasets import load_digits
  7from sklearn.model_selection import train_test_split
  8from sklearn.preprocessing import StandardScaler
  9
 10SEED = 591
 11random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 12try:
 13    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 14    if device.type == 'cuda': torch.set_per_process_memory_fraction(0.16)
 15except Exception:
 16    device = torch.device('cpu')
 17
 18
 19def math_check():
 20    # x'=-lambda*x: compare explicit Euler with denominator Euler.
 21    lam, alpha = 1.0, 0.5
 22    hs = np.array([0.5, 1., 2., 4., 8., 32.])
 23    q = hs / (1 + alpha * hs)
 24    euler_amp = np.abs(1 - hs * lam)
 25    ns_amp = np.abs(1 - q * lam)
 26    return {
 27        'h': hs.tolist(), 'q': q.tolist(),
 28        'euler_amplification': euler_amp.tolist(),
 29        'nonstandard_amplification': ns_amp.tolist(),
 30        'q_bound_1_over_alpha': 1 / alpha,
 31        'nonstandard_nonexpansive_all_tested': bool(np.all(ns_amp <= 1 + 1e-12)),
 32        'euler_nonexpansive_all_tested': bool(np.all(euler_amp <= 1 + 1e-12))
 33    }
 34
 35
 36class MLP(nn.Module):
 37    def __init__(self, d):
 38        super().__init__()
 39        self.net = nn.Sequential(nn.Linear(d, 32), nn.Tanh(), nn.Linear(32, d))
 40    def forward(self, x):
 41        return self.net(x)
 42
 43
 44class ResidualModel(nn.Module):
 45    def __init__(self, d, blocks, h, kind, alpha=0.5, fixed_iters=6):
 46        super().__init__()
 47        self.blocks = nn.ModuleList([MLP(d) for _ in range(blocks)])
 48        self.h, self.kind, self.alpha, self.fixed_iters = h, kind, alpha, fixed_iters
 49        self.head = nn.Linear(d, 10)
 50        self.last_stats = {'max_activation': 0., 'max_residual': 0., 'failed_solves': 0}
 51
 52    def forward(self, x):
 53        max_act, max_res, failures = 0., 0., 0
 54        for f in self.blocks:
 55            if self.kind == 'baseline':
 56                y = x + self.h * f(x)
 57                residual = (y - x).norm(dim=1).mean().item()
 58            else:
 59                # Two-stage explicit Heun tableau, with denominator q replacing h.
 60                q = self.h / (1 + self.alpha * self.h)
 61                g1 = x
 62                g2 = x + q * f(g1)
 63                for _ in range(self.fixed_iters - 1):
 64                    old = g2
 65                    g2 = x + q * f(g1)  # stage-2 is explicit in this concrete SSP tableau
 66                    err = (g2 - old).norm(dim=1).mean().item()
 67                    max_res = max(max_res, err)
 68                y = x + q * 0.5 * (f(g1) + f(g2))
 69                residual = (y - x).norm(dim=1).mean().item()
 70                if not torch.isfinite(y).all(): failures += 1
 71            x = y
 72            max_act = max(max_act, x.detach().norm(dim=1).max().item())
 73            max_res = max(max_res, residual)
 74        self.last_stats = {'max_activation': max_act, 'max_residual': max_res,
 75                           'failed_solves': failures}
 76        return self.head(x)
 77
 78
 79def run_one(kind, h, Xtr, ytr, Xte, yte, epochs=18):
 80    torch.manual_seed(SEED + int(h * 10) + (0 if kind == 'baseline' else 100))
 81    model = ResidualModel(Xtr.shape[1], blocks=16, h=h, kind=kind).to(device)
 82    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
 83    loss_fn = nn.CrossEntropyLoss()
 84    max_grad, nan = 0., False
 85    t0 = time.time()
 86    for _ in range(epochs):
 87        perm = torch.randperm(len(Xtr), device=device)
 88        for ix in perm.split(128):
 89            opt.zero_grad(set_to_none=True)
 90            out = model(Xtr[ix]); loss = loss_fn(out, ytr[ix])
 91            if not torch.isfinite(loss): nan = True; break
 92            loss.backward()
 93            grad = torch.nn.utils.clip_grad_norm_(model.parameters(), 10.)
 94            max_grad = max(max_grad, float(grad))
 95            opt.step()
 96        if nan: break
 97    with torch.no_grad():
 98        logits = model(Xte); test_loss = float(loss_fn(logits, yte)); acc = float((logits.argmax(1) == yte).float().mean())
 99    return {'loss': test_loss, 'accuracy': acc, 'max_activation': model.last_stats['max_activation'],
100            'max_update_proxy': model.last_stats['max_residual'], 'max_grad': max_grad,
101            'failed_solves': model.last_stats['failed_solves'], 'nan': nan,
102            'seconds': time.time() - t0, 'parameters': sum(p.numel() for p in model.parameters())}
103
104
105def main():
106    data = load_digits()
107    X = StandardScaler().fit_transform(data.data).astype('float32')
108    y = data.target.astype('int64')
109    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=.25, random_state=SEED, stratify=y)
110    tensors = [torch.tensor(z, device=device) for z in (Xtr, ytr, Xte, yte)]
111    Xtr, ytr, Xte, yte = tensors
112    results = {'device': str(device), 'math_check': math_check(), 'runs': {}}
113    for h in (0.5, 1., 2., 4.):
114        for kind in ('baseline', 'idea'):
115            results['runs'][f'{kind}_h{h}'] = run_one(kind, h, Xtr, ytr, Xte, yte)
116            print(kind, h, results['runs'][f'{kind}_h{h}'])
117    Path('results.json').write_text(json.dumps(results, indent=2))
118    print('Wrote results.json')
119
120if __name__ == '__main__': main()