% Driver for the iterative-solver chapter.
%
% The editable script supplies
%
%     names = method_names()
%     H = run_iter(mid, N, U0, UE, nit, tol)
%
% where H has, in its three rows and one column per iteration, the squared
% relative residual, the squared maximum error, and the elapsed milliseconds.
% The squares are not an oversight: taking the square roots in the MATLAB layer
% means either a sqrt inside the iteration loop or a call after it, and either
% one makes numbl compile that loop through its slow path. The page roots them.
%
% and the driver supplies the starting array U0, which holds the exact
% boundary data on the boundary and zeros inside, and the exact solution UE for
% measurement only. For reference it also solves the same system exactly with
% the sine transform of Chapter 6, which gives both the discretisation error,
% the accuracy no iterative method can beat, and the time a direct solve takes.
p = lp_params();
kind = p.kind;
par = p.par;
N = p.N;
names = method_names();
nm = numel(names);

x = (0:N)' / N;
[X, Y] = meshgrid(x, x);
UE = lp_exact2d(kind, X, Y, par);

U0 = zeros(N+1, N+1);
U0(1, :) = UE(1, :);
U0(N+1, :) = UE(N+1, :);
U0(:, 1) = UE(:, 1);
U0(:, N+1) = UE(:, N+1);

% the exact solution of the discrete system, by the fast Poisson solver
t0 = tic;
Ud = drv_fast(U0, N);
ms_direct = toc(t0) * 1000;
disc_err = max(max(abs(Ud - UE)));

nit = p.nit;
tol = p.tol;   % hoisted: a struct field read inside a loop puts numbl's
               % compiler into its slow path, and the loop below is the timed one
R = NaN(nm, nit+1);
E = NaN(nm, nit+1);
T = NaN(nm, nit+1);
for k = 1:nm
    % the copy goes through a function call: a slice assignment written here in
    % the loop body would put the loop, and everything it calls, on numbl's
    % slow path
    H = run_iter(k, N, U0, UE, nit, tol);
    R = drv_store(R, H, 1, k);
    E = drv_store(E, H, 2, k);
    T = drv_store(T, H, 3, k);
end

out = struct();
out.names = names;
out.N = N;
out.its = 0:nit;
out.R2 = lp_rows(R);
out.E2 = lp_rows(E);
out.T = lp_rows(T);
out.discErr = disc_err;
out.directMs = ms_direct;
out.unknowns = (N-1)^2;
lp_write(out);

function A = drv_store(A, H, row, k)
for c = 1:size(A, 2)
    A(k, c) = H(row, c);
end
end

function U = drv_fast(U, N)
n = N - 1;
B = zeros(n, n);
B(1, :) = B(1, :) + U(1, 2:N);
B(n, :) = B(n, :) + U(N+1, 2:N);
B(:, 1) = B(:, 1) + U(2:N, 1);
B(:, n) = B(:, n) + U(2:N, N+1);
lam = 2 * (1 - cos(pi * (1:n)' / N));
D = lam * ones(1, n) + ones(n, 1) * lam';
Bh = lp_dst1(lp_dst1(B).').';
U(2:N, 2:N) = lp_dst1(lp_dst1(Bh ./ D).').';
end