% Five ways to solve the five-point system without factorising anything.
%
% Every method works only through the action of
%
%     (A u)(i,j) = 4 u(i,j) - u(i-1,j) - u(i+1,j) - u(i,j-1) - u(i,j+1)
%
% on the interior points, with the boundary values held fixed in the array. No
% matrix is formed, so the storage is the grid and nothing else. What separates
% the methods is how fast the error goes away: Jacobi and Gauss-Seidel need
% O(N^2) sweeps, optimal SOR and conjugate gradients O(N), and multigrid a
% number of cycles that does not grow with N at all.
%
% UE is passed in only so that the true error can be recorded next to the
% residual. No method uses it.
%
% A note on the shape of this file. Each method has its own loop rather than
% one loop with a switch inside it, and multigrid keeps every level in one flat
% vector rather than a cell array of grids. Both keep the *timed* loop
% type-stable for numbl's compiler: an iterate assigned from several different
% functions, or a body that indexes a cell, is not something it can specialise,
% and here that loop is the measurement.

function names = method_names()
names = {'Jacobi', 'Gauss-Seidel', 'SOR', 'conjugate gradients', 'multigrid V-cycles'};
end

function H = run_iter(mid, N, U, UE, nit, tol)
if mid == 1
    H = loop_jacobi(N, U, UE, nit, tol);
elseif mid == 2
    H = loop_sor(N, U, UE, nit, tol, 1);
elseif mid == 3
    H = loop_sor(N, U, UE, nit, tol, 2 / (1 + sin(pi/N)));
elseif mid == 4
    H = loop_cg(N, U, UE, nit, tol);
else
    H = loop_mg(N, U, UE, nit, tol);
end
end

% ------------------------------------------------------------ the methods ---

function H = loop_jacobi(N, U, UE, nit, tol)
% Every point replaced by the average of its four old neighbours. Error
% components with wavelength comparable to the domain are damped by a factor
% 1 - O(h^2) per sweep, so the sweep count grows like N^2.
res = NaN(1, nit+1);
err = NaN(1, nit+1);
tms = NaN(1, nit+1);
nf = 1;
ss0 = res_sumsq(U, N);
res(1) = ss0;
err(1) = maxsq(U - UE);
tms(1) = 0;
stop = tol * tol * ss0;
tacc = 0;
for it = 1:nit
    t1 = tic;
    U = jacobi_sweep(U, N);
    tacc = tacc + toc(t1);
    res(it+1) = res_sumsq(U, N);
    err(it+1) = maxsq(U - UE);
    tms(it+1) = tacc * 1000;
    nf = it + 1;
    if res(it+1) < stop
        break;
    end
end
H = hist_pack(res, err, tms, ss0, nf);
end

function H = loop_sor(N, U, UE, nit, tol, omega)
% omega = 1 is Gauss-Seidel: new values are used as soon as they are available,
% which costs nothing and halves the sweep count. Over-relaxing with the
% optimal omega = 2/(1 + sin(pi h)) changes the rate from 1 - O(h^2) to
% 1 - O(h). That optimal value is known only because the eigenvalues of this
% particular matrix are known.
res = NaN(1, nit+1);
err = NaN(1, nit+1);
tms = NaN(1, nit+1);
nf = 1;
ss0 = res_sumsq(U, N);
res(1) = ss0;
err(1) = maxsq(U - UE);
tms(1) = 0;
stop = tol * tol * ss0;
tacc = 0;
for it = 1:nit
    t1 = tic;
    U = sor_sweep(U, N, omega);
    tacc = tacc + toc(t1);
    res(it+1) = res_sumsq(U, N);
    err(it+1) = maxsq(U - UE);
    tms(it+1) = tacc * 1000;
    nf = it + 1;
    if res(it+1) < stop
        break;
    end
end
H = hist_pack(res, err, tms, ss0, nf);
end

function H = loop_cg(N, U, UE, nit, tol)
% Conjugate gradients. A is symmetric positive definite, so CG applies; it
% needs one operator application per iteration and four vectors of storage.
% Its convergence is governed by sqrt(kappa) rather than kappa, which for this
% matrix means O(N) iterations, like optimal SOR, but without having to know
% anything about the spectrum in advance.
res = NaN(1, nit+1);
err = NaN(1, nit+1);
tms = NaN(1, nit+1);
nf = 1;
R = neg_residual(U, N);      % the right-hand side is zero at interior points
P = R;
rs = dot2(R, R);
ss0 = rs;
res(1) = ss0;
err(1) = maxsq(U - UE);
tms(1) = 0;
stop = tol * tol * ss0;
tacc = 0;
for it = 1:nit
    t1 = tic;
    AP = applyA(P, N);
    al = rs / dot2(P, AP);
    U = U + al * P;
    R = R - al * AP;
    rsn = dot2(R, R);
    P = R + (rsn / rs) * P;
    rs = rsn;
    tacc = tacc + toc(t1);
    res(it+1) = rsn;
    err(it+1) = maxsq(U - UE);
    tms(it+1) = tacc * 1000;
    nf = it + 1;
    if res(it+1) < stop
        break;
    end
