Fejer reflection accelerator for fixed-point layers / fejer_experiment.py
Failed on benchmark
1import json, math, time
2import numpy as np
3
4
5def rot(theta):
6 return np.array([[math.cos(theta), -math.sin(theta)],
7 [math.sin(theta), math.cos(theta)]], dtype=float)
8
9
10def J_from_R(R):
11 return (np.eye(R.shape[0]) + R) / 2.0
12
13
14def fejer(J, y0, K):
15 z = y0.copy()
16 a = y0.copy()
17 for _ in range(K):
18 z = 2.0 * (J @ z) - z
19 a += z
20 return a / (K + 1)
21
22
23def ppm(J, y0, n):
24 y = y0.copy()
25 for _ in range(n):
26 y = J @ y
27 return y
28
29
30def residual(J, y):
31 return np.linalg.norm(J @ y - y)
32
33
34def run():
35 y0 = np.array([1.0, 0.0])
36 rows = []
37 # Prediction 1: worst-case over rotation angle equals 1/(K+1).
38 for K in [1, 3, 7, 15, 31]:
39 th = np.linspace(0, 2*np.pi, 200001)
40 # exact closed form for the residual of the averaged reflection
41 vals = np.abs(np.sin((K + 1) * th / 2.0)) / (K + 1)
42 observed = float(vals.max())
43 predicted = 1.0 / (K + 1)
44 rows.append({'test':'worst_case_bound','K':K,'observed':observed,
45 'predicted':predicted,'relative_error':abs(observed-predicted)/predicted})
46 # Prediction 2: transition at (K+1) theta/2 = pi/2, i.e. theta*=pi/(K+1).
47 # Sweep first crossing of 90% of worst case for K=15.
48 K = 15
49 angles = np.linspace(1e-7, np.pi, 200000)
50 vals = np.abs(np.sin((K + 1) * angles / 2.0)) / (K + 1)
51 target = .9 / (K + 1)
52 idx = np.where(vals >= target)[0][0]
53 observed_theta = float(angles[idx])
54 predicted_theta = math.asin(.9) * 2.0 / (K + 1)
55 rows.append({'test':'transition_90pct','K':K,'observed_theta':observed_theta,
56 'predicted_theta':predicted_theta,'observed_Ktheta':observed_theta*K,
57 'predicted_Ktheta':predicted_theta*K,
58 'relative_error':abs(observed_theta-predicted_theta)/predicted_theta})
59 # Prediction 3: low-frequency residual is approximately theta/2 and does not
60 # improve with K while K theta << 1; report ratios over a parameter sweep.
61 low_rows = []
62 for K in [3, 7, 15, 31]:
63 theta = 0.01 / (K + 1) # safely in the low-frequency regime
64 R = rot(theta); J = J_from_R(R)
65 got = residual(J, fejer(J, y0, K))
66 predicted = theta / 2.0
67 low_rows.append({'K':K,'theta':theta,'observed':got,'predicted_theta_over_2':predicted,
68 'ratio':got/predicted})
69 # Direct numerical identity check on several angles and K.
70 identity_err = 0.0
71 for theta in np.linspace(0.03, 5.9, 12):
72 R = rot(float(theta)); J = J_from_R(R)
73 for K in [0,1,2,5,11]:
74 yh = fejer(J, y0, K)
75 rhs = (np.linalg.matrix_power(R, K+1) @ y0 - y0) / (2*(K+1))
76 identity_err = max(identity_err, np.linalg.norm((J @ yh-yh)-rhs))
77 # Equal oracle evaluations: Fejer K uses K J calls; PPM uses K calls.
78 comparison = []
79 theta = 0.35
80 R = rot(theta); J = J_from_R(R)
81 for K in [3,7,15,31]:
82 yf = fejer(J,y0,K); yp = ppm(J,y0,K)
83 comparison.append({'evaluations':K,'fejer_residual':residual(J,yf),
84 'ppm_residual':residual(J,yp),
85 'ratio_fejer_over_ppm':residual(J,yf)/residual(J,yp)})
86 out = {'identity_max_abs_error':identity_err,'predictions':rows,
87 'low_frequency':low_rows,'equal_eval_rotation_theta':theta,
88 'comparison':comparison,
89 'notes':'R is a rotation, hence nonexpansive; J=(I+R)/2 is an exact resolvent of a linear maximal monotone skew operator when -1 is absent.'}
90 print(json.dumps(out, indent=2))
91
92if __name__ == '__main__':
93 run()