% Four ways to end a grid, for waves that ought to leave it.
%
% The region of interest is the square of (N+1)-by-(N+1) nodes that u0
% covers.  Around it, a method may add a layer of W extra nodes on every
% side, which costs work but is never looked at.  The interior is Chapter 4's
% leapfrog scheme; the methods differ only at the edge:
%
%   'rigid'       a hard wall, du/dn = 0, by mirroring the first interior node
%                 into a ghost.  Everything comes back, with the same sign.
%   'mur'         Mur's (1981) first-order absorbing condition, a discrete
%                 u_t + c u_n = 0 on each edge.  It is exact for a wave
%                 arriving square-on and reflects a fraction
%                 (1 - cos a) / (1 + cos a) of one arriving at angle a.
%   'sponge'      a damping layer: the equation u_tt + 2 s(x) u_t = c^2 Lap u,
%                 with s rising from zero to smax across W nodes as the square
%                 of the depth, and a zero-displacement edge behind it.  A wave
%                 is damped on the way in and again on the way back out, and
%                 part of it is reflected by the change of s itself.
%   'pml'         a perfectly matched layer (Berenger 1994).  The scheme is
%                 rewritten in pressure and velocity on a staggered grid,
%                 which in the region is the same leapfrog scheme exactly;
%                 in the layer the pressure is split into the parts driven by
%                 the x and the y velocity, and each is damped only in its own
%                 direction.  In the continuum this absorbs a wave at any
%                 angle and any frequency without reflecting any of it; on
%                 the grid, the profile's steps reflect a little.
%
% u0 and uold are the region's field at steps 0 and -1.  The result U holds
% the region's field at the nodes `keep` (linear indices into the region)
% after each of the M steps, one column per step.

function U = openbox(u0, uold, lam, M, method, W, keep)
n0 = size(u0, 1);
if strcmp(method, 'rigid') || strcmp(method, 'mur')
    W = 0;
end
if strcmp(method, 'pml')
    U = drv_pml(u0, uold, lam, M, W, keep);
    return;
end
n1 = n0 + 2 * W;
u = zeros(n1, n1);
uo = zeros(n1, n1);
u(W + 1:W + n0, W + 1:W + n0) = u0;
uo(W + 1:W + n0, W + 1:W + n0) = uold;
un = zeros(n1, n1);
l2 = lam ^ 2;
% the damping, per step: s dt at each node, zero in the region itself
sd = zeros(n1, n1);
if W > 0
    % s dt at the outer edge.  A stronger sponge damps more but reflects more
    % from its own gradient; this strength did best of the few tried.
    smax = 3 * lam / (2 * W) * log(100);
    for j = 1:n1
        for i = 1:n1
            di = max([W + 1 - i, i - (W + n0), 0]);
            dj = max([W + 1 - j, j - (W + n0), 0]);
            d = max(di, dj) / W;
            sd(i, j) = smax * d ^ 2;
        end
    end
end
rigid = strcmp(method, 'rigid');
mur = strcmp(method, 'mur');
km = (lam - 1) / (lam + 1);
U = zeros(numel(keep), M);
for n = 1:M
    for j = 2:n1 - 1
        for i = 2:n1 - 1
            a = sd(i, j);
            un(i, j) = (2 * u(i, j) - (1 - a) * uo(i, j) ...
                + l2 * (u(i + 1, j) + u(i - 1, j) + u(i, j + 1) + u(i, j - 1) - 4 * u(i, j))) / (1 + a);
        end
    end
    if rigid
        % the edge node sees its inner neighbour twice
        for k = 2:n1 - 1
            un(1, k) = 2 * u(1, k) - uo(1, k) + l2 * (2 * u(2, k) + u(1, k + 1) + u(1, k - 1) - 4 * u(1, k));
            un(n1, k) = 2 * u(n1, k) - uo(n1, k) + l2 * (2 * u(n1 - 1, k) + u(n1, k + 1) + u(n1, k - 1) - 4 * u(n1, k));
            un(k, 1) = 2 * u(k, 1) - uo(k, 1) + l2 * (2 * u(k, 2) + u(k + 1, 1) + u(k - 1, 1) - 4 * u(k, 1));
            un(k, n1) = 2 * u(k, n1) - uo(k, n1) + l2 * (2 * u(k, n1 - 1) + u(k + 1, n1) + u(k - 1, n1) - 4 * u(k, n1));
        end
        un(1, 1) = 2 * u(1, 1) - uo(1, 1) + l2 * (2 * u(2, 1) + 2 * u(1, 2) - 4 * u(1, 1));
        un(1, n1) = 2 * u(1, n1) - uo(1, n1) + l2 * (2 * u(2, n1) + 2 * u(1, n1 - 1) - 4 * u(1, n1));
        un(n1, 1) = 2 * u(n1, 1) - uo(n1, 1) + l2 * (2 * u(n1 - 1, 1) + 2 * u(n1, 2) - 4 * u(n1, 1));
        un(n1, n1) = 2 * u(n1, n1) - uo(n1, n1) + l2 * (2 * u(n1 - 1, n1) + 2 * u(n1, n1 - 1) - 4 * u(n1, n1));
    elseif mur
        for k = 2:n1 - 1
            un(1, k) = u(2, k) + km * (un(2, k) - u(1, k));
            un(n1, k) = u(n1 - 1, k) + km * (un(n1 - 1, k) - u(n1, k));
            un(k, 1) = u(k, 2) + km * (un(k, 2) - u(k, 1));
            un(k, n1) = u(k, n1 - 1) + km * (un(k, n1 - 1) - u(k, n1));
        end
        un(1, 1) = 0.5 * (un(2, 1) + un(1, 2));
        un(1, n1) = 0.5 * (un(2, n1) + un(1, n1 - 1));
        un(n1, 1) = 0.5 * (un(n1 - 1, 1) + un(n1, 2));
        un(n1, n1) = 0.5 * (un(n1 - 1, n1) + un(n1, n1 - 1));
    end
    uo = u;
    u = un;
    R = u(W + 1:W + n0, W + 1:W + n0);
    U(:, n) = R(keep);
