Robust Physics-Sparse Neural Dynamics / experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json
  2import numpy as np
  3from pathlib import Path
  4
  5RNG = np.random.default_rng(1623)
  6
  7# Damped oscillator: eigenvalues -alpha +/- i beta.  The known relaxation is -alpha*x;
  8# the sparse library learns the coupling terms [x1, x2].
  9alpha, beta, h = 0.20, 0.90, 0.08
 10M = np.array([[-alpha, -beta], [beta, -alpha]], dtype=float)
 11
 12
 13def trajectory(x0, n=260, noise=0.0):
 14    x = np.zeros((n, 2)); x[0] = x0
 15    D = np.eye(2) + h*M
 16    for k in range(n-1): x[k+1] = D @ x[k]
 17    if noise: x += RNG.normal(0, noise, x.shape)
 18    return x
 19
 20
 21def make_data(ntraj=24, n=180, obs_noise=0.003, outlier_frac=0.05):
 22    X, Y = [], []
 23    for _ in range(ntraj):
 24        clean = trajectory(RNG.normal(0, 1, 2), n, 0)
 25        obs = clean + RNG.normal(0, obs_noise, clean.shape)
 26        # centered derivative, with the same observation corruption mechanism
 27        der = (obs[2:] - obs[:-2])/(2*h)
 28        lib = obs[1:-1]
 29        # Known physics: f_known=-alpha*x. Learn only cross-state interactions.
 30        A_lib = np.column_stack([lib[:, 1], lib[:, 0]])
 31        residual_target = der + alpha*lib
 32        X.append(A_lib); Y.append(residual_target)
 33    A, b = np.vstack(X), np.vstack(Y)
 34    bad = RNG.choice(len(A), int(outlier_frac*len(A)), replace=False)
 35    b[bad] += RNG.normal(0, 0.7, (len(bad),2))
 36    return A, b, bad
 37
 38
 39def ols(A,b):
 40    return np.linalg.lstsq(A, b, rcond=None)[0]
 41
 42
 43def tls(A,b):
 44    # Total least squares independently for each derivative coordinate.
 45    out = np.zeros((A.shape[1], b.shape[1]))
 46    Z = np.column_stack([A, b])
 47    for j in range(b.shape[1]):
 48        # For multivariate b, the standard TLS null vector is applied to [A,b_j].
 49        _, _, vh = np.linalg.svd(np.column_stack([A, b[:,j]]), full_matrices=False)
 50        v = vh[-1]
 51        out[:,j] = -v[:-1]/v[-1]
 52    return out
 53
 54
 55def robust_tls(A,b, seed=4, rounds=500, threshold=0.08):
 56    rng = np.random.default_rng(seed); p=A.shape[1]
 57    best_in = None; best_score=-1
 58    m=min(len(A), p+3)
 59    for _ in range(rounds):
 60        idx=rng.choice(len(A),m,replace=False)
 61        try: c=tls(A[idx],b[idx])
 62        except np.linalg.LinAlgError: continue
 63        residual=np.linalg.norm(b-A@c,axis=1)
 64        # normalized residual makes threshold comparable across trials
 65        inliers=residual < threshold
 66        score=int(inliers.sum())
 67        if score>best_score: best_score, best_in=score,inliers
 68    if best_in is None: return tls(A,b), np.ones(len(A),bool)
 69    # TLS consensus refit, followed by iterative hard thresholding (sparsity step).
 70    c=tls(A[best_in],b[best_in])
 71    scale=np.maximum(np.max(np.abs(c),axis=0,keepdims=True),1e-9)
 72    c[np.abs(c)<0.08*scale]=0
 73    return c,best_in
 74
 75
 76def spectral_gain_boundary(M, hh):
 77    # J(g)=I+g*h*M; solve rho(J)=1 for the first positive crossing.
 78    eig=np.linalg.eigvals(M)
 79    candidates=[]
 80    for z in eig:
 81        if abs(z.imag)>1e-10:
 82            candidates.append(-2*z.real/(hh*abs(z)**2))
 83        elif z.real<0: candidates.append(-2/(hh*z.real))
 84    return min(candidates)
 85
 86
 87def rho(M, gain): return max(abs(np.linalg.eigvals(np.eye(2)+gain*h*M)))
 88
 89
 90def rollout_error(C, gain=1.0, n=100):
 91    # C maps row features to derivative targets; state-column dynamics use C.T.
 92    D=np.eye(2)+gain*h*C
 93    x=np.array([1.,-.4]); norms=[]
 94    for _ in range(n):
 95        norms.append(np.linalg.norm(x)); x=D@x
 96    return float(norms[-1]), float(max(norms))
 97
 98
 99def main():
