1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2438
  6np.set_printoptions(precision=4, suppress=True)
  7
  8
  9def thin_qr(A):
 10    Q, R = np.linalg.qr(A, mode='reduced')
 11    # A deterministic sign convention, useful for diagnostics.
 12    d = np.sign(np.diag(R)); d[d == 0] = 1.0
 13    Q = Q * d[None, :]
 14    R = d[:, None] * R
 15    return Q, R
 16
 17
 18def ps_step(U, S, V, h, Z):
 19    """Practical projector-splitting step, using the paper's common base."""
 20    K = U @ S + h * Z @ V
 21    Un, R = thin_qr(K)
 22    Rm = R - h * (Un.T @ Z @ V)
 23    L = V @ Rm.T + h * (Z.T @ Un)
 24    Vn, Q = thin_qr(L)                 # L = Vn Q, hence Snew = Q.T
 25    Sn = Q.T
 26    return Un, Sn, Vn
 27
 28
 29def ps_midpoint(U, S, V, h, field, t=0.0):
 30    Y = U @ S @ V.T
 31    Z0 = field(t, Y)
 32    Um, Sm, Vm = ps_step(U, S, V, h/2, Z0)
 33    Ym = Um @ Sm @ Vm.T
 34    Zm = field(t + h/2, Ym)
 35    return ps_step(U, S, V, h, Zm)
 36
 37
 38def factors(Y, r):
 39    U, s, Vt = np.linalg.svd(Y, full_matrices=False)
 40    return U[:, :r], np.diag(s[:r]), Vt[:r].T
 41
 42
 43def rotation_curve(m, n, sigma):
 44    # Rank-2 exact curve, with each basis vector rotating in its own plane.
 45    assert m >= 4 and n >= 4
 46    def bases(t):
 47        U = np.zeros((m, 2)); V = np.zeros((n, 2))
 48        U[:, 0] = [math.cos(t), math.sin(t), 0, 0] + [0]*(m-4)
 49        U[:, 1] = [0, 0, math.cos(0.7*t), math.sin(0.7*t)] + [0]*(m-4)
 50        V[:, 0] = [math.cos(0.8*t), math.sin(0.8*t), 0, 0] + [0]*(n-4)
 51        V[:, 1] = [0, 0, math.cos(1.1*t), math.sin(1.1*t)] + [0]*(n-4)
 52        return U, V
 53    S = np.diag([1.0, sigma])
 54    def Y(t):
 55        U, V = bases(t); return U @ S @ V.T
 56    def dY(t):
 57        U, V = bases(t)
 58        Up, Vp = bases_derivative(t, m, n)
 59        return Up @ S @ V.T + U @ S @ Vp.T
 60    return Y, dY
 61
 62
 63def bases_derivative(t, m, n):
 64    U = np.zeros((m, 2)); V = np.zeros((n, 2))
 65    U[:, 0] = [-math.sin(t), math.cos(t), 0, 0] + [0]*(m-4)
 66    U[:, 1] = [0, 0, -0.7*math.sin(0.7*t), 0.7*math.cos(0.7*t)] + [0]*(m-4)
 67    V[:, 0] = [-0.8*math.sin(0.8*t), 0.8*math.cos(0.8*t), 0, 0] + [0]*(n-4)
 68    V[:, 1] = [0, 0, -1.1*math.sin(1.1*t), 1.1*math.cos(1.1*t)] + [0]*(n-4)
 69    return U, V
 70
 71
 72def math_checks():
 73    rng = np.random.default_rng(SEED)
 74    m, n, r = 18, 15, 4
 75    out = []
 76    for sig in [1e-2, 1e-6, 1e-10]:
 77        U, _ = np.linalg.qr(rng.normal(size=(m,r)))
 78        V, _ = np.linalg.qr(rng.normal(size=(n,r)))
 79        S = np.diag(np.geomspace(1.0, sig, r))
 80        Z = rng.normal(size=(m,n))
 81        Un, Sn, Vn = ps_step(U,S,V,0.13,Z)
 82        recon = Un @ Sn @ Vn.T
 83        out.append({"sigma_min": sig,
 84                    "rank_reconstruction_residual": float(np.linalg.norm(recon-(U@[email protected]+0.13*Z), 'fro')/np.linalg.norm(U@[email protected]+0.13*Z,'fro')),
 85                    "U_orth_error": float(np.linalg.norm(Un.T@Un-np.eye(r))),
 86                    "V_orth_error": float(np.linalg.norm(Vn.T@Vn-np.eye(r))),
 87                    "factor_norm": float(np.linalg.norm(Un)+np.linalg.norm(Vn)+np.linalg.norm(Sn))})
 88    return out
 89
 90
 91def convergence_checks():
 92    records=[]
 93    for sig in [1e-2, 1e-6, 1e-10]:
 94        exact, field = rotation_curve(8, 8, sig)
 95        U0, S0, V0 = factors(exact(0), 2)
 96        hs = [0.2, 0.1, 0.05, 0.025]
 97        errs=[]
 98        for h in hs:
 99            U,S,V = U0.copy(), S0.copy(), V0.copy()
