import json, math, os import numpy as np import torch from scipy.special import logsumexp SEED = 2227 np.random.seed(SEED) torch.manual_seed(SEED) torch.set_num_threads(min(8, os.cpu_count() or 1)) def sinkhorn(cost, eps=0.15, iters=250, a=None, b=None): # Log-domain Sinkhorn avoids false gauge violations from kernel clipping. n, m = cost.shape if a is None: a = np.ones(n) / n if b is None: b = np.ones(m) / m logK = -cost / eps logu = np.zeros(n); logv = np.zeros(m) for _ in range(iters): logu = np.log(a) - logsumexp(logK + logv[None, :], axis=1) logv = np.log(b) - logsumexp(logK + logu[:, None], axis=0) return np.exp(logu[:, None] + logK + logv[None, :]) def gauge_matrix(k): # r_i + s_j, with one target potential removed. G = np.zeros((k*k, 2*k-1)) for i in range(k): for j in range(k): q = i*k+j G[q, i] = 1 if j < k-1: G[q, k+j] = 1 return G def gauge_project(Phi, G): # Orthogonal projection in pair/cost space. Q = np.eye(G.shape[0]) - G @ np.linalg.pinv(G) A = Q @ Phi U0, s, Vt = np.linalg.svd(A, full_matrices=False) rank = int(np.sum(s > max(A.shape)*np.finfo(float).eps*s[0])) if s[0] else 0 U = Vt[:rank].T return Q, A, U, s, rank def make_features(k=8, scales=(1.0, .1, .01), gauge_amp=1.0, seed=0): rng = np.random.default_rng(seed) G = gauge_matrix(k) # Three identifiable random pair features with prescribed anisotropy. R = rng.normal(size=(k*k, 3)) @ np.diag(scales) # Remove accidental gauge components so the quotient is controlled. Q = np.eye(k*k) - G @ np.linalg.pinv(G) R = Q @ R # Exact row/column gauge directions, mixed into columns to mimic nuisance biases. B = G[:, :5] * gauge_amp Phi = np.concatenate([B, R], axis=1) return Phi, G def jacobian_plan(Phi, theta, k, eps=.15): # Finite difference is intentionally independent of the whitening code. base = sinkhorn((-Phi @ theta).reshape(k,k), eps) J = np.zeros((k*k, Phi.shape[1])) h = 1e-5 for f in range(Phi.shape[1]): t = theta.copy(); t[f] += h p = sinkhorn((-Phi @ t).reshape(k,k), eps) J[:, f] = ((p-base)/h).ravel() return J def prediction_sweeps(): # P1: exact gauge invariance over several amplitudes. k=8; eps=.15 Phi, G = make_features(k, gauge_amp=1.0, seed=4) rng=np.random.default_rng(5); theta=rng.normal(size=Phi.shape[1]) gauge_theta=np.zeros(Phi.shape[1]); gauge_theta[:5]=rng.normal(size=5) inv=[] for amp in [0., .1, 1., 10., 100.]: p=sinkhorn((-Phi@(theta+amp*gauge_theta)).reshape(k,k),eps) p0=sinkhorn((-Phi@theta).reshape(k,k),eps) inv.append(float(np.max(np.abs(p-p0)))) # P2: quotient removes exactly the gauge-induced zero Jacobian spectrum. J=jacobian_plan(Phi,theta,k,eps) sv=np.linalg.svd(J,compute_uv=False) Q,A,U,s,rank=gauge_project(Phi,G) Jr=jacobian_plan(A@U, np.linalg.lstsq(A@U, Q@Phi@theta,rcond=None)[0], k, eps) svr=np.linalg.svd(Jr,compute_uv=False) zero_count=int(np.sum(sv < 1e-8)) # P3: covariance conditioning scales quadratically with feature anisotropy, # while whitening gives approximately unit covariance (delta regularizer). rows=[] for alpha in [1., .3, .1, .03, .01]: P,G2=make_features(k, scales=(1.,alpha,alpha*alpha), gauge_amp=3., seed=8) Q,A,U,_,rank=gauge_project(P,G2) C=(A@U).T@(A@U)/(k*k) ev=np.linalg.eigvalsh(C) cond=float(ev[-1]/max(ev[0],1e-15)) delta=1e-6 # Symmetric eigendecomposition implements (C+delta I)^(-1/2). ce, cv=np.linalg.eigh(C+delta*np.eye(rank)) W=(cv*(1/np.sqrt(ce)))@cv.T Cw=W.T@C@W ew=np.linalg.eigvalsh(Cw) ridge_pred=float(ev[-1]*(ev[0]+delta)/(ev[0]*(ev[-1]+delta))) rows.append({'anisotropy_alpha':alpha,'raw_cov_condition':cond, 'predicted_condition_scaling':float(1/(alpha**4)), 'ridge_aware_whitening_prediction':ridge_pred, 'whitened_cov_condition':float(ew[-1]/ew[0]), 'cov_eigenvalues':ev.tolist(),'rank':rank}) return {'gauge_invariance_max_abs_plan_change':inv, 'gauge_jacobian_singular_values':sv.tolist(), 'gauge_jacobian_near_zero_count':zero_count, 'quotient_rank_predicted':rank, 'quotient_jacobian_singular_values':svr.tolist(), 'conditioning_sweep':rows} def train_compare(steps=180, trials=6): # Small supervised matching: target is diagonal assignment, uniform marginals. k=8; eps=.12 losses={'raw':[], 'gauge_free_whitened':[]} final={'raw':[], 'gauge_free_whitened':[]} for trial in range(trials): Phi,G=make_features(k, scales=(1.,.15,.02), gauge_amp=5., seed=100+trial) Q,A,U,_,rank=gauge_project(Phi,G) C=(A@U).T@(A@U)/(k*k); delta=1e-5 # symmetric eig whitening, with theta=U W z and cost Phi theta. ew,ev=np.linalg.eigh(C+delta*np.eye(rank)); W=(ev*(1/np.sqrt(ew)))@ev.T target=torch.tensor(np.eye(k)/k,dtype=torch.float32) def run(kind): z=torch.zeros(rank if kind=='gauge_free_whitened' else Phi.shape[1], requires_grad=True) opt=torch.optim.Adam([z],lr=.08) vals=[] Pt=torch.tensor(Phi,dtype=torch.float32) if kind=='gauge_free_whitened': T=torch.tensor(U@W,dtype=torch.float32) for st in range(steps): opt.zero_grad() theta=(T@z if kind=='gauge_free_whitened' else z) cost=-(Pt@theta).reshape(k,k) # log-domain-ish stabilized Sinkhorn in torch. logK=-cost/eps logu=torch.zeros(k); logv=torch.zeros(k) for _ in range(35): logu=-torch.logsumexp(logK+logv[None,:],dim=1)-math.log(k) logv=-torch.logsumexp(logK+logu[:,None],dim=0)-math.log(k) plan=torch.exp(logu[:,None]+logK+logv[None,:]) loss=((plan-target)**2).mean() loss.backward(); opt.step() vals.append(float(loss.detach())) return vals for kind in losses: losses[kind].append(run(kind)) for kind in losses: final[kind].append(losses[kind][-1][-1]) mean_curves={k:np.mean(v,axis=0).tolist() for k,v in losses.items()} return {'final_loss_mean':{k:float(np.mean(v)) for k,v in final.items()}, 'final_loss_std':{k:float(np.std(v)) for k,v in final.items()}, 'loss_at_steps':{k:{str(s):float(np.mean([x[s-1] for x in v])) for s in [10,30,60,120,180]} for k,v in losses.items()}, 'curves':mean_curves} def main(): out={'seed':SEED,'predictions':prediction_sweeps(),'training':train_compare()} with open('results.json','w') as f: json.dump(out,f,indent=2) print(json.dumps(out,indent=2)) if __name__=='__main__': main()