% Driver for Chapter 6: the round drum.
%
% The editable script supplies
%
%     [C, E, W, Nn, S, in] = diskstencil(N, method)
%
% the five-point Laplacian (times h^2) on the unknowns of an N-by-N grid over
% [-1, 1]^2 that lie inside the unit disk, with some treatment of the rim.
%
% The exact modes of the unit disk are J_m(j_mk r) cos(m theta), with
% frequency j_mk / (2 pi) when c = 1, where j_mk is the k-th zero of the
% Bessel function J_m; wv_bjzeros computes them to rounding error.
%
% mode 'hear'  the exact round drum, struck at radius r0, heard at a fixed
%              point, by modal synthesis; and a square drum with the same
%              fundamental, also exact, for comparison.
% mode 'grid'  the pitch error of the lowest modes on grids of increasing
%              size, for each rim treatment.  The driver runs leapfrog itself,
%              at a Courant number of 0.95 times the largest the stencil
%              allows (found by power iteration on the stencil, and never
%              more than 1/sqrt(2)); the grid starts in the exact mode, and
%              the frequency is read off the projection of the field onto it.
p = wv_params();
out = struct();

if strcmp(p.mode, 'hear')
    fs = p.fs;
    f01 = p.f01;
    T = p.T;
    t = (0:round(T * fs) - 1)' / fs;
    c = 2 * pi * f01 / 2.404825557695773;    % so that the lowest mode is f01, R = 1
    sw = p.width;
    y = zeros(size(t));
    F = [];
    Amp = [];
    Lab = zeros(0, 2);
    for m = 0:p.mmax
        z = wv_bjzeros(m, p.kmax);
        for k = 1:numel(z)
            j = z(k);
            f = c * j / (2 * pi);
            if f > p.fmax
                continue;
            end
            % mode shape normalised: int J_m(j r)^2 r dr = J_{m+1}(j)^2 / 2,
            % times pi (m > 0) or 2 pi (m = 0) for the angle
            nrm = besselj(m + 1, j) ^ 2 / 2 * pi * (1 + (m == 0));
            ps = besselj(m, j * p.r0);                   % strike at angle 0
            pr = besselj(m, j * p.rr) * cos(m * p.thr);  % receiver
            a = ps * pr / nrm * exp(-0.5 * (j * sw) ^ 2) / (2 * pi * f);
            if abs(a) < 1e-12
                continue;
            end
            y = y + a * sin(2 * pi * f * t);
            F(end + 1) = f;
            Amp(end + 1) = abs(a);
            Lab(end + 1, :) = [m, k];
        end
    end
    % the square drum with the same fundamental: unit side, c chosen so that
    % f11 = f01, struck and heard at comparable points.  Partials more than
    % 80 dB below the loudest are left out; they cannot be heard over it.
    cs = 2 * f01 / sqrt(2);
    ysq = zeros(size(t));
    nq = floor(2 * p.fmax / cs);
    [nn, mm] = meshgrid(1:nq, 1:nq);
    fq = cs / 2 * sqrt(mm .^ 2 + nn .^ 2);
    aq = sin(mm * pi * 0.5 * (1 + p.r0 / 1.2)) .* sin(nn * pi * 0.5) ...
        .* sin(mm * pi * 0.77) .* sin(nn * pi * 0.63) ...
        .* exp(-0.5 * (pi * sqrt(mm .^ 2 + nn .^ 2) * sw / 2) .^ 2) ./ (2 * pi * fq);
    aq(fq > p.fmax) = 0;
    for q = find(abs(aq(:)) > 1e-4 * max(abs(aq(:))))'
        ysq = ysq + aq(q) * sin(2 * pi * fq(q) * t);
    end
    % one decay for every partial of both drums, so that neither rings forever
    env = exp(-t / p.tau);
    y = y .* env;
    ysq = ysq .* env;
    [F, ord] = sort(F);
    Amp = Amp(ord);
    Lab = Lab(ord, :);
    out.f = F;
    out.adb = 20 * log10(Amp / max(Amp));
    out.m = Lab(:, 1)';
    out.k = Lab(:, 2)';
    out.ratio = F / f01;
    [fr, db] = wv_spectrum(y, fs, p.fmax, 800);
    out.fr = fr;
    out.db = db;
    out.clipRound = wv_audio(y, fs);
    out.clipSquare = wv_audio(ysq, fs);
    % the first dozen partials, whatever the strike, for the table
    tbl = zeros(0, 3);
    for m = 0:5
        z = wv_bjzeros(m, 3);
        for k = 1:3
            tbl(end + 1, :) = [m, k, z(k)];
        end
    end
    [~, o] = sort(tbl(:, 3));
    tbl = tbl(o(1:12), :);
    out.tm = tbl(:, 1)';
    out.tk = tbl(:, 2)';
    out.tratio = (tbl(:, 3) / tbl(1, 3))';
