%% m13_sos_tomography_limited_angle.m
% Echo-mode (CUTE-type) sound-speed tomography toy with limited angular aperture.
% Purpose : Unknown: slowness deviation ds = 1/c - 1/c0 on a 24 x 20 grid of
%           1.6 mm pixels (x = +-19.2 mm under a 38.4 mm, 128 x 0.3 mm array;
%           z = 3 to 35 mm). For every pixel r (echoes come from speckle
%           everywhere) and steering angle th, the beamforming delay error is
%           the line integral of ds along the straight path from r to the
%           array at angle th:  L_th(r) = int ds dl.
%           Data = common-midangle differences between transmit/receive angle
%           pairs:  d = L_i + L_j - 2 L_m, with m = (i+j)/2, K = 7 angles in
%           [-thmax, thmax]. The absolute delay is never measured.
%           Inversion: Tikhonov with a gradient penalty (Hansen convention)
%              min ||A ds - d||^2 + lam^2 ||G ds||^2
%           and anisotropic TV by ADMM. Data are generated by fine ray
%           sampling of a smooth-edged truth (0.05 mm steps) and inverted with
%           bilinear 0.4 mm steps on the coarse grid, to avoid the inverse crime.
%           The 0-3 mm band above the grid is lumped into the top row (clamped weights).
%           The model-resolution matrix R = (A'A + lam^2 G'G)^-1 A'A shows
%           what the inversion can see: row k = how pixel k mixes its neighbours.
% Change  : thmax_list, noise_ns, the truth (kind = 'vpair', 'disk', 'hpair').
% Expect  : two inclusions 12 mm apart in depth merge into one vertical streak
%           at 10 deg and separate at 30 deg; the resolution row is several
%           times longer in z than in x; a horizontal layer is only weakly
%           sensed (its data scale as 2(sec th - 1) ~ th^2).
% Runtime : about 10 s in Octave.
clear; close all; rng(1);

c0 = 1540;                                   % reference sound speed (m/s)
Nx = 24; Nz = 20; dx = 1.6;                  % grid (pixels, mm)
xc = ((0:Nx-1) - (Nx-1)/2)*dx;               % lateral centres (mm)
zc = 3 + ((0:Nz-1) + 0.5)*dx;                % depth centres (mm)
Xh = 19.2;                                   % aperture half-width (mm)
np = Nx*Nz; K = 7; noise_ns = 5;             % unknowns, angles, delay noise std (ns)
kind = 'vpair';
sg = @(r, r0) 1./(1 + exp((r - r0)/0.4));    % smooth edge, 0.4 mm
switch kind
    case 'vpair', truth = @(x, z) 1540 + 50*(sg(hypot(x, z - 13), 3) + sg(hypot(x, z - 25), 3));
    case 'hpair', truth = @(x, z) 1540 + 50*(sg(hypot(x - 6, z - 19), 3) + sg(hypot(x + 6, z - 19), 3));
    otherwise,    truth = @(x, z) 1540 + 50*sg(hypot(x, z - 19), 5);