end
H = hist_pack(res, err, tms, ss0, nf);
end

function H = loop_mg(N, U, UE, nit, tol)
% Geometric multigrid. Gauss-Seidel is a poor solver but an excellent smoother:
% two sweeps remove the oscillatory part of the error and leave something
% smooth, which is well represented on a grid twice as coarse and much cheaper
% to correct there. Recursing on that idea gives a V-cycle whose cost is
% proportional to the number of unknowns and whose error reduction per cycle is
% a constant independent of N.
nlev = 1;
while mod(N, 2^nlev) == 0 && N / 2^nlev >= 2
    nlev = nlev + 1;
end
ns = zeros(nlev, 1);
offs = zeros(nlev, 1);
tot = 0;
nn = N;
for l = 1:nlev
    ns(l) = nn;
    offs(l) = tot;
    tot = tot + (nn+1)^2;
    nn = nn / 2;
end

% one flat vector holds every level: the iterate in the first half, the
% right-hand side in the second
W = zeros(2*tot, 1);
W = flatten_into(W, U, 0, N);
WE = flatten_into(zeros((N+1)^2, 1), UE, 0, N);

res = NaN(1, nit+1);
err = NaN(1, nit+1);
tms = NaN(1, nit+1);
nf = 1;
ss0 = res_sumsq_flat(W, 0, tot, N);
res(1) = ss0;
err(1) = maxsq_flat(W, WE, N);
tms(1) = 0;
stop = tol * tol * ss0;
tacc = 0;
for it = 1:nit
    t1 = tic;
    W = vcycle(W, offs, ns, nlev, tot, 2, 2);
    tacc = tacc + toc(t1);
    res(it+1) = res_sumsq_flat(W, 0, tot, N);
    err(it+1) = maxsq_flat(W, WE, N);
    tms(it+1) = tacc * 1000;
    nf = it + 1;
    if res(it+1) < stop
        break;
    end
end
H = hist_pack(res, err, tms, ss0, nf);
end

% ---------------------------------------------------------------- kernels ---

function U = jacobi_sweep(U, N)
V = U;
for j = 2:N
    for i = 2:N
        U(i,j) = 0.25 * (V(i-1,j) + V(i+1,j) + V(i,j-1) + V(i,j+1));
    end
end
end

function U = sor_sweep(U, N, omega)
for j = 2:N
    for i = 2:N
        U(i,j) = (1-omega)*U(i,j) + 0.25*omega*(U(i-1,j) + U(i+1,j) + U(i,j-1) + U(i,j+1));
    end
end
end

function Y = applyA(U, N)
Y = zeros(N+1, N+1);
for j = 2:N
    for i = 2:N
        Y(i,j) = 4*U(i,j) - U(i-1,j) - U(i+1,j) - U(i,j-1) - U(i,j+1);
    end
end
end

function Y = neg_residual(U, N)
Y = zeros(N+1, N+1);
for j = 2:N
    for i = 2:N
        Y(i,j) = -(4*U(i,j) - U(i-1,j) - U(i+1,j) - U(i,j-1) - U(i,j+1));
    end
end
end

function s = dot2(A, B)
s = 0;
n = numel(A);
for k = 1:n
    s = s + A(k) * B(k);
end
end

function ss = res_sumsq(U, N)
% The sum of squares, not the norm: a square root taken here would widen the
% inferred type of the accumulator and slow this loop by a factor of 20.
ss = 0;
for j = 2:N
    for i = 2:N
        r = 4*U(i,j) - U(i-1,j) - U(i+1,j) - U(i,j-1) - U(i,j+1);
        ss = ss + r*r;
    end
end
end

function H = hist_pack(res, err, tms, ss0, nf)
% Both histories are accumulated as squares and rooted here, over the first nf
% columns — the part the iteration actually filled; the rest stays NaN, which
% is what the page draws as a gap. Rooting inside the loop instead costs about
% 60 per cent on the sweep itself, because sqrt may return a complex value and
% the history array is then no longer a plain real array to numbl.
r = res;
e = err;
r(1:nf) = sqrt(res(1:nf) / ss0);
e(1:nf) = sqrt(err(1:nf));
H = [r; e; tms];
end

function m = maxsq(E)
m = max(max(E .* E));
end


% -------------------------------------------------------------- multigrid ---

function W = flatten_into(W, U, off, N)
m = N + 1;
for j = 1:m
    b = off + (j-1)*m;
    for i = 1:m
        W(b+i) = U(i,j);
    end
end
end

