1import json, random
  2import numpy as np
  3import torch
  4
  5SEED = 7
  6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_num_threads(4)
  8DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
  9try:
 10    torch.zeros(1, device=DEVICE)
 11except Exception:
 12    DEVICE = 'cpu'
 13
 14# A 10D system with an exactly rank-2 invariant state manifold.
 15# x = c + Bq, S(x) = c + B M q.  Thus the faithful latent dimension is 2.
 16def make_map(gamma):
 17    rng = np.random.RandomState(23)
 18    B, _ = np.linalg.qr(rng.randn(10, 2))
 19    B = B.astype(np.float32)
 20    c = np.array([.4,-.2,.3,.1,-.1,.2,.05,-.15,.12,.08], np.float32)
 21    M = np.diag([gamma, .72 * gamma]).astype(np.float32)
 22    A = B @ M @ B.T
 23    b = c - A @ c
 24    return A.astype(np.float32), b.astype(np.float32), B, c, M
 25
 26def apply(A, b, x):
 27    return x @ A.T + b
 28
 29def manifold_states(B, c, n, seed=11):
 30    q = np.random.RandomState(seed).randn(n, 2).astype(np.float32) * 1.4
 31    return c[None, :] + q @ B.T
 32
 33def train_latent(A, b, B, c, m, epochs=1300):
 34    x = manifold_states(B, c, 1800, 11 + m)
 35    y = apply(A, b, x)
 36    E = torch.nn.Linear(10, m, device=DEVICE)
 37    D = torch.nn.Linear(m, 10, device=DEVICE)
 38    T = torch.nn.Linear(m, m, device=DEVICE)
 39    opt = torch.optim.Adam(list(E.parameters()) + list(D.parameters()) + list(T.parameters()), lr=.025)
 40    xt, yt = torch.tensor(x, device=DEVICE), torch.tensor(y, device=DEVICE)
 41    for _ in range(epochs):
 42        idx = torch.randint(0, len(x), (128,), device=DEVICE)
 43        xv, yv = xt[idx], yt[idx]
 44        z, zn = E(xv), T(E(xv))
 45        loss = ((E(yv)-zn)**2).mean() + ((D(zn)-yv)**2).mean() + .5*((D(z)-xv)**2).mean()
 46        opt.zero_grad(); loss.backward(); opt.step()
 47    return E, D, T
 48
 49def learned_metrics(A, b, B, c, E, D, T, gamma):
 50    x = torch.tensor(manifold_states(B, c, 300, 99), device=DEVICE)
 51    y = torch.tensor(apply(A, b, x.cpu().numpy()), device=DEVICE)
 52    with torch.no_grad():
 53        enc = ((E(y)-T(E(x)))**2).mean().sqrt().item()
 54        dec = ((D(T(E(x)))-y)**2).mean().sqrt().item()
 55        rec = ((D(E(x))-x)**2).mean().sqrt().item()
 56        z = E(x); gaps=[]
 57        for _ in range(12):
 58            gaps.append(torch.linalg.vector_norm(D(T(z))-D(z), dim=1).mean().item())
 59            z = T(z)
 60        ratio = float(np.median(np.array(gaps[6:]) / np.array(gaps[5:11])))
 61        fp = c + np.zeros(10, np.float32)  # exact fixed point is c
 62        fp_t = torch.tensor(fp[None, :], device=DEVICE)
 63        fp_res = torch.linalg.vector_norm(D(T(E(fp_t)))-fp_t, dim=1).item()
 64    return {'enc_residual':enc, 'decode_residual':dec, 'reconstruction':rec,
 65            'decoded_fixed_point_residual':fp_res, 'decoded_rate':ratio,
 66            'predicted_rate':gamma}
 67
 68def exact_checks():
 69    # Prediction 1: scalar contraction boundary is gamma*Lambda=gamma < 1.
 70    boundary=[]
 71    for g in [.6,.9,.99,1.,1.01,1.2]:
 72        d=1.;
 73        for _ in range(12): d=g*d
 74        boundary.append({'gamma':g, 'predicted_contracts':g < 1,
 75                         'observed_contracts':d < 1, '12_step_ratio':d})
 76    # Prediction 2: an m-dimensional faithful representation needs m >= rank(M)=2.
 77    A,b,B,c,M=make_map(.82)
 78    x=manifold_states(B,c,500,41)
 79    # exact best rank-m linear reconstruction of centered manifold states
 80    _,s,_=np.linalg.svd((x-c).T, full_matrices=False)
 81    rank=[]
 82    for m in [1,2,3]:
 83        residual=np.sqrt(np.maximum(0., np.sum(s[m:]**2) / x.shape[0])) if m < len(s) else 0.
 84        rank.append({'m':m, 'predicted_zero_residual':m>=2,
 85                     'observed_reconstruction_rmse':float(residual)})
 86    # Prediction 3: ideal decoded trajectory contracts at dominant latent rate gamma.
 87    rates=[]
 88    for g in [.55,.70,.82,.92]:
 89        _,_,_,_,M=make_map(g); q=1.; vals=[]
 90        for _ in range(12): vals.append(abs(q)); q=g*q
 91        rates.append({'gamma':g, 'predicted_asymptotic_rate':g,
 92                      'observed_rate':float(np.median(np.array(vals[6:]) / np.array(vals[5:11])))})
 93    return boundary,rank,rates
 94
 95def main():
 96    boundary, rank, rates = exact_checks()
 97    A,b,B,c,_=make_map(.82)
 98    learned=[]
 99    for m in [1,2,3]:
100        E,D,T=train_latent(A,b,B,c,m)
101        row=learned_metrics(A,b,B,c,E,D,T,.82); row['m']=m; learned.append(row)
102    # Same manifold initial states and 12 update evaluations.
103    x=manifold_states(B,c,300,123); fixed=c
104    direct=x.copy()
105    for _ in range(12): direct=apply(A,b,direct)
106    E,D,T=train_latent(A,b,B,c,2)
107    with torch.no_grad():
108        z=E(torch.tensor(x,device=DEVICE))
109        for _ in range(12): z=T(z)
110        idea=D(z).cpu().numpy()
111    out={'device':DEVICE, 'exact_predictions':{'contraction_boundary':boundary,
112          'rank_threshold':rank, 'rate_transfer':rates}, 'learned_sweep':learned,
113          'comparison':{'baseline_S_evals':12,'idea_S_evals':0,'idea_T_evals':12,
114            'baseline_fixed_point_rmse':float(np.sqrt(np.mean((direct-fixed)**2))),
115            'idea_fixed_point_rmse':float(np.sqrt(np.mean((idea-fixed)**2))),
116            'baseline_output_residual':float(np.sqrt(np.mean((apply(A,b,direct)-direct)**2))),
117            'idea_output_residual':float(np.sqrt(np.mean((apply(A,b,idea)-idea)**2)))}}
118    with open('results.json','w') as f: json.dump(out,f,indent=2)
119    print(json.dumps(out, indent=2))
120
121if __name__ == '__main__': main()