end
dsf = @(c) (1./c - 1/c0)*1e6;                % slowness deviation in ns/mm
[XX, ZZ] = meshgrid(xc, zc); ctrue = truth(XX, ZZ); ctrue = reshape(ctrue', [], 1);  % row-major: index = iz*Nx + ix

% gradient operator G (horizontal and vertical neighbour differences)
I = speye(np); G = [];
for iz = 0:Nz-1, for ix = 0:Nx-2, k = iz*Nx + ix + 1; G = [G; I(k + 1, :) - I(k, :)]; end, end %#ok<AGROW>
for iz = 0:Nz-2, for ix = 0:Nx-1, k = iz*Nx + ix + 1; G = [G; I(k + Nx, :) - I(k, :)]; end, end %#ok<AGROW>
GG = full(G'*G);                              % G stays sparse for fast products

fprintf('Layer sensitivity 2(sec th - 1) for the symmetric pair at +-th: 10 deg %.3f, 20 deg %.3f, 30 deg %.3f\n', ...
    2*(sec(pi/18) - 1), 2*(sec(pi/9) - 1), 2*(sec(pi/6) - 1));
kz0 = find(abs(zc - 19) == min(abs(zc - 19)), 1); kx0 = find(abs(xc - 0.8) == min(abs(xc - 0.8)), 1);
kpix = (kz0 - 1)*Nx + kx0;                   % pixel whose resolution row we inspect (x = 0.8, z = 19.0 mm)
fwhm = @(p, dd) dd*0.01*sum(interp1(1:numel(p), p, 1:0.01:numel(p)) >= max(p)/2);
thmax_list = [10 20 30];
figure('Position', [40 40 1100 800]); colormap(jet);
for q = 1:numel(thmax_list)
    tic;
    th = linspace(-thmax_list(q), thmax_list(q), K)*pi/180;
    Lr = zeros(np, np, K); Lt = zeros(np, K); ok = false(np, K);
    for a = 1:K
        S = []; W = []; Rw = [];
        for iz = 1:Nz
            rows = (iz - 1)*Nx + (1:Nx)'; z = zc(iz); x = xc(:);              % one depth row at a time
            ok(rows, a) = abs(x - z*tan(th(a))) <= Xh;
            Lp = z/cos(th(a));
            ns = ceil(Lp/0.4); tt = ((1:ns) - 0.5)*Lp/ns; w = Lp/ns;          % coarse model, 0.4 mm steps
            gx = min(max((x - tt*sin(th(a)) - xc(1))/dx, 0), Nx - 1 - 1e-9);  % Nx x ns grid coordinates
            gz = min(max((z - tt*cos(th(a)) - zc(1))/dx, 0), Nz - 1 - 1e-9);
            gz = repmat(gz, Nx, 1);
            i0 = floor(gx); j0 = floor(gz); fx = gx - i0; fz = gz - j0;
            i1 = min(i0 + 1, Nx - 1); j1 = min(j0 + 1, Nz - 1);
            S  = [S;  j0(:)*Nx + i0(:) + 1; j0(:)*Nx + i1(:) + 1; j1(:)*Nx + i0(:) + 1; j1(:)*Nx + i1(:) + 1]; %#ok<AGROW>
            W  = [W;  w*[(1 - fx(:)).*(1 - fz(:)); fx(:).*(1 - fz(:)); (1 - fx(:)).*fz(:); fx(:).*fz(:)]];    %#ok<AGROW>
            Rw = [Rw; repmat(repmat(rows, ns, 1), 4, 1)];                                                      %#ok<AGROW>
            nf = ceil(Lp/0.05); tf = ((1:nf) - 0.5)*Lp/nf;                     % fine truth, 0.05 mm steps
            Lt(rows, a) = sum(dsf(truth(x - tf*sin(th(a)), z - tf*cos(th(a)))), 2)*Lp/nf;
        end
        Lr(:, :, a) = accumarray([Rw S], W, [np np]);
    end
    A = []; d0 = [];
    for i = 1:K
        for j = i + 2:2:K
            m = (i + j)/2; use = ok(:, i) & ok(:, j) & ok(:, m);
            A  = [A;  Lr(use, :, i) + Lr(use, :, j) - 2*Lr(use, :, m)];   %#ok<AGROW>
            d0 = [d0; Lt(use, i) + Lt(use, j) - 2*Lt(use, m)];            %#ok<AGROW>
        end
    end
    d = d0 + noise_ns*randn(size(d0));
    AtA = A'*A; Atd = A'*d; s1 = sqrt(max(eig(AtA)));
    lr = logspace(-3, 0, 31); e = zeros(size(lr));
    for j = 1:numel(lr)
        x = (AtA + (lr(j)*s1)^2*GG)\Atd; e(j) = sqrt(mean((1./(1/c0 + x*1e-6) - ctrue).^2));
    end
    [eb, jb] = min(e); M = AtA + (lr(jb)*s1)^2*GG;
    x = M\Atd; cT = 1./(1/c0 + x*1e-6);
    Rm = M\AtA; row = reshape(Rm(kpix, :), Nx, Nz)';            % resolution row as an Nz x Nx image
    fprintf('\nthmax = %2d deg: %d data, rms data %.0f ns, sigma_max %.0f\n', thmax_list(q), size(A, 1), sqrt(mean(d0.^2)), s1);
    fprintf('  Tikhonov-gradient, oracle lambda = %.4f sigma_max: RMSE %.1f m/s, max c %.0f m/s (true 1590), trace(R) = %.0f of %d\n', ...
        lr(jb), eb, max(cT), trace(Rm), np);
    fprintf('  resolution row at (x, z) = (%.1f, %.1f) mm: FWHM axial %.1f mm, lateral %.1f mm\n', ...
        xc(kx0), zc(kz0), fwhm(row(:, kx0), dx), fwhm(row(kz0, :), dx));
    % ---- TV by ADMM: min 0.5||A x - d||^2 + mu ||G x||_1
    rho = 1e-2*s1^2; Rc = chol(AtA + rho*GG); mus = logspace(0, 2.5, 8); Gt = G'; ev = zeros(size(mus));
    for j = 1:numel(mus)
        zz = zeros(size(G, 1), 1); u = zz;
        for it = 1:200
            x = Rc\(Rc'\(Atd + rho*Gt*(zz - u))); w = G*x + u;
            zz = sign(w).*max(abs(w) - mus(j)/rho, 0); u = w - zz;
        end
        ev(j) = sqrt(mean((1./(1/c0 + x*1e-6) - ctrue).^2));
        if j == 1 || ev(j) < min(ev(1:j-1)), cV = 1./(1/c0 + x*1e-6); end
    end
    [evb, jv] = min(ev);
    fprintf('  TV (ADMM, 200 it), oracle mu = %.3g: RMSE %.1f m/s, max c %.0f m/s   [%.1f s]\n', mus(jv), evb, max(cV), toc);
    subplot(3, 4, 4*(q - 1) + 1); imagesc(xc, zc, reshape(ctrue, Nx, Nz)', [1520 1600]); axis image; title('truth c (m/s)'); ylabel(sprintf('%d deg: z (mm)', thmax_list(q)));
    subplot(3, 4, 4*(q - 1) + 2); imagesc(xc, zc, reshape(cT, Nx, Nz)', [1520 1600]); axis image; title('Tikhonov-gradient');
    subplot(3, 4, 4*(q - 1) + 3); imagesc(xc, zc, reshape(cV, Nx, Nz)', [1520 1600]); axis image; title('TV');
    subplot(3, 4, 4*(q - 1) + 4); imagesc(xc, zc, row); axis image; title('resolution row R(k,:)'); colorbar;
end
xlabel('x (mm)');