Adaptive Physics-Lifted Koopman State Space / experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
  1import json, math
  2import numpy as np
  3
  4SEED = 7
  5
  6
  7def rls_update(theta, P, r, y, lam):
  8    """Forgetting-factor RLS, theta has shape (outputs, inputs)."""
  9    Pr = P @ r
 10    Pn = (P - np.outer(Pr, Pr) / (lam + r @ Pr)) / lam
 11    en = y - theta @ r
 12    return theta + np.outer(en, r @ Pn), Pn
 13
 14
 15def lift(x):
 16    return np.array([x, x * x], dtype=float)
 17
 18
 19def dynamics(x, regime):
 20    return (0.62 * x + 0.22 * x*x - 0.10 * x**3) if regime == 0 else (0.35 * x + 0.42 * x*x - 0.10 * x**3)
 21
 22
 23def stream(n=500, change=250):
 24    rng = np.random.default_rng(SEED)
 25    x, data = 0.15, []
 26    for t in range(n):
 27        u = .10 * math.sin(.17*t) + .025 * rng.normal()
 28        reg = int(t >= change)
 29        y = dynamics(x, reg) + u
 30        data.append((x, u, y, reg))
 31        x = np.clip(y, -1.4, 1.4)
 32    return data
 33
 34
 35def train(data, lam, lifted, updates=True):
 36    # Linear baseline models x' from [x,u,1]; lifted model z' from [x,x^2,u,1].
 37    def inp(x, u): return np.r_[lift(x) if lifted else x, u, 1.0]
 38    def out(y): return lift(y) if lifted else np.array([y])
 39    din, dout = (4, 2) if lifted else (3, 1)
 40    th, P = np.zeros((dout, din)), np.eye(din) * 100.
 41    for x,u,y,_ in data[:250]:
 42        if updates: th,P = rls_update(th,P,inp(x,u),out(y),lam)
 43    before = th.copy()
 44    errs, rhos = [], []
 45    for x,u,y,_ in data[250:]:
 46        if updates: th,P = rls_update(th,P,inp(x,u),out(y),lam)
 47        pred = th @ inp(x,u)
 48        errs.append((pred[0] - y)**2)
 49        rhos.append(float(max(abs(np.linalg.eigvals(th[:,:2])))) if lifted else abs(float(th[0,0])))
 50    return np.asarray(errs), np.asarray(rhos), before, th
 51
 52
 53def rollout(th, x0, horizon, regime):
 54    z = lift(x0); truth = x0; pred = []
 55    for _ in range(horizon):
 56        # no external input; last column is the intercept
 57        zp = th @ np.r_[z, 0.0, 1.0]
 58        z = zp
 59        truth = dynamics(truth, regime)
 60        pred.append((z[0], truth))
 61    return float(np.mean([(a-b)**2 for a,b in pred]))
 62
 63
 64def math_checks():
 65    # Exact finite-horizon identity: prior precision is lambda^N P0^-1.
 66    rng = np.random.default_rng(11); d, n, lam = 3, 40, .93
 67    th, P = np.zeros((1,d)), np.eye(d)*4.0
 68    X, Y = [], []
 69    for _ in range(n):
 70        r = rng.normal(size=d); y = np.array([.4*r[0]-.2*r[1]+.1*r[2]])
 71        th,P = rls_update(th,P,r,y,lam); X.append(r); Y.append(y)
 72    X, Y = np.asarray(X), np.asarray(Y).reshape(-1,1)
 73    W = np.diag(lam**np.arange(n-1,-1,-1))
 74    batch = (Y.T@W@X) @ np.linalg.inv(X.T@W@X + (lam**n)*np.eye(d)/4.0)
 75    # 63% exponential forgetting mass and geometric contraction.
 76    timescale = 1/(1-lam)
 77    measured = int(np.argmin(abs((1-lam**np.arange(1,200))-(1-1/math.e)))+1)
 78    return {"rls_batch_max_error": float(np.max(abs(th-batch))),
 79            "adaptation_timescale": timescale, "measured_63pct_steps": measured,
 80            "rho_0.8_power29": .8**29}
 81
 82
 83def stability_check():
 84    # Scaling a stable scalar lifted operator gives the claimed rho=1 boundary.
 85    vals=[]
 86    for s in [.8, .99, 1.01, 1.2]:
 87        rho=abs(s); norm=abs(s**40)
 88        vals.append({"scale":s,"rho":rho,"power40_norm":norm})
 89    return vals
 90
 91
 92def main():
 93    data=stream(); result={"seed":SEED,"math":math_checks(),"stability_boundary":stability_check()}
 94    for lam in (.98,.95,.90):
 95        for lifted in (False,True):
 96            e,r,_,_=train(data,lam,lifted,True)
 97            result[f"adaptive_{'lifted' if lifted else 'linear'}_lambda_{lam}"]={
 98                "post_shift_rmse_first20":float(np.sqrt(np.mean(e[:20]))),
 99                "post_shift_rmse_last100":float(np.sqrt(np.mean(e[-100:]))),
100                "max_rho":float(np.max(r))}
101    for lifted in (False,True):
102        e,_,_,_=train(data,.95,lifted,False)
103        result[f"frozen_{'lifted' if lifted else 'linear'}"]={
104            "post_shift_rmse_first20":float(np.sqrt(np.mean(e[:20]))),
105            "post_shift_rmse_last100":float(np.sqrt(np.mean(e[-100:]))) }
106    _,_,frozen,adapted=train(data,.95,True,True)
107    result["50step_post_shift_rollout_mse"]={"adaptive_lifted":rollout(adapted,.45,50,1),"frozen_lifted":rollout(frozen,.45,50,1)}
108    with open("results.json","w") as f: json.dump(result,f,indent=2)
109    print(json.dumps(result,indent=2))
110
111if __name__ == '__main__': main()