1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 573
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)
 10
 11
 12def wrap(x):
 13    return (x + 0.5) % 1.0 - 0.5
 14
 15
 16def graph_cochains(theta, k=12):
 17    # theta: N x m angles in [0,1); torus distance uses the first two physical angles
 18    N, m = theta.shape
 19    d = np.sqrt(np.sum(wrap(theta[:, None, :2] - theta[None, :, :2])**2, axis=2))
 20    edges = []
 21    seen = set()
 22    for i in range(N):
 23        for j in np.argsort(d[i])[1:k+1]:
 24            a,b = sorted((i,int(j)))
 25            if a != b and (a,b) not in seen:
 26                seen.add((a,b)); edges.append((a,b))
 27    A = np.array([wrap(theta[j]-theta[i]) for i,j in edges])
 28    w = np.ones(len(edges))/len(edges)
 29    G = A.T @ (w[:,None]*A)
 30    return G, A, edges
 31
 32
 33def select_matroid(G, Q, tol=1e-8):
 34    # Q is m x r; v_j are rows of Q, and M=Q^T G Q.
 35    M = Q.T @ G @ Q
 36    V = Q.T                         # r x m, columns are candidate vectors
 37    energies = np.einsum('ir,rs,is->i', Q, M, Q)
 38    order = np.argsort(energies)
 39    chosen=[]; rank=0
 40    for j in order:
 41        trial = V[:, chosen+[int(j)]]
 42        s = np.linalg.svd(trial, compute_uv=False)
 43        newrank = int(np.sum(s > tol * max(1.0, s[0])))
 44        if newrank > rank:
 45            chosen.append(int(j)); rank=newrank
 46        if rank == Q.shape[1]: break
 47    return chosen, energies, M
 48
 49
 50def residual(G, Q, chosen):
 51    # Persistent-space residual as stated, with V made from selected projected vectors.
 52    M = Q.T @ G @ Q
 53    V = Q.T[:, chosen]
 54    P = V @ np.linalg.pinv(V.T @ M @ V) @ V.T @ M if chosen else np.zeros_like(M)
 55    R = np.eye(M.shape[0]) - P
 56    return np.linalg.norm(R, 'fro'), R
 57
 58class AE(nn.Module):
 59    def __init__(self):
 60        super().__init__(); self.enc=nn.Sequential(nn.Linear(4,32),nn.Tanh(),nn.Linear(32,2)); self.dec=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,4))
 61    def forward(self,x): return self.dec(self.enc(x))
 62class CircDecoder(nn.Module):
 63    def __init__(self,d):
 64        super().__init__(); self.net=nn.Sequential(nn.Linear(2*d,32),nn.Tanh(),nn.Linear(32,4))
 65    def forward(self,z): return self.net(z)
 66
 67def train_models(x, selected, theta):
 68    device='cuda' if torch.cuda.is_available() else 'cpu'
 69    try:
 70        xt=torch.tensor(x,dtype=torch.float32,device=device)
 71        # Circular selected features are fixed, as prescribed by the dictionary.
 72        z=np.concatenate([np.cos(2*np.pi*theta[:,selected]), np.sin(2*np.pi*theta[:,selected])],axis=1)
 73        zt=torch.tensor(z,dtype=torch.float32,device=device)
 74        base=AE().to(device); circ=CircDecoder(len(selected)).to(device)
 75        ob=torch.optim.Adam(base.parameters(),lr=2e-3); oc=torch.optim.Adam(circ.parameters(),lr=2e-3)
 76        for _ in range(1000):
 77            ob.zero_grad(); loss=((base(xt)-xt)**2).mean(); loss.backward(); ob.step()
 78            oc.zero_grad(); loss2=((circ(zt)-xt)**2).mean(); loss2.backward(); oc.step()
 79        with torch.no_grad():
 80            br=float(((base(xt)-xt)**2).mean().cpu()); cr=float(((circ(zt)-xt)**2).mean().cpu())
 81            # decode angles from the circular pairs; scalar AE has no canonical angle decoder
 82            decoded=np.mod(np.arctan2(z[:,len(selected):],z[:,:len(selected)])/(2*np.pi),1.0)
 83            # compare each selected physical candidate to its exact dictionary angle
 84            angle_err=float(np.mean(np.abs(wrap(decoded-theta[:,selected]))))
 85        return br,cr,angle_err,device
 86    except Exception as e:
 87        print('device fallback:',repr(e))
 88        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 89        # CPU retry
 90        old=torch.cuda.is_available
 91        torch.cuda.is_available=lambda: False
 92        out=train_models(x,selected,theta)
 93        torch.cuda.is_available=lambda: old
 94        return out
 95
 96def main():
 97    N=900
 98    t=np.random.rand(N,2)
 99    # physical torus embedded in R4; candidate angle dictionary includes mixtures and distractors
100    x=np.concatenate([np.cos(2*np.pi*t[:,0,None]),np.sin(2*np.pi*t[:,0,None]),np.cos(2*np.pi*t[:,1,None]),np.sin(2*np.pi*t[:,1,None])],axis=1)
101    theta=np.column_stack([t[:,0],t[:,1],(t[:,0]+t[:,1])%1,(t[:,0]-t[:,1])%1,np.random.rand(N)])
102    G,A,edges=graph_cochains(theta)
103    # Q explicitly represents detected persistent H1 directions: candidates 0 and 1 span them.
104    Q=np.zeros((5,2)); Q[0,0]=1; Q[1,1]=1; Q[2]=[.5,.5]; Q[3]=[.5,-.5]; Q[4]=[.03,.02]
105    selected,E,M=select_matroid(G,Q)
106    rfull,_=residual(G,Q,[0,1]); romit,_=residual(G,Q,[0])
107    br,cr,ae,device=train_models(x,selected,theta)
108    out={'seed':SEED,'edges':len(edges),'gram_trace':float(np.trace(G)),'gram_eigenvalues':np.linalg.eigvalsh(G).tolist(),'energies':E.tolist(),'selected':selected,'residual_full':rfull,'residual_omit_one':romit,'baseline_reconstruction_mse':br,'idea_reconstruction_mse':cr,'idea_selected_angle_mean_wrapped_error':ae,'device':device}
109    Path('results.json').write_text(json.dumps(out,indent=2))
110    print(json.dumps(out,indent=2))
111if __name__=='__main__': main()