1import json
 2import numpy as np
 3
 4GAMMA = 0.25
 5A = np.array([[0, 0, 0, 0, 0], [0.5, 0, 0, 0, 0],
 6              [17/50, -1/25, 0, 0, 0],
 7              [371/1360, -137/2720, 15/544, 0, 0],
 8              [25/24, -49/48, 125/16, -85/12, 0]], float)
 9
10
11def residual(z, b, h, rates):
12    # r=b-Az for A=I+h*gamma*diag(rates), i.e. F(y)=-diag(rates)y.
13    return b - (z + h * GAMMA * rates * z)
14
15
16def chebyshev_solve(b, h, rates, K, inflate=1.10, supplied_beta=False):
17    """Chebyshev semi-iteration for the diagonal implicit stage system.
18
19    supplied_beta=True reproduces the idea_context literally.  The default uses
20    the dimensionally consistent momentum beta=d^2*alpha_k*alpha_{k-1}/4.
21    """
22    lo = 1.0
23    hi = (1.0 + h * GAMMA * float(np.max(rates))) * inflate
24    c, d = (lo + hi) / 2.0, (hi - lo) / 2.0
25    z = b.copy(); zprev = z.copy(); alpha_prev = None
26    rs = []
27    for k in range(K):
28        r = residual(z, b, h, rates); rs.append(float(np.linalg.norm(r)))
29        alpha = 1.0 / c if k == 0 else 1.0 / (c - d*d*alpha_prev/4.0)
30        beta = 0.0 if k == 0 else d*d*alpha/4.0
31        if not supplied_beta and k > 0:
32            beta *= alpha_prev
33        znew = z + alpha*r + beta*(z-zprev)
34        zprev, z, alpha_prev = z, znew, alpha
35    rs.append(float(np.linalg.norm(residual(z, b, h, rates))))
36    return z, np.asarray(rs)
37
38
39def fixed_point_solve(b, h, rates, K):
40    z = b.copy(); rs = []
41    for _ in range(K):
42        r = residual(z, b, h, rates); rs.append(float(np.linalg.norm(r))); z = z + r
43    rs.append(float(np.linalg.norm(residual(z, b, h, rates))))
44    return z, np.asarray(rs)
45
46
47def sdirk_step(y, h, rates, method="cheb", K=24):
48    fvals = []
49    for i in range(5):
50        b = y + h * sum(A[i, j] * fvals[j] for j in range(i)) if i else y.copy()
51        z, _ = (chebyshev_solve(b, h, rates, K) if method == "cheb"
52                else fixed_point_solve(b, h, rates, K))
53        fvals.append(-rates * z)
54    return z
55
56
57def rk4_step(y, h, rates):
58    f = lambda x: -rates*x
59    k1 = f(y); k2 = f(y+h*k1/2); k3 = f(y+h*k2/2); k4 = f(y+h*k3)
60    return y + h*(k1+2*k2+2*k3+k4)/6
61
62
63def main():
64    np.random.seed(7)
65    rates = np.geomspace(1.0, 1000.0, 16)
66    y0 = np.random.randn(rates.size)
67    h = 0.20; b = y0.copy()
68    _, literal = chebyshev_solve(b, h, rates, 20, supplied_beta=True)
69    _, corrected = chebyshev_solve(b, h, rates, 20)
70    _, fixed = fixed_point_solve(b, h, rates, 20)
71    report = {"rates": rates.tolist(), "checks": {
72        "h": h, "literal_beta_final_over_initial": float(literal[-1]/literal[0]),
73        "corrected_beta_final_over_initial": float(corrected[-1]/corrected[0]),
74        "fixed_final_over_initial": float(fixed[-1]/fixed[0]),
75        "corrected_residuals": corrected.tolist(),
76        "literal_residuals": literal.tolist(),
77        "fixed_residuals": fixed.tolist(),
78        "corrected_monotone_after_transients": bool(np.all(np.diff(corrected[2:]) <= 1e-12))}}
79    report["steps"] = []
80    for h in [0.001, 0.01, 0.05, 0.2]:
81        exact = y0 * np.exp(-rates*h)
82        yc = sdirk_step(y0, h, rates, "cheb")
83        yf = sdirk_step(y0, h, rates, "fixed")
84        yr = rk4_step(y0, h, rates)
85        report["steps"].append({"h": h, "cheb_error": float(np.linalg.norm(yc-exact)),
86            "fixed_error": float(np.linalg.norm(yf-exact)), "rk4_error": float(np.linalg.norm(yr-exact)),
87            "cheb_norm": float(np.linalg.norm(yc)), "fixed_norm": float(np.linalg.norm(yf)),
88            "rk4_norm": float(np.linalg.norm(yr)), "rk4_finite": bool(np.all(np.isfinite(yr)))})
89    print(json.dumps(report, indent=2))
90
91if __name__ == "__main__":
92    main()