function ss = res_sumsq_flat(W, off, tot, N)
m = N + 1;
ss = 0;
for j = 2:N
    b = off + (j-1)*m;
    for i = 2:N
        r = W(tot + b + i) - (4*W(b+i) - W(b+i-1) - W(b+i+1) - W(b+i-m) - W(b+i+m));
        ss = ss + r*r;
    end
end
end

function m = maxsq_flat(W, WE, N)
mm = N + 1;
m = 0;
for k = 1:mm*mm
    d = W(k) - WE(k);
    if d*d > m
        m = d*d;
    end
end
end

function W = vcycle(W, offs, ns, nlev, tot, nu1, nu2)
for l = 1:nlev-1
    W = mg_down(W, offs(l), offs(l+1), tot, ns(l), nu1);
end
W = gs_flat(W, offs(nlev), tot, ns(nlev), 8);   % the coarsest grid: one unknown
for l = nlev-1:-1:1
    W = mg_up(W, offs(l), offs(l+1), tot, ns(l), nu2);
end
end

function W = mg_down(W, off, offc, tot, n, nu)
% Smooth, then carry the residual to the next grid and start the coarse
% iterate at zero. The residual is formed inline, one 9-point full-weighting
% stencil per coarse point, rather than materialised on the fine grid first.
W = gs_flat(W, off, tot, n, nu);
m = n + 1;
nc = n / 2;
mc = nc + 1;
for jc = 2:nc
    j = 2*jc - 1;
    for ic = 2:nc
        i = 2*ic - 1;
        k = off + (j-1)*m + i;
        r00 = W(tot+k) + W(k+1) + W(k-1) + W(k+m) + W(k-m) - 4*W(k);
        rm0 = W(tot+k-1) + W(k) + W(k-2) + W(k-1+m) + W(k-1-m) - 4*W(k-1);
        rp0 = W(tot+k+1) + W(k+2) + W(k) + W(k+1+m) + W(k+1-m) - 4*W(k+1);
        r0m = W(tot+k-m) + W(k-m+1) + W(k-m-1) + W(k) + W(k-2*m) - 4*W(k-m);
        r0p = W(tot+k+m) + W(k+m+1) + W(k+m-1) + W(k+2*m) + W(k) - 4*W(k+m);
        rmm = W(tot+k-m-1) + W(k-m) + W(k-m-2) + W(k-1) + W(k-2*m-1) - 4*W(k-m-1);
        rpm = W(tot+k-m+1) + W(k-m+2) + W(k-m) + W(k+1) + W(k-2*m+1) - 4*W(k-m+1);
        rmp = W(tot+k+m-1) + W(k+m) + W(k+m-2) + W(k+2*m-1) + W(k-1) - 4*W(k+m-1);
        rpp = W(tot+k+m+1) + W(k+m+2) + W(k+m) + W(k+2*m+1) + W(k+1) - 4*W(k+m+1);
        s = 4*r00 + 2*(rm0 + rp0 + r0m + r0p) + (rmm + rpm + rmp + rpp);
        % full weighting is s/16; the extra factor of 4 is because the coarse
        % operator is written in the same scaled form and so stands for (2h)^2
        % times the Laplacian. Leave it out and the correction comes back four
        % times too small, which shows up as a V-cycle that reduces the
        % residual by a factor of two instead of twenty.
        W(tot + offc + (jc-1)*mc + ic) = s / 4;
    end
end
for k = 1:mc*mc
    W(offc + k) = 0;
end
end

function W = mg_up(W, off, offc, tot, n, nu)
% Interpolate the coarse correction, add it, smooth again. The correction is
% built in a scratch array first: interpolating in place would mix it with the
% fine iterate that is already there.
m = n + 1;
nc = n / 2;
mc = nc + 1;
Z = zeros(m, m);
for jc = 1:mc
    for ic = 1:mc
        Z(2*ic-1, 2*jc-1) = W(offc + (jc-1)*mc + ic);
    end
end
for j = 1:2:m
    for i = 2:2:n
        Z(i,j) = 0.5 * (Z(i-1,j) + Z(i+1,j));
    end
end
for j = 2:2:n
    for i = 1:m
        Z(i,j) = 0.5 * (Z(i,j-1) + Z(i,j+1));
    end
end
for j = 1:m
    b = off + (j-1)*m;
    for i = 1:m
        W(b+i) = W(b+i) + Z(i,j);
    end
end
W = gs_flat(W, off, tot, n, nu);
end

function W = gs_flat(W, off, tot, n, nu)
% Gauss-Seidel on one level of the flat hierarchy: nu sweeps over the interior
% of the level at `off`, reading its right-hand side from the second half of W.
m = n + 1;
for s = 1:nu
    for j = 2:n
        b = off + (j-1)*m;
        for i = 2:n
            W(b+i) = 0.25 * (W(b+i-1) + W(b+i+1) + W(b+i-m) + W(b+i+m) + W(tot+b+i));
        end
    end
end
end