Optimization & automatic differentiation

%% m14_gd_conditioning.m  --  gradient descent, momentum and the condition number
% Purpose : minimize the quadratic f(x) = 1/2 x'*A*x (minimum at 0) with eigenvalues
%           mu = 1 and L = kappa, and count iterations to reduce the error norm by 1e-6.
%           Compares gradient descent with the optimal fixed step 2/(L+mu), with step 1/L,
%           Polyak heavy-ball with its optimal (alpha, beta), and Nesterov's method
%           (step 1/L, momentum (sqrt(kappa)-1)/(sqrt(kappa)+1)) against the rate formulas.
%           Then shows that diagonal rescaling of badly scaled parameters (here a sound
%           speed in m/s and a dimensionless amplitude) removes most of kappa.
% Change  : kappa_list, tol, the rotation angle th, the momentum settings.
% Expect  : kappa = 100: GD(2/(L+mu)) 691 iterations (theory 691), GD(1/L) 1341 (1375; this
%           start puts only 1/sqrt(2) of the error on the slow mode), heavy ball 93 and
%           Nesterov 155 (rate-only predictions 69 and 131: both tuned iterations have a
%           double root, so the error decays as k*rho^k, not rho^k). Jacobi scaling of a
%           (m/s, amplitude) problem drops kappa(J'J) from about 1e6 to about 1.
% Base MATLAB only; runs in MATLAB R2019b+ and GNU Octave 8.
rng(1);

%% 1. Iterations to tolerance vs kappa
tol = 1e-6;  kappa_list = [10 100 1000];  maxit = 1e5;
th = pi/6;  Q = [cos(th) -sin(th); sin(th) cos(th)];          % rotate so the axes are not the eigenvectors
fprintf('kappa | GD 2/(L+mu)  theory | GD 1/L  theory | heavy ball  theory | Nesterov  theory\n');
hist100 = {};
for kappa = kappa_list
    mu = 1;  L = kappa;  A = Q * diag([mu L]) * Q.';
    x0 = Q * [1; 1];
    sk = sqrt(kappa);
    cfg = { 2/(L+mu), 0, 'hb';                                  % {alpha, beta, type}
            1/L,      0, 'hb';
            4/(sqrt(L)+sqrt(mu))^2, ((sk-1)/(sk+1))^2, 'hb';
            1/L,      (sk-1)/(sk+1), 'nest' };
    its = zeros(1, 4);
    for m = 1:4
        a = cfg{m,1};  b = cfg{m,2};  x = x0;  xp = x0;  e = zeros(1, maxit);
        for k = 1:maxit
            e(k) = norm(x) / norm(x0);
            if e(k) < tol, break; end
            if strcmp(cfg{m,3}, 'nest')
                y = x + b*(x - xp);  xn = y - a*(A*y);
            else
                xn = x - a*(A*x) + b*(x - xp);
            end
            xp = x;  x = xn;
        end
        its(m) = k - 1;
        if kappa == 100, hist100{m} = e(1:k); end
    end
    th_gd = log(1/tol) / log((kappa+1)/(kappa-1));
    th_1L = log(1/tol) / -log(1 - 1/kappa);
    th_hb = log(1/tol) / log((sk+1)/(sk-1));
    th_ne = log(1/tol) / -log(1 - 1/sk);
    fprintf('%5d | %6d  %6.0f     | %6d  %6.0f | %6d  %6.0f     | %6d  %6.0f\n', ...
        kappa, its(1), th_gd, its(2), th_1L, its(3), th_hb, its(4), th_ne);
end

%% 2. Divergence threshold: step just above 2/L
kappa = 100;  A = Q * diag([1 kappa]) * Q.';  x = Q * [1; 1];
for k = 1:50, x = x - (2.05/kappa) * (A*x); end
fprintf('step 2.05/L after 50 iterations: |x| = %.3g (diverges; factor |1 - 2.05| = 1.05 per step)\n', norm(x));

%% 3. Badly scaled parameters: sound speed (m/s) and amplitude (unitless)
% Linearized model y = J*[dc; dA] with column sensitivities that differ by 3 orders.
J = [randn(50,1) * 1e-3, randn(50,1)];          % d y / d c (per m/s) and d y / d A
H = J.' * J;  D = diag(1 ./ sqrt(diag(H)));     % Jacobi (diagonal) preconditioner
fprintf('kappa(J''J) raw %.3g  ->  after column scaling %.3g\n', cond(H), cond(D*H*D));
kr = cond(H);  ks = cond(D*H*D);
fprintf('GD iterations to 1e-6 (optimal step): raw %.3g, scaled %.0f\n', ...
    log(1e6)/log((kr+1)/(kr-1)), log(1e6)/log((ks+1)/(ks-1)));

%% 4. Figures: error histories (kappa = 100) and paths on the contours
figure('Position', [100 100 1050 420]);
subplot(1,2,1);
cl = {[0 0.3 0.8], [0.6 0.6 0.6], [0.8 0.2 0.5], [0 0.55 0.4]};  hh = zeros(1,6);
for m = 1:4
    hh(m) = semilogy(0:numel(hist100{m})-1, hist100{m}, 'Color', cl{m}, 'LineWidth', 1.5); hold on;
end
kk = 0:700;
hh(5) = semilogy(kk, ((100-1)/(100+1)).^kk, '--', 'Color', cl{1});
hh(6) = semilogy(kk, ((10-1)/(10+1)).^kk, '--', 'Color', cl{3});
grid on; ylim([1e-7 2]); xlim([0 700]); xlabel('iteration k'); ylabel('||x_k|| / ||x_0||');
legend(hh, {'GD 2/(L+\mu)', 'GD 1/L', 'heavy ball', 'Nesterov', '((\kappa-1)/(\kappa+1))^k', '((\surd\kappa-1)/(\surd\kappa+1))^k'}, 'Location', 'northeast');
title('\kappa = 100: measured error vs rate formulas');
subplot(1,2,2);
kappa = 30;  A = Q * diag([1 kappa]) * Q.';
[g1, g2] = meshgrid(linspace(-1.6, 1.6, 200), linspace(-1.6, 1.6, 200));
F = 0.5 * (A(1,1)*g1.^2 + 2*A(1,2)*g1.*g2 + A(2,2)*g2.^2);
contour(g1, g2, log10(F + 1e-3), 14); hold on; colormap(gray);
x0 = Q * [1.2; 0.25];  hp = zeros(1,2);
for m = 1:2
    if m == 1, a = 2/(kappa+1); b = 0; col = [0 0.3 0.8]; else, sk = sqrt(kappa); a = 4/(sk+1)^2; b = ((sk-1)/(sk+1))^2; col = [0.8 0.2 0.5]; end
    x = x0; xp = x0; P = x0;
    for k = 1:60, xn = x - a*(A*x) + b*(x - xp); xp = x; x = xn; P(:, end+1) = x; end
    hp(m) = plot(P(1,:), P(2,:), '.-', 'Color', col, 'LineWidth', 1.2, 'MarkerSize', 10);
end
axis image; xlabel('x_1'); ylabel('x_2'); title('\kappa = 30: GD zig-zags, heavy ball spirals in');
legend(hp, {'GD 2/(L+\mu)', 'heavy ball (optimal \alpha, \beta)'}, 'Location', 'southoutside');

%% m14_gauss_newton_lm.m  --  gradient descent vs Gauss-Newton vs Levenberg-Marquardt
% Purpose : fit a Gaussian-pulse echo model to a noisy RF echo by nonlinear least squares
%             s(t; A, t0, sg, f) = A exp(-(t-t0)^2/(2 sg^2)) cos(2 pi f (t - t0))
%           (units: t, t0, sg in us; f in MHz) and compare
%             GD + Armijo backtracking, Gauss-Newton (J'J d = -J'r), and
%             Levenberg-Marquardt ((J'J + lam diag(J'J)) d = -J'r, lam /10 on success, x10 on failure).
%           Prints the iterations, the estimates, the linearized standard deviations
%           sqrt(diag(sn^2 (J'J)^-1)) (the Gaussian-noise CRLB of Module 12 evaluated at the
%           solution), kappa(J'J) before and after column scaling, and shows that a start
%           more than half a period off in t0 converges to the wrong cycle unless the model
%           has a free carrier phase.
% Change  : theta_true, sn (noise), the starts, lam0.
% Expect  : good start: GN and LM about 10 iterations, GD + Armijo about 10^4 in raw units
%           and tens after Jacobi scaling; bad start (t0 off by 100 ns, more than T/4 = 56 ns):
%           LM lands half a period late with the polarity flipped (A < 0, t0 near 2.11 us,
%           |r| three times the noise level); the free-phase model recovers t0 = 2.00 us.
% Base MATLAB only; runs in MATLAB R2019b+ and GNU Octave 8. Runtime about 2 s.
rng(1);

%% 1. Data
fs = 50;                          % sampling rate (MHz) -> dt = 0.02 us
t = (1:1/fs:3).';                 % time axis (us), 101 samples
th_true = [1, 2.00, 0.15, 4.5];   % [A, t0 (us), sg (us), f (MHz)]
sn = 0.05;                        % noise std (peak SNR = 20 log10(1/0.05) = 26 dB)
model = @(p) p(1) * exp(-(t-p(2)).^2/(2*p(3)^2)) .* cos(2*pi*p(4)*(t-p(2)));
y = model(th_true) + sn * randn(size(t));
% analytic Jacobian of the residual r = model - y (columns: A, t0, sg, f)
jac = @(p) [exp(-(t-p(2)).^2/(2*p(3)^2)) .* cos(2*pi*p(4)*(t-p(2))), ...
     p(1)*exp(-(t-p(2)).^2/(2*p(3)^2)) .* ((t-p(2))/p(3)^2 .* cos(2*pi*p(4)*(t-p(2))) + 2*pi*p(4)*sin(2*pi*p(4)*(t-p(2)))), ...
     p(1)*exp(-(t-p(2)).^2/(2*p(3)^2)) .* cos(2*pi*p(4)*(t-p(2))) .* (t-p(2)).^2/p(3)^3, ...
    -p(1)*exp(-(t-p(2)).^2/(2*p(3)^2)) .* sin(2*pi*p(4)*(t-p(2))) .* 2*pi.*(t-p(2))];
cost = @(p) 0.5 * sum((model(p) - y).^2);
% check the Jacobian against central differences
p = [0.9 2.03 0.13 4.8];  Jn = zeros(numel(t), 4);
for k = 1:4, e = zeros(1,4); e(k) = 1e-6; Jn(:,k) = (model(p+e) - model(p-e)) / 2e-6; end
fprintf('Jacobian check: max |analytic - numeric| / max|J| = %.1e\n', max(max(abs(jac(p)-Jn))) / max(abs(Jn(:))));

%% 2. Solvers from a good start and from a start beyond a quarter period
% GD+Armijo in raw units (step capped at 0.01), GD+Armijo after Jacobi (column) scaling,
% Gauss-Newton (least-squares solve of J d = -r) and Levenberg-Marquardt.
starts = {[0.8, 2.04, 0.12, 4.8], [0.8, 2.10, 0.12, 4.5]};
labels = {'good start (dt0 = +40 ns, df = +0.3 MHz)', 'bad start (dt0 = +100 ns > T/4)'};
nm = {'GD+Armijo', 'GD+Armijo+scale', 'Gauss-Newton', 'LM'};
hist = cell(2, 4);
for s = 1:2
    fprintf('\n%s; period T = %.0f ns\n', labels{s}, 1e3/th_true(4));
    for meth = 1:4
        p = starts{s};  lam = 1e-3;  al = 1e-2;  h = cost(p);  it = 0;
        if meth == 2, Dsc = 1 ./ sqrt(sum(jac(p).^2, 1)); amax = 1; else, Dsc = ones(1,4); amax = 1e-2; end
        while it < 20000
            r = model(p) - y;  J = jac(p);  g = (J.' * r).';
            if meth <= 2                                  % (scaled) GD + Armijo backtracking
                d = -(Dsc.^2) .* g;  al = min(amax, 2*al);
                while cost(p + al*d) > cost(p) + 1e-4*al*(g*d.'), al = al/2; end
                pn = p + al*d;
            elseif meth == 3                              % Gauss-Newton
                pn = p - (J \ r).';                       % least-squares solve, J'J never formed
            else                                          % Levenberg-Marquardt
                H = J.'*J;
                while true
                    pn = p - ((H + lam*diag(diag(H))) \ g.').';
                    if cost(pn) < cost(p), lam = lam/10; break; end
                    lam = lam*10;  if lam > 1e12, pn = p; break; end
                end
            end
            it = it + 1;  step = norm(pn - p);  p = pn;  h(end+1) = cost(p);
            if step < 1e-9, break; end
        end
        hist{s, meth} = h;
        fprintf('%-16s %5d its  A = %7.4f  t0 = %.4f us  sg = %.4f us  f = %.4f MHz  |r| = %.4f\n', ...
            nm{meth}, it, p(1), p(2), abs(p(3)), abs(p(4)), sqrt(2*cost(p)));
        if s == 1 && meth == 4, p_hat = p; end
    end
end
fprintf('expected |r| for a correct fit: sn*sqrt(n) = %.3f\n', sn*sqrt(numel(t)));

%% 3. Precision at the solution and conditioning
J = jac(p_hat);  H = J.'*J;  sd = sqrt(diag(sn^2 * (H \ eye(4))));
fprintf('\nlinearized std (CRLB at the fit): A %.4f, t0 %.2f ns, sg %.2f ns, f %.4f MHz\n', sd(1), 1e3*sd(2), 1e3*sd(3), sd(4));
Dg = diag(1 ./ sqrt(diag(H)));
fprintf('kappa(J''J) = %.3g in (1, us, us, MHz) units; %.3g after column scaling\n', cond(H), cond(Dg*H*Dg));
Hs = diag([1 1e-6 1e-6 1e6]) * H * diag([1 1e-6 1e-6 1e6]);
fprintf('kappa(J''J) = %.3g if the same fit is written in SI units (1, s, s, Hz)\n', cond(Hs));

%% 4. Free carrier phase removes the cycle ambiguity in t0
modelp = @(q) q(1) * exp(-(t-q(2)).^2/(2*q(3)^2)) .* cos(2*pi*q(4)*(t-q(2)) + q(5));
costp = @(q) 0.5 * sum((modelp(q) - y).^2);
q = [starts{2} 0];  lam = 1e-3;
for it = 1:200
    r = modelp(q) - y;  Jq = zeros(numel(t), 5);
    for k = 1:5, e = zeros(1,5); e(k) = 1e-7; Jq(:,k) = (modelp(q+e) - modelp(q-e)) / 2e-7; end
    H = Jq.'*Jq;  g = Jq.'*r;
    while true
        qn = q - ((H + lam*diag(diag(H))) \ g).';
        if costp(qn) < costp(q), lam = lam/10; break; end
        lam = lam*10;  if lam > 1e12, qn = q; break; end
    end
    if norm(qn - q) < 1e-10, q = qn; break; end
    q = qn;
end
fprintf('free-phase LM from the bad start: t0 = %.4f us, phase = %.3f rad, |r| = %.4f (%d its)\n', q(2), angle(exp(1j*q(5))), sqrt(2*costp(q)), it);

%% 5. Figure
figure('Position', [100 100 1050 420]);
subplot(1,2,1);
cols = {[0.6 0.6 0.6], [0.85 0.55 0], [0 0.3 0.8], [0.8 0.2 0.5]};  hh = zeros(1,4);
for meth = 1:4
    hh(meth) = semilogy(1:numel(hist{1,meth}), hist{1,meth} - min(hist{1,4}) + 1e-12, 'Color', cols{meth}, 'LineWidth', 1.5); hold on;
end
set(gca, 'XScale', 'log'); grid on; xlim([1 2e4]);
xlabel('iteration'); ylabel('cost - final cost');
legend(hh, {'GD + Armijo (raw units)', 'GD + Armijo (Jacobi scaled)', 'Gauss-Newton', 'Levenberg-Marquardt'}, 'Location', 'southwest');
title('good start: linear vs quadratic convergence');
subplot(1,2,2);
plot(t, y, 'Color', [0.6 0.6 0.6]); hold on;
plot(t, model(p_hat), 'b', 'LineWidth', 1.5);
plot(t, model(starts{2}), 'r--');
xlabel('time (\mus)'); ylabel('amplitude'); grid on; xlim([1.4 2.6]);
legend('data', 'LM fit (good start)', 'bad start model', 'Location', 'southwest');
title('echo, fit, and a start that lands half a period late');

%% m14_cycle_skipping.m  --  why RF misfits trap gradient descent, and how to widen the basin
% Purpose : delay estimation by minimizing L(tau) = sum (d(t) - s(t - tau))^2 for a 5 MHz
%           Gaussian pulse (60 % bandwidth). Compares three misfits computed numerically
%           on sampled signals: RF, envelope (|analytic signal|, FFT Hilbert), and RF after
%           a Gaussian low-pass at 1.5 or 1.0 MHz. Lists the local minima, runs gradient
%           descent from 121 starting delays, and shows that a multiscale schedule (envelope
%           first, then RF) converges from starts the RF misfit alone cannot handle. Ends with
%           the travel-time form of the criterion used in full-waveform inversion.
% Change  : f0, sg, the low-pass corner fl, the starting delays.
% Expect  : RF local minima at 0, +-194, +-388 ns (just inside multiples of T0 = 200 ns);
%           RF basin |tau0| < 97 ns; envelope converges from every start (no local minima);
%           low-pass at 1.5 MHz: basin about +-160 ns at -29 dB of the pulse energy;
%           envelope first, then RF: converges from every start with RF curvature (32x
%           the envelope curvature at the minimum).
% Base MATLAB only; runs in MATLAB R2019b+ and GNU Octave 8. Runtime about 5 s.
rng(1);

%% 1. Signals
f0 = 5e6; sg = 0.125e-6;           % centre frequency (Hz), envelope std (s)
fs = 1e9; t = (-3e-6:1/fs:3e-6).';  % fine grid so delays are continuous (1 ns)
nT = numel(t);  nf = 2^nextpow2(2*nT);
fax = [0:nf/2-1, -nf/2:-1].' * fs / nf;              % FFT frequency axis (Hz)
pulse = @(tau) exp(-(t-tau).^2/(2*sg^2)) .* cos(2*pi*f0*(t-tau));
fl = 1.5e6;                                           % Gaussian low-pass corner (Hz)
lowp  = @(x) real(ifft(fft(x, nf) .* exp(-fax.^2/(2*fl^2))));
H     = @(x) ifft(fft(x, nf) .* (1 + sign(fax)));     % analytic signal (x + j Hilbert x)
crop  = @(x) x(1:nT);
envl  = @(x) abs(crop(H(x)));
d = pulse(0);                                          % data: true delay 0
fprintf('low-pass at %.1f MHz keeps %.1f dB of the pulse energy\n', fl/1e6, 10*log10(sum(crop(lowp(d)).^2)/sum(d.^2)));

%% 2. Misfit curves, computed numerically from the sampled signals
% L(tau) = sum (a - b_tau)^2 = sum a^2 + sum b^2 - 2 R_ab(tau): all shifts at once by FFT
% correlation (a shift commutes with the low-pass filter and with taking the envelope).
xc = @(a, b) real(ifft(fft(a, nf) .* conj(fft(b, nf))));
lagk = -600:600;  taus = lagk * 1e-9;                  % 1 ns lags
mis = @(a, b) sum(a.^2) + sum(b.^2) - 2 * subsref(xc(b, a), struct('type', '()', 'subs', {{mod(lagk, nf) + 1}})).';
dE = envl(d);  dL = crop(lowp(d));
fl2 = 1.0e6;  lowp2 = @(x) real(ifft(fft(x, nf) .* exp(-fax.^2/(2*fl2^2))));  dL2 = crop(lowp2(d));
L = [mis(d, d); mis(dE, dE); mis(dL, dL); mis(dL2, dL2)];
L = L ./ (2*[sum(d.^2); sum(dE.^2); sum(dL.^2); sum(dL2.^2)]);   % normalize by 2 x energy: L = 1 - R(tau)/R(0)
fprintf('low-pass at %.1f MHz keeps %.1f dB of the pulse energy\n', fl2/1e6, 10*log10(sum(dL2.^2)/sum(d.^2)));
names = {'RF', 'envelope', 'LP 1.5 MHz', 'LP 1.0 MHz'};
for m = 1:4
    loc = find(L(m,2:end-1) < L(m,1:end-2) & L(m,2:end-1) < L(m,3:end)) + 1;
    fprintf('%-11s local minima (ns): %s\n', names{m}, sprintf('%g ', round(taus(loc)*1e9)));
end
mx = find(L(1,2:end-1) > L(1,1:end-2) & L(1,2:end-1) > L(1,3:end)) + 1;
fprintf('RF local maxima nearest 0 (basin edges): %s ns; T0/2 = %.0f ns\n', sprintf('%g ', round(taus(mx(abs(taus(mx)) < 150e-9))*1e9)), 0.5e9/f0);

%% 3. Gradient descent from many starts on the tabulated misfits
% dL/dtau by central differences on the 1 ns table; step = 1/L''(0) (Newton step at the minimum)
dLdt = zeros(size(L));  dLdt(:,2:end-1) = (L(:,3:end) - L(:,1:end-2)) / 2e-9;
curv = (L(:,602) - 2*L(:,601) + L(:,600)) / 1e-18;    % L'' at tau = 0
fprintf('curvature at the minimum, RF / envelope = %.1f (theory 1 + 2 (2 pi f0 sg)^2 = %.1f)\n', curv(1)/curv(2), 1 + 2*(2*pi*f0*sg)^2);
starts = (-600:10:600) * 1e-9;
plan = {1, 2, 3, 4, [2 1]};                           % misfits used in sequence
lbl = {'RF', 'envelope', 'LP 1.5 MHz', 'LP 1.0 MHz', 'envelope -> RF'};
final = zeros(numel(plan), numel(starts));
for m = 1:numel(plan)
    tau = starts;
    for q = plan{m}
        for it = 1:400
            g = interp1(taus, dLdt(q,:), min(max(tau, taus(2)), taus(end-1)));
            tau = min(max(tau - g / curv(q), taus(1)), taus(end));   % stay on the table
        end
    end
    final(m, :) = tau;
end
ok = abs(final) < 2e-9;
for m = 1:numel(plan)
    k0 = find(starts == 0);  lo = k0;  hi = k0;
    while lo > 1 && ok(m, lo-1), lo = lo - 1; end
    while hi < numel(starts) && ok(m, hi+1), hi = hi + 1; end
    fprintf('%-15s converges to 0 from %3d of %d starts; contiguous basin [%4.0f, %4.0f] ns\n', ...
        lbl{m}, sum(ok(m,:)), numel(starts), starts(lo)*1e9, starts(hi)*1e9);
end

%% 4. Travel-time form of the criterion (transmission FWI / tomography)
% the starting model must predict the travel time within T/2: |dt| = L |dc| / c^2 < 1/(2 f)
fprintf('\nmax tolerable starting error |dc| (m/s) for c = 1500 m/s\n   f (MHz) | L = 50 mm | L = 100 mm | L = 200 mm\n');
for f = [0.5 1 2 5]
    fprintf('   %5.1f   | %8.1f  | %9.1f  | %9.1f\n', f, 1500^2./(2*f*1e6*[0.05 0.1 0.2]));
end

%% 5. Figures
figure('Position', [100 100 1050 420]);
subplot(1,2,1);
h = plot(taus*1e9, L(1,:), 'b', taus*1e9, L(2,:), 'r', taus*1e9, L(3,:), 'g');
set(h(3), 'Color', [0 0.55 0.4]); set(h, 'LineWidth', 1.4); grid on;
xlabel('trial delay \tau (ns), true delay 0'); ylabel('misfit / (2 x signal energy)');
xlim([-600 600]); legend(h, names(1:3), 'Location', 'southeast'); title('RF misfit: a minimum every period');
subplot(1,2,2);
h2 = plot(starts*1e9, final(1,:)*1e9, 'b.-', starts*1e9, final(3,:)*1e9, 'g.-', starts*1e9, final(5,:)*1e9 + 8, 'k.-');
set(h2(2), 'Color', [0 0.55 0.4]); grid on; xlim([-600 600]); xlabel('starting delay (ns)'); ylabel('converged delay (ns)');
legend(h2, {'RF', 'LP 1.5 MHz', 'envelope -> RF (offset +8 ns)'}, 'Location', 'northwest');
title('where gradient descent ends');

%% m14_complex_step_ad.m  --  finite differences, complex step, dual numbers, reverse mode, Wirtinger
% Purpose : 1) relative error of forward/central differences and of the complex-step
%              derivative Im f(x + i h)/h versus h, for f(x) = e^x / sqrt(sin^3 x + cos^3 x)
%              at x0 = 1.5 (a standard test function in the complex-step literature);
%           2) forward-mode AD with dual numbers [value, derivative] built from anonymous
%              functions; 3) reverse mode by hand on y = ln(x1) + x1 x2 - sin(x2) at (2, 5);
%           4) two ways the complex step silently fails (abs and the ' operator);
%           5) Wirtinger gradient dL/dw* = -A^H (y - A w) for L = ||y - A w||^2 checked
%              against finite differences, and what happens if you step along its conjugate.
% Change  : x0, the step list, the test functions.
% Expect  : best forward difference 1e-8 to 1e-9 relative error near h = 1e-8; best central
%           about 1e-11 near h = 3e-6; complex step at machine precision for every
%           h <= 1e-8 (including 1e-20 and 1e-100); dual number exact; reverse-mode
%           gradient (5.5, 1.7163); Wirtinger check about 1e-9; the conjugate step diverges.
% Base MATLAB only; runs in MATLAB R2019b+ and GNU Octave 8. Runtime about 1 s.
rng(1);

%% 1. Finite differences vs complex step
f  = @(x) exp(x) ./ sqrt(sin(x).^3 + cos(x).^3);
x0 = 1.5;
df_exact = f(x0) * (1 - (3*sin(x0)^2*cos(x0) - 3*cos(x0)^2*sin(x0)) / (2*(sin(x0)^3 + cos(x0)^3)));
fprintf('f(1.5) = %.6f, exact f''(1.5) = %.12f\n', f(x0), df_exact);
hs = 10.^(-1:-0.25:-20);
f0v = f(x0 + 0*hs);                       % same vector code path as f(x0+hs): avoids a 1-ulp mismatch
efd = abs((f(x0+hs) - f0v) ./ hs - df_exact) / abs(df_exact);
ecd = abs((f(x0+hs) - f(x0-hs)) ./ (2*hs) - df_exact) / abs(df_exact);
ecs = abs(imag(f(x0 + 1i*hs)) ./ hs - df_exact) / abs(df_exact);
pfd = max(efd, 1e-17);  pcd = max(ecd, 1e-17);  pcs = max(ecs, 1e-17);       % floor zeros for the log plot
fprintf('     h     | forward diff | central diff | complex step\n');
for h = [1e-2 1e-4 1e-6 1e-8 1e-10 1e-12 1e-16 1e-20]
    k = find(abs(log10(hs) - log10(h)) < 1e-9);
    fprintf('  %7.0e  |   %8.1e   |   %8.1e   |   %8.1e\n', h, efd(k), ecd(k), ecs(k));
end
[m1, i1] = min(efd);  [m2, i2] = min(ecd);
fprintf('best forward h = %.0e (err %.1e; sqrt(eps) = %.1e); best central h = %.0e (err %.1e; eps^(1/3) = %.1e)\n', ...
    hs(i1), m1, sqrt(eps), hs(i2), m2, eps^(1/3));
fprintf('complex step with h = 1e-100: relative error %.1e\n', abs(imag(f(x0 + 1e-100i))/1e-100 - df_exact)/df_exact);

%% 2. Forward mode with dual numbers  a = [value, derivative]
dmul = @(a, b) [a(1)*b(1), a(1)*b(2) + a(2)*b(1)];
ddiv = @(a, b) [a(1)/b(1), (a(2)*b(1) - a(1)*b(2)) / b(1)^2];
dadd = @(a, b) a + b;
dexp = @(a) [exp(a(1)), exp(a(1))*a(2)];
dsin = @(a) [sin(a(1)), cos(a(1))*a(2)];
dcos = @(a) [cos(a(1)), -sin(a(1))*a(2)];
dsqrt = @(a) [sqrt(a(1)), a(2) / (2*sqrt(a(1)))];
dcube = @(a) dmul(dmul(a, a), a);
fdual = @(a) ddiv(dexp(a), dsqrt(dadd(dcube(dsin(a)), dcube(dcos(a)))));
r = fdual([x0, 1]);                                      % seed dx/dx = 1
fprintf('\ndual number: f = %.12f, f'' = %.15f, relative error %.1e\n', r(1), r(2), abs(r(2) - df_exact)/df_exact);

%% 3. Reverse mode by hand: y = ln(x1) + x1*x2 - sin(x2) at (2, 5)
x1 = 2; x2 = 5;
v1 = log(x1); v2 = x1*x2; v3 = sin(x2); v4 = v1 + v2; y = v4 - v3;       % forward sweep (the tape)
yb = 1;                                                                   % reverse sweep: adjoints
v4b = yb;  v3b = -yb;  v1b = v4b;  v2b = v4b;
x1b = v1b/x1 + v2b*x2;  x2b = v2b*x1 + v3b*cos(x2);
fprintf('reverse mode: y = %.4f, grad = (%.4f, %.4f); analytic (1/x1 + x2, x1 - cos x2) = (%.4f, %.4f)\n', ...
    y, x1b, x2b, 1/x1 + x2, x1 - cos(x2));
fprintf('  one reverse sweep gave both partials; forward mode needs one sweep per input (2 here)\n');

%% 4. Silent failures of the complex step
g1 = @(x) x .* abs(x - 1);                % abs is not analytic: abs(z) returns |z|, a real number
fprintf('\nx|x-1| at 1.5: exact 2.0, complex step %.4f, dual %.4f\n', imag(g1(1.5 + 1e-20i))/1e-20, ...
    subsref(dmul([1.5 1], [abs(0.5), sign(0.5)*1]), struct('type', '()', 'subs', {{2}})));
g2 = @(x) x' * x;                         % ' is the CONJUGATE transpose: x'*x = |x|^2 for complex x
fprintf('x''*x at 1.5: exact 3.0, complex step %.4f;   with .'' : %.4f\n', imag(g2(1.5 + 1e-20i))/1e-20, ...
    imag((1.5 + 1e-20i).' * (1.5 + 1e-20i))/1e-20);

%% 5. Wirtinger gradient for a real loss of complex parameters
A = (randn(8,4) + 1i*randn(8,4)) / sqrt(2);  y = (randn(8,1) + 1i*randn(8,1)) / sqrt(2);
L = @(w) real((y - A*w)' * (y - A*w));
w = (randn(4,1) + 1i*randn(4,1)) / sqrt(2);
g = -A' * (y - A*w);                      % dL/dw*  (the steepest-ascent direction is 2*g in R^8)
h = 1e-6;  gfd = zeros(4,1);
for k = 1:4
    e = zeros(4,1); e(k) = h;
    gfd(k) = (L(w+e) - L(w-e))/(2*h) + 1i*(L(w+1i*e) - L(w-1i*e))/(2*h);    % dL/dRe + i dL/dIm
end
fprintf('\nWirtinger check: max |(dL/dRe + i dL/dIm) - 2 dL/dw*| / max|.| = %.1e\n', max(abs(gfd - 2*g))/max(abs(gfd)));
mu = 0.5 / max(eig(A'*A));  wa = zeros(4,1);  wb = zeros(4,1);
for it = 1:200
    wa = wa - mu * (-A' * (y - A*wa));            % correct: step along -dL/dw*
    wb = wb - mu * conj(-A' * (y - A*wb));        % bug: step along -dL/dw = -conj(dL/dw*)
end
wls = A \ y;
fprintf('after 200 steps: L(correct) = %.4f (LS optimum %.4f); L(conjugated step) = %.4g\n', L(wa), L(wls), L(wb));

%% 6. Figure
figure('Position', [100 100 640 440]);
loglog(hs, pfd, 'b.-', hs, pcd, 'r.-', hs, pcs, 'k.-'); hold on;
loglog([1e-20 1e-1], [1 1]*eps, 'k:');
set(gca, 'XDir', 'reverse'); grid on; ylim([1e-17 10]);
xlabel('step h'); ylabel('relative error of f''(1.5)');
legend('forward difference', 'central difference', 'complex step', 'machine epsilon', 'Location', 'northwest');
title('truncation error falls with h, round-off grows as eps/h; the complex step has no subtraction');

%% m14_autofocus_sound_speed.m  --  sound speed by gradient ascent on a focus metric
% Purpose : simulate plane-wave IQ channel data (64 elements, 5 MHz) from a point
%           target in speckle in a medium with c_true = 1480 m/s, beamform at an
%           assumed speed ch, and estimate c by gradient ascent on four metrics:
%             energy : sum of |b|^2 over an ROI fixed in two-way time (coherently compounded)
%             sharp  : P*sum(I.^2)/sum(I)^2 over the same ROI
%             envz   : |b|^2 at a pixel fixed in depth (x = 0, z = 30 mm)
%             rfz    : real part of b (the RF value) at that fixed pixel
%           The derivative dM/dch is computed exactly by the chain rule (forward mode):
%           d/dch [IQ(tau) e^{j w0 tau}] = (IQ'(tau) + j w0 IQ(tau)) e^{j w0 tau} dtau/dch.
% Change  : c_true, z, angles_deg (try a single 0 deg plane wave), starts, the scene.
% Expect  : 5 angles: energy peaks at 1484-1486 m/s (bias about +5 m/s) and every start
%           in 1380-1720 converges there; rfz oscillates with period ch^2/(2 z f0)
%           (7.3 m/s here) and every start stops on its nearest local maximum (cycle
%           skipping). 1 angle: energy is nearly flat and peaks far from c_true.
% Base MATLAB only; runs in MATLAB R2019b+ and GNU Octave 8. Runtime about 10 s.
rng(1);

%% 1. Geometry, pulse and scatterers
c_true = 1480;            % true sound speed (m/s)
z      = 30e-3;           % target depth (m)
f0 = 5e6; w0 = 2*pi*f0;   % centre frequency (Hz)
sig = 0.125e-6;           % Gaussian envelope std (s): -6 dB bandwidth 3 MHz (60 %)
fs = 40e6;                % IQ sampling rate (Hz)
N = 64; pitch = 0.3e-3;   % elements, pitch (m): aperture 19.2 mm, F/1.56 at 30 mm
xe = ((0:N-1) - (N-1)/2) * pitch;                   % 1 x N element positions (m)
K = 500;                                            % speckle scatterers
sx = [0; (rand(K,1) - 0.5) * 8e-3];                 % lateral (m), first one is the point
sz = [z; z + (rand(K,1) - 0.5) * 3e-3];             % depth (m)
amp = [3; 0.25 * (randn(K,1) + 1j*randn(K,1)) / sqrt(2)];
angle_sets = {0, (-8:4:8)};                         % plane-wave angles (deg)

for set_id = 1:2
th = angle_sets{set_id} * pi/180;  A = numel(th);

%% 2. IQ channel data: IQ(t) = sum_k a_k g(t - t_k) e^{-j w0 t_k}
tk = zeros(K+1, N, A);
for a = 1:A
    tk(:,:,a) = (sz*cos(th(a)) + sx*sin(th(a)))/c_true + sqrt((sx - xe).^2 + sz.^2)/c_true;
end
t0 = min(tk(:)) - 1.5e-6;  T = ceil((max(tk(:)) + 1.5e-6 - t0) * fs);
tg = t0 + (0:T-1).' / fs;                            % T x 1 fast time (s)
IQ = zeros(T, N, A);
for a = 1:A
    for e = 1:N
        tt = tk(:, e, a).';                          % 1 x (K+1)
        IQ(:, e, a) = exp(-(tg - tt).^2 / (2*sig^2)) * (amp .* exp(-1j*w0*tt.'));
    end
end

%% 3. Beamformer with exact derivative (anonymous functions only)
% ROI in time coordinates: pixel (x, t), depth at assumed speed = ch*t/2
nx = 21; nt = 25;
[PX, PT] = meshgrid(((0:nx-1) - (nx-1)/2) * 0.15e-3, 2*z/c_true + ((0:nt-1) - (nt-1)/2) * 50e-9);
Px = PX(:); Pt = PT(:); P = numel(Px);
U   = @(tau) (tau - t0) * fs;
I0  = @(tau) min(max(floor(U(tau)), 0), T-2);
% core: q0 = IQ at the sample before tau, dq = next minus q0, fr = fraction, E = e^{j w0 tau}
core = @(q0, dq, fr, E, dt) [sum((q0 + dq.*fr) .* E, 2), sum((dq*fs + 1j*w0*(q0 + dq.*fr)) .* E .* dt, 2)];
bfi  = @(k, fr, tau, dt) core(IQ(k), IQ(k+1) - IQ(k), fr, exp(1j*w0*tau), dt);   % k = linear index
bfa  = @(tau, dt, a) bfi(I0(tau) + 1 + (0:N-1)*T + (a-1)*T*N, U(tau) - I0(tau), tau, dt);
% time-ROI delays and d tau / d ch
tauT  = @(ch, a) Pt/2*cos(th(a)) + Px*sin(th(a))/ch + sqrt((Px - xe).^2/ch^2 + (Pt/2).^2);
dtauT = @(ch, a) -Px*sin(th(a))/ch^2 - ((Px - xe).^2/ch^3) ./ sqrt((Px - xe).^2/ch^2 + (Pt/2).^2);
% fixed-depth pixel (x = 0, z)
tauZ  = @(ch, a) (z*cos(th(a)) + sqrt(xe.^2 + z^2)) / ch;
BT = @(ch) cell2mat(arrayfun(@(a) bfa(tauT(ch,a), dtauT(ch,a), a), 1:A, 'UniformOutput', false));
BZ = @(ch) cell2mat(arrayfun(@(a) bfa(tauZ(ch,a), -tauZ(ch,a)/ch, a), 1:A, 'UniformOutput', false));
bsum = @(Bm) [sum(Bm(:,1:2:end), 2), sum(Bm(:,2:2:end), 2)];       % compound angles: [b, db]
% metrics: each returns [M, dM/dch]
mE = @(bd) [sum(abs(bd(:,1)).^2), sum(2*real(conj(bd(:,1)).*bd(:,2)))];
mS = @(bd) P * [sum(abs(bd(:,1)).^4) / sum(abs(bd(:,1)).^2)^2, ...
     (sum(4*abs(bd(:,1)).^2 .* real(conj(bd(:,1)).*bd(:,2))) * sum(abs(bd(:,1)).^2) ...
      - 2*sum(abs(bd(:,1)).^4) * sum(2*real(conj(bd(:,1)).*bd(:,2)))) / sum(abs(bd(:,1)).^2)^3];
metric = {@(ch) mE(bsum(BT(ch))), @(ch) mS(bsum(BT(ch))), ...
          @(ch) mE(bsum(BZ(ch))), @(ch) real(bsum(BZ(ch)))};
names = {'energy', 'sharp', 'envz', 'rfz'};
S0 = [40 40 8 1.5];       % initial trial step of the line search (m/s)
cmin = 1350; cmax = 1750; % box constraint: every step is projected onto [cmin, cmax]

%% 4. Derivative check at 1540 m/s (exact chain rule vs central difference)
if set_id == 2
    for m = 1:4
        r = metric{m}(1540);  h = 1e-3;
        fd = (metric{m}(1540+h) - metric{m}(1540-h)) / (2*h);
        fprintf('dM/dc check %-6s  chain rule / central difference = %.6f\n', names{m}, r(2)/fd(1));
    end
end

%% 5. Sweep and gradient ascent from several starts
fprintf('\n%d plane wave(s): c_true = %d m/s, z = %d mm\n', A, c_true, round(z*1e3));
fprintf('metric | argmax (m/s) | gradient ascent: start -> end (iterations)\n');
starts = [1380 1440 1540 1620 1720];
cs_roi = 1350:5:1750;  cs_pix = 1350:1:1750;
curves = cell(1,4);  ends = zeros(4, numel(starts));
for m = 1:4
    if m <= 2, cs = cs_roi; else, cs = cs_pix; end
    M = zeros(size(cs));  for i = 1:numel(cs), r = metric{m}(cs(i)); M(i) = r(1); end
    curves{m} = [cs; M];  [~, im] = max(M);
    s_out = '';
    for j = 1:numel(starts)
        c = starts(j);  r = metric{m}(c);  n = 0;
        while n < 60
            g = r(2);  if g == 0, break; end
            s = S0(m);  ok = false;
            while s > 0.05
                cn = min(max(c + sign(g)*s, cmin), cmax);  rn = metric{m}(cn);   % projected step
                if cn ~= c && rn(1) >= r(1) + 1e-4*abs(cn-c)*abs(g), c = cn; r = rn; ok = true; break; end
                s = s/2;
            end
            if ~ok, break; end
            n = n + 1;
        end
        ends(m, j) = c;
        s_out = [s_out sprintf('%d->%.1f(%d)  ', starts(j), c, n)];
    end
    fprintf('%-6s | %7.1f      | %s\n', names{m}, cs(im), s_out);
end
fprintf('RF-at-fixed-depth period predicted ch^2/(2 z f0) = %.1f m/s at ch = c_true\n', c_true^2/(2*z*f0));
if set_id == 1, curves1 = curves; ends1 = ends; else, curves5 = curves; ends5 = ends; end
end

%% 6. Figures
figure('Position', [80 80 1100 460]);
cols = {[0 0.3 0.8], [0.8 0.2 0.5], [0.85 0.55 0], [0 0.55 0.4]};
for set_id = 1:2
    subplot(1,2,set_id);
    if set_id == 1, cv = curves1; en = ends1; ttl = '1 plane wave (0 deg)'; else, cv = curves5; en = ends5; ttl = '5 plane waves (-8:4:8 deg)'; end
    h = zeros(1,4);
    for m = 1:4
        y = cv{m}(2,:);  y = y / max(abs(y));
        h(m) = plot(cv{m}(1,:), y, 'Color', cols{m}, 'LineWidth', 1.3); hold on;
        yi = interp1(cv{m}(1,:), y, en(m,:));
        plot(en(m,:), yi, 'o', 'Color', cols{m}, 'MarkerSize', 7, 'LineWidth', 1.5);
    end
    plot([c_true c_true], [-1 1.05], 'k--');
    xlim([1350 1750]); ylim([-1 1.05]); grid on;
    xlabel('assumed sound speed c_h (m/s)'); ylabel('metric / max |metric|');
    title([ttl ': circles = gradient-ascent end points']);
    legend(h, {'energy (time ROI)', 'sharpness (time ROI)', '|b|^2 at fixed depth', 'Re b at fixed depth'}, 'Location', 'southwest');
end