Fejer reflection accelerator for fixed-point layers / fejer_solver.py

Failed on benchmark

Raw ⬇ ZIP
 1"""Fejer averaged-reflection solver for callable fixed-point maps J.
 2
 3J is expected to be a resolvent (or a numerically stable approximation to one).
 4The core macro-step uses K evaluations of J and averages y, R y, ..., R^K y,
 5where R=2J-I. The optional safeguard costs one additional J evaluation.
 6"""
 7import numpy as np
 8
 9
10def residual(J, y):
11    """Euclidean fixed-point residual ||J(y)-y||."""
12    y = np.asarray(y, dtype=float)
13    return float(np.linalg.norm(np.asarray(J(y), dtype=float) - y))
14
15
16def fejer_step(J, y, K):
17    """One K-reflection Fejer macro-step; exactly K calls to J."""
18    if K < 0:
19        raise ValueError("K must be nonnegative")
20    y = np.asarray(y, dtype=float)
21    z = y.copy()
22    acc = y.copy()
23    for _ in range(K):
24        z = 2.0 * np.asarray(J(z), dtype=float) - z
25        acc += z
26    return acc / (K + 1)
27
28
29def safeguarded_step(J, y, K, eta=1.0):
30    """Accept Fejer candidate if normalized residual does not increase.
31
32    Returns (new_y, accepted, resulting_rho, old_rho). Rejection falls back
33    to one ordinary J update, matching the proposed safeguard.
34    """
35    y = np.asarray(y, dtype=float)
36    old_rho = residual(J, y) / max(1.0, float(np.linalg.norm(y)))
37    candidate = fejer_step(J, y, K)
38    eta = float(np.clip(eta, 0.0, 1.0))
39    trial = (1.0 - eta) * y + eta * candidate
40    trial_rho = residual(J, trial) / max(1.0, float(np.linalg.norm(trial)))
41    if trial_rho <= old_rho:
42        return trial, True, trial_rho, old_rho
43    fallback = np.asarray(J(y), dtype=float)
44    fallback_rho = residual(J, fallback) / max(1.0, float(np.linalg.norm(fallback)))
45    return fallback, False, fallback_rho, old_rho
46
47
48def solve(J, y0, macro_steps, K, safeguard=True):
49    """Run macro_steps and return iterate plus history of normalized residuals."""
50    y = np.asarray(y0, dtype=float).copy()
51    history = []
52    accepted = []
53    for _ in range(macro_steps):
54        if safeguard:
55            y, ok, rho, _ = safeguarded_step(J, y, K)
56            accepted.append(ok)
57        else:
58            y = fejer_step(J, y, K)
59            rho = residual(J, y) / max(1.0, float(np.linalg.norm(y)))
60            accepted.append(True)
61        history.append(rho)
62    return y, np.asarray(history), accepted