100    # Math sanity: exact boundary and a sweep around it.
101    pred=spectral_gain_boundary(M,h)
102    gain_grid=np.linspace(0.25, 1.35*pred, 45)
103    measured=[]
104    for g in gain_grid:
105        # observed boundedness over 1000 steps, with a tiny nonzero initial state
106        D=np.eye(2)+g*h*M; x=np.array([1.,0.]);
107        for _ in range(1000): x=D@x
108        measured.append(np.linalg.norm(x)<1e6)
109    crossing=next((gain_grid[i] for i,v in enumerate(measured) if not v), gain_grid[-1])
110    # Refine empirical boundary by direct spectral test on a dense grid.
111    dense=np.linspace(0.1,1.3*pred,5000)
112    obs=dense[np.argmax([rho(M,g)>=1 for g in dense])]
113
114    A,b,bad=make_data()
115    c_ols=ols(A,b); c_tls=tls(A,b); c_rob,inliers=robust_tls(A,b)
116    q=np.array([[-beta,0.0],[0.0,beta]])
117    errs={"OLS":float(np.linalg.norm(c_ols-q)/np.linalg.norm(q)),
118          "TLS":float(np.linalg.norm(c_tls-q)/np.linalg.norm(q)),
119          "TLS_RANSAC":float(np.linalg.norm(c_rob-q)/np.linalg.norm(q))}
120    # Held-out clean 100-step trajectory one-step model rollout.
121    def hybrid_matrix(c):
122        # A=[x2,x1], so the learned interaction matrix is reconstructed here.
123        return np.array([[-alpha, c[0,0]], [c[1,1], -alpha]])
124    roll={}
125    for name,c in [("OLS",c_ols),("TLS_RANSAC",c_rob)]: roll[name]=rollout_error(hybrid_matrix(c))
126    roll["true"]=rollout_error(M)
127
128    # Step-size scaling prediction: g* = 2 alpha/(h(alpha^2+beta^2)), measured by eigenvalue sweep.
129    scaling=[]
130    for hh in [0.04,0.08,0.16,0.24]:
131        p=spectral_gain_boundary(M,hh)
132        gs=np.linspace(.1,1.15*p,4000)
133        measured_h=gs[np.argmax([max(abs(np.linalg.eigvals(np.eye(2)+g*hh*M)))>=1 for g in gs])]
134        scaling.append({"h":hh,"predicted":p,"observed":float(measured_h),"relative_error":float(abs(measured_h-p)/p)})
135    result={"parameters":{"alpha":alpha,"beta":beta,"h":h,"outlier_fraction":len(bad)/len(A)},
136      "stability_boundary":{"predicted_gain":pred,"observed_gain":float(obs),"coarse_rollout_crossing":float(crossing),"relative_error":float(abs(obs-pred)/pred),"stable_at_0.9":bool(rho(M,.9*pred)<1),"unstable_at_1.1":bool(rho(M,1.1*pred)>1)},
137      "step_scaling":scaling,
138      "coefficient_relative_error":errs,
139      "inlier_fraction":float(inliers.mean()),"rollout_norm_100":roll}
140    Path("results.json").write_text(json.dumps(result,indent=2))
141    print(json.dumps(result,indent=2))
142
143if __name__=="__main__": main()