end
end

function U = drv_pml(u0, uold, lam, M, W, keep)
% Leapfrog in pressure p and scaled velocity w = rho c v, staggered: w_x sits
% between p(i, j) and p(i, j+1), w_y between p(i, j) and p(i+1, j).
%
%   w <- w - lam grad p,    p <- p - lam div w
%
% Eliminating w gives Chapter 4's scheme for p, so in the region this is the
% same computation.  In the layer, p = px + py and every field is damped by
% the profile of its own direction, s(d) = smax (d/W)^2 at depth d, with
% smax chosen for a nominal reflection of 1e-5 at normal incidence.
n0 = size(u0, 1);
n1 = n0 + 2 * W;
smax = 3 * lam / (2 * W) * log(1e5);
% damping per step at whole and half nodes, along one axis
dw = zeros(1, n1);
dh = zeros(1, n1 - 1);
for k = 1:n1
    d = max([W + 1 - k, k - (W + n0), 0]) / W;
    dw(k) = smax * d ^ 2;
end
for k = 1:n1 - 1
    kk = k + 0.5;
    d = max([W + 1 - kk, kk - (W + n0), 0]) / W;
    dh(k) = smax * d ^ 2;
end
p0 = zeros(n1, n1);
p0(W + 1:W + n0, W + 1:W + n0) = u0;
pm = zeros(n1, n1);
pm(W + 1:W + n0, W + 1:W + n0) = uold;
% start at rest: w at step -1/2 is lam/2 times the gradient of p at step 0,
% which makes p at step 1 equal to what leapfrog would give from p0 and pm
wx = zeros(n1, n1 - 1);
wy = zeros(n1 - 1, n1);
for j = 1:n1 - 1
    for i = 1:n1
        wx(i, j) = 0.5 * lam * (p0(i, j + 1) - p0(i, j));
    end
end
for j = 1:n1
    for i = 1:n1 - 1
        wy(i, j) = 0.5 * lam * (p0(i + 1, j) - p0(i, j));
    end
end
% the split: all of p0 starts in px, which is harmless because the split
% only matters inside the layer, where p0 is zero
px = p0;
py = zeros(n1, n1);
pp = p0;
% the damped update x <- (1 - a) x / (1 + a) - lam (...) / (1 + a), with its
% two coefficients worked out once per row or column
fw = (1 - dw / 2) ./ (1 + dw / 2);
gw = lam ./ (1 + dw / 2);
fh = (1 - dh / 2) ./ (1 + dh / 2);
gh = lam ./ (1 + dh / 2);
U = zeros(numel(keep), M);
for n = 1:M
    for j = 1:n1 - 1
        f = fh(j);
        g = gh(j);
        for i = 1:n1
            wx(i, j) = f * wx(i, j) - g * (pp(i, j + 1) - pp(i, j));
        end
    end
    for j = 1:n1
        for i = 1:n1 - 1
            wy(i, j) = fh(i) * wy(i, j) - gh(i) * (pp(i + 1, j) - pp(i, j));
        end
    end
    for j = 2:n1 - 1
        f = fw(j);
        g = gw(j);
        for i = 2:n1 - 1
            px(i, j) = f * px(i, j) - g * (wx(i, j) - wx(i, j - 1));
            py(i, j) = fw(i) * py(i, j) - gw(i) * (wy(i, j) - wy(i - 1, j));
            pp(i, j) = px(i, j) + py(i, j);
        end
    end
    R = pp(W + 1:W + n0, W + 1:W + n0);
    U(:, n) = R(keep);
end
end