Adaptive Physics-Lifted Koopman State Space / experiment.py
Beats tuned baseline
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()