1import json
  2import numpy as np
  3from scipy.linalg import eigh
  4
  5
  6def make_gram(n=120, seed=7):
  7    rng = np.random.default_rng(seed)
  8    rows = []
  9    for i in range(n - 1):
 10        row = np.zeros(n)
 11        row[i], row[i + 1] = 1.0, -1.0
 12        rows.append(row)
 13    for i in range(n):
 14        row = np.zeros(n)
 15        row[i] = 0.08
 16        rows.append(row)
 17    G = np.asarray(rows)
 18    b = G.T @ rng.normal(size=G.shape[0])
 19    return G, b
 20
 21
 22def row_closures(G, blocks, tol=1e-12):
 23    supports = [set(np.flatnonzero(np.abs(row) > tol)) for row in G]
 24    closures, touching = [], []
 25    for block in blocks:
 26        js = [j for j, s in enumerate(supports) if s.intersection(block)]
 27        om = sorted(set().union(*(supports[j] for j in js)))
 28        touching.append(np.asarray(js, dtype=int))
 29        closures.append(np.asarray(om, dtype=int))
 30    return closures, touching
 31
 32
 33def build_preconditioners(G, lam=1e-4, block_size=10, keep=1):
 34    m, n = G.shape
 35    A0 = G.T @ G
 36    A = A0 + lam * np.eye(n)
 37    blocks = [np.arange(i, min(i + block_size, n)) for i in range(0, n, block_size)]
 38    oms, touching = row_closures(G, blocks)
 39    mult = np.zeros(m, dtype=int)
 40    for js in touching:
 41        mult[js] += 1
 42    # Exact Gram splitting, used here as a first numerical verification.
 43    split = np.zeros((n, n))
 44    local_data = []
 45    for om, js in zip(oms, touching):
 46        H = G[np.ix_(js, om)] / np.sqrt(mult[js])[:, None]
 47        local0 = H.T @ H
 48        local = local0 + lam * np.eye(len(om))
 49        split[np.ix_(om, om)] += local0
 50        local_data.append((om, local, local0))
 51    # P contains low generalized-energy modes on each aggregate, with the
 52    # row-closure Gram Schur complement providing the local energy metric.
 53    cols = []
 54    for block, (om, local, local0) in zip(blocks, local_data):
 55        pos = {v: k for k, v in enumerate(om)}
 56        ii = np.asarray([pos[v] for v in block])
 57        jj = np.asarray([k for k, v in enumerate(om) if v not in set(block)])
 58        Lbb = local0[np.ix_(ii, ii)]
 59        if len(jj):
 60            Lbg = local0[np.ix_(ii, jj)]
 61            Lgg = local0[np.ix_(jj, jj)]
 62            schur = local0[np.ix_(ii, ii)] - Lbg @ np.linalg.pinv(Lgg) @ Lbg.T
 63        else:
 64            schur = Lbb
 65        # D is the damped local diagonal, as in the generalized local problem.
 66        Dfull = np.diag(np.diag(local))
 67        D = Dfull[np.ix_(ii, ii)]
 68        vals, vecs = eigh(schur + lam * np.eye(len(ii)), D)
 69        take = np.argsort(vals)[:min(keep, len(ii))]
 70        for q in take:
 71            v = np.zeros(n)
 72            v[block] = vecs[:, q]
 73            v /= np.linalg.norm(v)
 74            cols.append(v)
 75    P = np.column_stack(cols) if cols else np.zeros((n, 0))
 76    # Columns have disjoint aggregate support, hence are already independent;
 77    # QR makes the construction robust to any future duplicate modes.
 78    P, _ = np.linalg.qr(P, mode='reduced')
 79    Ac = P.T @ A @ P
 80
 81    def block_jacobi(v):
 82        z = np.zeros(n)
 83        counts = np.zeros(n)
 84        for om, local, _ in local_data:
 85            z[om] += np.linalg.solve(local, v[om])
 86            counts[om] += 1.0
 87        return z / counts
 88
 89    def two_level(v):
 90        # Additive Schwarz plus Galerkin coarse correction. The coarse solve
 91        # uses A_c=(GP)^T(GP)+lambda P^T P exactly, without Hessian assembly.
 92        z = block_jacobi(v)
 93        if P.shape[1]:
 94            z += P @ np.linalg.solve(Ac, P.T @ v)
 95        return z
 96
 97    return A, split, P, Ac, block_jacobi, two_level, mult
 98
 99
100def pcg(A, b, M=None, tol=1e-9, maxit=500):
101    x = np.zeros_like(b)
102    r = b - A @ x
103    z = M(r) if M else r.copy()
104    p = z.copy()
105    rz = float(r @ z)
106    history = [np.linalg.norm(r)]
107    for k in range(1, maxit + 1):
108        Ap = A @ p
109        alpha = rz / float(p @ Ap)
110        x += alpha * p
111        r -= alpha * Ap
112        history.append(np.linalg.norm(r))
113        if history[-1] <= tol * history[0]:
114            return x, k, history
115        z = M(r) if M else r.copy()
116        rz_new = float(r @ z)
117        p = z + (rz_new / rz) * p
118        rz = rz_new
119    return x, maxit, history
120
121
122def main():
123    G, b = make_gram()
124    A, split, P, Ac, bj, tl, mult = build_preconditioners(G)
125    n = A.shape[0]
126    gram_error = np.linalg.norm(split - G.T @ G) / np.linalg.norm(G.T @ G)
127    coarse_error = np.linalg.norm(Ac - ((G @ P).T @ (G @ P) + 1e-4 * (P.T @ P))) / max(1, np.linalg.norm(Ac))
128    results = {}
129    for name, M in [('unpreconditioned', None), ('block_jacobi', bj), ('two_level', tl)]:
130        _, it, hist = pcg(A, b, M)
131        results[name] = {'iterations': int(it), 'final_relative_residual': float(hist[-1] / hist[0]), 'history': hist}
132    # A small damping sensitivity check at fixed construction settings.
133    damping = {}
134    for lam in [1e-6, 1e-4, 1e-2]:
135        A2, _, _, _, bj2, tl2, _ = build_preconditioners(G, lam=lam)
136        damping[str(lam)] = {}
137        for name, M in [('block_jacobi', bj2), ('two_level', tl2)]:
138            _, it, hist = pcg(A2, b, M)
139            damping[str(lam)][name] = int(it)
140    out = {'n': n, 'm': G.shape[0], 'coarse_dimension': int(P.shape[1]),
141           'max_row_touch_multiplicity': int(mult.max()),
142           'relative_gram_split_error': float(gram_error),
143           'relative_coarse_gram_error': float(coarse_error),
144           'results': results, 'damping_iterations': damping}
145    print(json.dumps(out, indent=2))
146
147
148if __name__ == '__main__':
149    main()