1import json
 2from pathlib import Path
 3import numpy as np
 4from experiment import phi
 5
 6
 7def make_problem(eps_scale, seed=2802):
 8    rng = np.random.default_rng(seed)
 9    blocks = [np.diag([1.0, 1.4]), np.diag([2.0, 2.4]),
10              np.diag([4.0, 4.5]), np.diag([7.0, 7.8])]
11    n = 8
12    H = np.zeros((n, n))
13    for i, B in enumerate(blocks): H[2*i:2*i+2, 2*i:2*i+2] = B
14    for i in range(3):
15        eps = eps_scale * (1.0 if i != 1 else 0.7)
16        u = np.array([1., .3]); u /= np.linalg.norm(u)
17        v = np.array([.4, 1.]); v /= np.linalg.norm(v)
18        E = eps * np.outer(v, u)
19        H[2*i+2:2*i+4, 2*i:2*i+2] = E
20        H[2*i:2*i+2, 2*i+2:2*i+4] = E.T
21    return H, rng.normal(size=n), blocks
22
23
24def groups_fixed(blocks):
25    return [list(range(2*i, 2*i+2)) for i in range(len(blocks))]
26
27
28def groups_adaptive(H, blocks, tau):
29    groups = groups_fixed(blocks)
30    while True:
31        for j in range(len(groups)-1):
32            a, b = groups[j], groups[j+1]
33            E = H[np.ix_(b, a)]
34            eps = np.linalg.norm(E, 2)
35            ea = np.linalg.eigvalsh(H[np.ix_(a, a)])
36            eb = np.linalg.eigvalsh(H[np.ix_(b, b)])
37            eta = np.min(np.abs(ea[:, None] - eb[None, :]))
38            cert = phi(eta, eps) * np.sqrt(2) * eps
39            if cert > tau:
40                groups[j] = a + b
41                groups.pop(j+1)
42                break
43        else:
44            return groups
45
46
47def preconditioner(H, groups):
48    P = np.zeros_like(H)
49    for g in groups:
50        P[np.ix_(g, g)] = np.linalg.inv(H[np.ix_(g, g)])
51    return P
52
53
54def run(H, b, groups, steps=8):
55    P = preconditioner(H, groups)
56    x = np.zeros_like(b)
57    optimum = np.linalg.solve(H, b)
58    opt_loss = .5*optimum@H@optimum - b@optimum
59    for _ in range(steps): x -= 0.9 * P @ (H @ x - b)
60    loss = .5*x@H@x - b@x
61    return float(max(loss - opt_loss, 0.0))
62
63
64def main():
65    out = []
66    for coupling in [0.03, 0.3, 0.8]:
67        H, b, blocks = make_problem(coupling)
68        fixed = groups_fixed(blocks)
69        adaptive = groups_adaptive(H, blocks, tau=0.08)
70        full = [list(range(8))]
71        out.append({'coupling': coupling, 'steps': 8,
72                    'fixed_groups': fixed, 'adaptive_groups': adaptive,
73                    'full_loss': run(H,b,full), 'fixed_loss': run(H,b,fixed),
74                    'adaptive_loss': run(H,b,adaptive), 'full_matrix_entries': 64,
75                    'fixed_preconditioner_entries': sum(len(g)**2 for g in fixed),
76                    'adaptive_preconditioner_entries': sum(len(g)**2 for g in adaptive)})
77    Path('optimizer_results.json').write_text(json.dumps(out, indent=2))
78    print(json.dumps(out, indent=2))
79
80if __name__ == '__main__': main()