"""Gold reference: Modified Craig-Sneyd (MCS) ADI solver for a 2D convection-free parabolic PDE *with a mixed derivative term*. u_t = a u_xx + 2c u_xy + b u_yy, (x, y) in D = (0,1)^2, t in [0, T], T = 1 with diffusion matrix M = [[a, c], [c, b]] symmetric positive definite (a > 0, b > 0, a*b - c^2 > 0). This is the prototype of the correlated-diffusion operators that appear in financial PDEs (e.g. Heston / correlated multi-asset Black-Scholes), where the cross term comes from the correlation rho. Why this is hard ---------------- The classic Peaceman-Rachford / Douglas ADI schemes split the operator into unidirectional implicit sweeps and CANNOT make the mixed term 2c u_xy implicit. Treated naively (dropped, or lagged in an unstable way) the method either converges to the wrong PDE (order 0) or blows up. Second order *with* a mixed derivative requires a modern ADI scheme that keeps the mixed term explicit but corrects it to second order -- Craig-Sneyd (CS), Modified Craig-Sneyd (MCS), or Hundsdorfer-Verwer (HV). This reference uses MCS with theta = 1/3, which is unconditionally stable in 2D with a mixed derivative (in 't Hout & Welfert). Manufactured exact solution (homogeneous problem) ------------------------------------------------- u_e(x, y, t) = sum_m A_m exp(-lambda_m t) sin(alpha_m x + beta_m y + phi_m), lambda_m = a alpha_m^2 + 2c alpha_m beta_m + b beta_m^2 (>= 0 by SPD). Each rotated wave solves the PDE exactly, so the sum does. Dirichlet boundary data and the initial condition are u_e. A single separable product sin(px)sin(qy) is NOT a solution (the cross term turns sin*sin into cos*cos), so the mixed derivative genuinely matters. Operator splitting (in 't Hout & Foulon framework) -------------------------------------------------- F = F0 + F1 + F2, F0 V = 2c * delta_xy V (mixed term; explicit only), F1 V = a * delta_xx V (x-implicit sweep), F2 V = b * delta_yy V (y-implicit sweep), each F_j(t, .) also carries the Dirichlet boundary contribution at time t. Required solver interface (see problem_spec.md) ----------------------------------------------- solve(h, T=1.0, problem=None) -> (X, Y, U) holding the full (m+1)x(m+1) grid with indexing='ij'. The grader imports this module in an isolated subprocess and calls solve() to regenerate the convergence study, so the submitted arrays can be checked against the submitted code. Dependencies: Python + NumPy + SciPy only. """ import numpy as np from scipy.linalg import solve_banded # ---- problem definition ----------------------------------------------------- A_COEF, B_COEF, C_COEF = 1.0, 0.8, 0.5 # SPD: a*b - c^2 = 0.55 > 0 THETA = 1.0 / 3.0 T_FINAL = 1.0 # (amplitude, alpha, beta, phase) for each rotated decaying wave MODES = [ (1.0, 0.7, 1.3, 0.0), (0.7, 1.4, -0.6, 0.3), (0.5, 1.0, 0.8, 0.7), ] assert A_COEF > 0 and B_COEF > 0 and A_COEF * B_COEF - C_COEF**2 > 0 def u_exact(x, y, t): """Manufactured exact solution / boundary data. Works on arrays.""" s = np.zeros_like(x * y * 1.0) if np.ndim(x * y) else 0.0 for (amp, al, be, ph) in MODES: lam = A_COEF * al * al + 2 * C_COEF * al * be + B_COEF * be * be s = s + amp * np.exp(-lam * t) * np.sin(al * x + be * y + ph) return s def benchmark_problem(): """The problem defined in problem_spec.md, as a data dict.""" return {"a": A_COEF, "b": B_COEF, "c": C_COEF, "u0": lambda X, Y: u_exact(X, Y, 0.0), "g": lambda X, Y, t: u_exact(X, Y, t)} def adi_solve(h, T=T_FINAL, problem=None): """Advance from t=0 to t=T with k=h using MCS ADI. `problem` supplies the diffusion coefficients and the data: a, b, c : floats (SPD) u0(X, Y) -> initial field on the grid g(X, Y, t) -> Dirichlet data (only its boundary values are used) None means the benchmark problem. Nothing here is specialised to the manufactured solution, so the same code solves any admissible instance. Returns (X, Y, U) on the full (m+1)x(m+1) grid (indexing 'ij'). """ if problem is None: problem = benchmark_problem() a_c, b_c, c_c = problem["a"], problem["b"], problem["c"] u0_f, g_f = problem["u0"], problem["g"] m = int(round(1.0 / h)) k = h N = int(round(T / k)) h2 = h * h x = np.linspace(0.0, 1.0, m + 1) y = np.linspace(0.0, 1.0, m + 1) X, Y = np.meshgrid(x, y, indexing="ij") def fill(Vint, t): """Full grid with interior = Vint and boundary = g(., ., t).""" G = np.array(g_f(X, Y, t), dtype=float, copy=True) G[1:m, 1:m] = Vint return G def Dxx(G): return (G[2:m + 1, 1:m] - 2 * G[1:m, 1:m] + G[0:m - 1, 1:m]) / h2 def Dyy(G): return (G[1:m, 2:m + 1] - 2 * G[1:m, 1:m] + G[1:m, 0:m - 1]) / h2 def Dxy(G): return (G[2:m + 1, 2:m + 1] - G[2:m + 1, 0:m - 1] - G[0:m - 1, 2:m + 1] + G[0:m - 1, 0:m - 1]) / (4 * h2) def F0(V, t): return 2 * c_c * Dxy(fill(V, t)) def F1(V, t): return a_c * Dxx(fill(V, t)) def F2(V, t): return b_c * Dyy(fill(V, t)) def F(V, t): G = fill(V, t) return 2 * c_c * Dxy(G) + a_c * Dxx(G) + b_c * Dyy(G) # implicit tridiagonal operators (I - theta k A_j) rx = THETA * k * a_c / h2 ry = THETA * k * b_c / h2 abx = np.zeros((3, m - 1)); abx[0, 1:] = -rx; abx[1, :] = 1 + 2 * rx; abx[2, :-1] = -rx aby = np.zeros((3, m - 1)); aby[0, 1:] = -ry; aby[1, :] = 1 + 2 * ry; aby[2, :-1] = -ry def solve_x(R, t_imp): gb = np.asarray(g_f(X, Y, t_imp), dtype=float) Rb = R.copy() Rb[0, :] += rx * gb[0, 1:m] Rb[-1, :] += rx * gb[m, 1:m] return solve_banded((1, 1), abx, Rb) def solve_y(R, t_imp): gb = np.asarray(g_f(X, Y, t_imp), dtype=float) Rb = R.copy() Rb[:, 0] += ry * gb[1:m, 0] Rb[:, -1] += ry * gb[1:m, m] return solve_banded((1, 1), aby, Rb.T).T U = np.asarray(u0_f(X, Y), dtype=float)[1:m, 1:m].copy() for n in range(N): tn, tn1 = n * k, (n + 1) * k F0n, F1n, F2n = F0(U, tn), F1(U, tn), F2(U, tn) Fn = F0n + F1n + F2n # predictor (explicit Euler) + two implicit unidirectional corrections Y0 = U + k * Fn Y1 = solve_x(Y0 - THETA * k * F1n, tn1) Y2 = solve_y(Y1 - THETA * k * F2n, tn1) # MCS mixed-term correction (this is what buys 2nd order with u_xy) Yhat0 = (Y0 + THETA * k * (F0(Y2, tn1) - F0n) + (0.5 - THETA) * k * (F(Y2, tn1) - Fn)) Yh1 = solve_x(Yhat0 - THETA * k * F1n, tn1) Yh2 = solve_y(Yh1 - THETA * k * F2n, tn1) U = Yh2 Ufull = np.array(g_f(X, Y, T), dtype=float, copy=True) Ufull[1:m, 1:m] = U return X, Y, Ufull def solve(h, T=T_FINAL, problem=None): """Required public entry point. See the module docstring. Parameters ---------- h : float mesh size; the grid is (m+1)x(m+1) with m = 1/h, and k = h. T : float final time. problem : dict or None None -> the benchmark problem. Otherwise a dict with keys 'a', 'b', 'c' (floats) and callables 'u0(X, Y)' and 'g(X, Y, t)'. The grader calls solve() on problems it chooses, so the solver must not be specialised to the manufactured solution. Returns ------- (X, Y, U) : three (m+1, m+1) float arrays, indexing='ij', U at time T. """ return adi_solve(h, T, problem) def max_norm_error(h, T=T_FINAL): X, Y, U = solve(h, T) return float(np.max(np.abs(U - u_exact(X, Y, T)))) if __name__ == "__main__": print(f"MCS ADI, theta={THETA:.4f}, M=[[{A_COEF},{C_COEF}],[{C_COEF},{B_COEF}]]") print(f"{'h':>10} {'e_h':>16} {'order':>9}") prev = None for p in range(5): h = 0.1 * 2.0**(-p) e = max_norm_error(h) o = np.log2(prev / e) if prev is not None else float("nan") print(f"{h:>10.5f} {e:>16.6e} {o:>9.4f}") prev = e