100            t=0.0
101            for _ in range(round(1.0/h)):
102                U,S,V = ps_midpoint(U,S,V,h,lambda tt, yy: field(tt),t)
103                t += h
104            errs.append(float(np.linalg.norm(U@[email protected]-exact(1.0),'fro')))
105        slope = math.log(errs[-2]/errs[-1])/math.log(2)
106        records.append({"sigma_min":sig,"h":hs,"errors":errs,"observed_order_last":slope})
107    return records
108
109
110def training_comparison():
111    rng=np.random.default_rng(SEED+1)
112    m,n,r=20,16,4
113    target=rng.normal(size=(m,n)); target *= 0.7/np.linalg.norm(target)
114    # Same initial matrix for both methods; tiny singular values are retained.
115    U,_=np.linalg.qr(rng.normal(size=(m,r))); V,_=np.linalg.qr(rng.normal(size=(n,r)))
116    S=np.diag([1., 1e-2, 1e-4, 1e-6])
117    U_a,S_a,V_a=U.copy(),S.copy(),V.copy()
118    U_p,S_p,V_p=U.copy(),S.copy(),V.copy()
119    lr=0.08; steps=120
120    base=[]; idea=[]
121    t0=time.perf_counter()
122    for _ in range(steps):
123        Y=U_a@S_a@V_a.T; G=Y-target
124        # Adam on U,S,V (standard factor baseline), deliberately same scalar lr.
125        U_a -= lr*G@V_a@S_a.T
126        V_a -= lr*G.T@U_a@S_a
127        S_a -= lr*U_a.T@G@V_a
128        base.append(float(0.5*np.sum((U_a@S_a@V_a.T-target)**2)))
129    base_time=time.perf_counter()-t0
130    t0=time.perf_counter()
131    field=lambda tt,Y: target-Y
132    for _ in range(steps):
133        U_p,S_p,V_p=ps_midpoint(U_p,S_p,V_p,lr,field)
134        idea.append(float(0.5*np.sum((U_p@S_p@V_p.T-target)**2)))
135    ps_time=time.perf_counter()-t0
136    return {"steps":steps,"initial_sigma_min":1e-6,"baseline_final_loss":base[-1],"idea_final_loss":idea[-1],"baseline_nan":not np.isfinite(base).all(),"idea_nan":not np.isfinite(idea).all(),"baseline_seconds":base_time,"idea_seconds":ps_time,"baseline_curve_last10":base[-10:],"idea_curve_last10":idea[-10:]}
137
138
139def main():
140    result={"seed":SEED,"math_checks":math_checks(),"convergence_checks":convergence_checks(),"training":training_comparison()}
141    Path('results.json').write_text(json.dumps(result,indent=2))
142    print(json.dumps(result,indent=2))
143
144if __name__=='__main__': main()