else
    Ns = p.Ns(:)';
    methods = p.methods;
    modes = reshape(p.modes, 2, [])';
    nm = numel(methods);
    nq = size(modes, 1);
    CE = NaN(nm * nq, numel(Ns));
    LAM = NaN(nm, numel(Ns));
    TM = NaN(nm, numel(Ns));
    WORST = NaN(nm, numel(Ns));
    for s = 1:nm
        for j = 1:numel(Ns)
            N = Ns(j);
            h = 2 / N;
            [C, E, W, Nn, S, in] = diskstencil(N, methods{s});
            rho = drv_radius(C, E, W, Nn, S, in);
            lam = 0.95 * min(2 / sqrt(rho), 1 / sqrt(2));
            LAM(s, j) = lam;
            dt = lam * h;
            x = -1 + (0:N) * h;
            [X, Y] = meshgrid(x, x);
            R = sqrt(X .^ 2 + Y .^ 2);
            TH = atan2(Y, X);
            wst = 0;
            el = 0;
            for q = 1:nq
                m = modes(q, 1);
                z = wv_bjzeros(m, modes(q, 2));
                jm = z(end);
                phi = besselj(m, jm * R) .* cos(m * TH) .* in;
                w = jm;                               % c = 1
                M = ceil(p.periods * 2 * pi / (w * dt));
                t0 = tic;
                a = drv_leapfrog(phi, phi * cos(w * dt), C, E, W, Nn, S, lam, M, phi);
                el = el + toc(t0);
                fm = wv_peakfreq(a, dt, 0, 0.5 / dt);
                cq = wv_cents(fm, w / (2 * pi));
                CE((s - 1) * nq + q, j) = cq;
                wst = max(wst, abs(cq));
            end
            TM(s, j) = el * 1000;
            WORST(s, j) = wst;
        end
    end
    out.Ns = Ns;
    out.methods = methods;
    out.modes = wv_rows(modes);
    out.cents = wv_rows(CE);
    out.worst = wv_rows(WORST);
    out.lam = wv_rows(LAM);
    out.ms = wv_rows(TM);
    Q = NaN(nm, numel(Ns));
    for s = 1:nm
        Q(s, :) = wv_orders(WORST(s, :), Ns);
    end
    out.orders = wv_rows(Q);
end

wv_write(out);

function a = drv_leapfrog(u, uold, C, E, W, Nn, S, lam, M, phi)
% leapfrog with the script's stencil; returns the projection of the field
% onto phi after each step
n1 = size(u, 1);
l2 = lam ^ 2;
unew = zeros(n1, n1);
a = zeros(M, 1);
for n = 1:M
    s = 0;
    for j = 2:n1 - 1
        for i = 2:n1 - 1
            v = 2 * u(i, j) - uold(i, j) + l2 * (C(i, j) * u(i, j) + E(i, j) * u(i, j + 1) ...
                + W(i, j) * u(i, j - 1) + Nn(i, j) * u(i + 1, j) + S(i, j) * u(i - 1, j));
            unew(i, j) = v;
            s = s + phi(i, j) * v;
        end
    end
    uold = u;
    u = unew;
    a(n) = s;
end
end

function r = drv_radius(C, E, W, Nn, S, in)
% the spectral radius of the stencil, by power iteration: leapfrog with this
% stencil is stable when lam^2 r <= 4
n1 = size(C, 1);
v = double(in) .* (mod((1:n1)' + (1:n1), 2) * 2 - 1);   % a checkerboard: the roughest vector
r = 0;
w = zeros(n1, n1);
for it = 1:300
    for j = 2:n1 - 1
        for i = 2:n1 - 1
            w(i, j) = C(i, j) * v(i, j) + E(i, j) * v(i, j + 1) + W(i, j) * v(i, j - 1) ...
                + Nn(i, j) * v(i + 1, j) + S(i, j) * v(i - 1, j);
        end
    end
    nv = sqrt(sum(w(:) .^ 2));
    rn = nv / sqrt(sum(v(:) .^ 2));
    v = w / nv;
    if abs(rn - r) < 1e-6 * rn
        r = rn;
        break;
    end
    r = rn;
end
end