% deblur2_nesterov.m
% compare sqs and nesterov for deblurring with nonnegativity constraint

%% read predefined data
if ~isvar('y')
	load deblur2_data
	[nx ny] = size(y);

	%% system matrix for blur, with non-periodic edge conditions
	A = Gblur(true([nx ny]), 'psf', psf);

	C = Cfd2_abs(nx,ny); % 2D finite differences with "abs" support
end


figure(1), im plc 2 3
clim = [-10 130]; % display all images on same gray scale
im(1, xtrue, clim)
im(2, y, clim)
im(3, pad_into_center(psf, 8), 'psf')
drawnow


%% image deblurring parameters

reg = 2^-1;
niter = 200; % max # of iter

grad0 = A' * (A * (0*xtrue) - y); % gradient at "0" image to normalize tolerance
tol_grad = 8e-4 * norm(grad0(:)); % stop tolerance for gradient

delta = 2; % small compared to max(xtrue) - min(xtrue).
pot = @(t) delta^2 * (sqrt(1 + (t / delta).^2) - 1); % hyperbola potential
dpot = @(t) t ./ sqrt(1 + (t / delta).^2); % derivative of hyperbola potential
wpot = @(t) 1 ./ sqrt(1 + (t / delta).^2); % Huber weighting function

cost = @(x) 0.5 * norm(col(y - A * x)).^2 + reg * sum(col(pot(C * x)));
snr = @(x) 20*log10(norm(xtrue(:)) / norm(x(:) - xtrue(:)));

if 1 % wiener filter
	e0 = zeros(nx,ny); e0(nx/2+1,ny/2+1) = 1; % unit impulse in "middle"
%	Hpsf = (A' * A + reg * C' * C) * e0;
	Hpsf = (A' * A) * e0 + reg * (C' * C) * e0;
	filter = 1 ./ reale(fft2(fftshift(Hpsf)));
	forw = @(arg, x) reale(ifft2(fft2(x) .* filter));
	P = fatrix2('forw', forw, 'idim', [nx ny], 'odim', [nx ny]);
	x0 = max(P * (A' * y), 0); % wiener filter on noisy image
else
	x0 = 1 * y; % just noisy image
end
im(3, x0, clim)
xlabelf('SNR = %4.1f dB', snr(x0))
drawnow

denomL = abs(A)' * abs(A) * ones(nx,ny);

%% gradient projection (GP) with *fixed* diagonal precon
if ~isvar('cost_gp')
x = x0;
cost_gp = nan(niter, 1);
Ca1 = abs(C) * ones(nx,ny);
denomR = abs(C)' * Ca1;
D = Gdiag(1 ./ (denomL + reg * denomR));
for iter=1:niter
	ticker(mfilename, iter, niter)
	cost_gp(iter) = cost(x);
	grad = A' * (A * x - y) + reg * (C' * dpot(C * x)); % gradient
	x = x - D * grad; % step=1 because D is a majorizer
	x = max(x, 0); % nonnegativity constraint
	if (norm(grad(:) .* (x(:) > 0)) < tol_grad)
		printm('GP stopped at %d iterations', iter)
		break
	end
end
im(4, x, 'GP', clim)
xlabelf('SNR = %4.1f dB', snr(x))
ylabelf('%d iter', iter)
end


%% SQS, where "diagonal preconditioner" changes each iteration
if ~isvar('cost_sqs')
x = x0;
cost_sqs = nan(niter, 1);
for iter=1:niter
	ticker(mfilename, iter, niter)
	cost_sqs(iter) = cost(x);
	denomR = abs(C)' * (wpot(C * x) .* Ca1); % Huber curvatures for SQS
	D = Gdiag(1 ./ (denomL + reg * denomR));
	grad = A' * (A * x - y) + reg * (C' * dpot(C * x)); % gradient
	x = x - D * grad; % SQS update
	x = max(x, 0); % nonnegativity constraint
	if (norm(grad(:) .* (x(:) > 0)) < tol_grad)
		printm('SQS stopped at %d iterations', iter)
		break
	end
end
im(5, x, 'SQS', clim)
xlabelf('SNR = %4.1f dB', snr(x))
ylabelf('%d iter', iter)
end


%% SQS+Nesterov
if ~isvar('cost_mom')
x = x0;
z = x;
t = 1;
cost_mom = nan(niter, 1);
for iter=1:niter
	ticker(mfilename, iter, niter)
	cost_mom(iter) = cost(x);
	denomR = abs(C)' * (wpot(C * z) .* Ca1); % Huber curvatures for SQS
	D = Gdiag(1 ./ (denomL + reg * denomR));
	grad = A' * (A * z - y) + reg * (C' * dpot(C * z)); % gradient (in z)
	xnew = z - D * grad; % SQS update (in z)
	xnew = max(xnew, 0); % nonnegativity constraint
	tnew = (1 + sqrt(1 + 4 * t^2)) / 2;
	z = xnew + (t-1)/tnew * (xnew - x);
	t = tnew;
	x = xnew;
	if (norm(grad(:) .* (x(:) > 0)) < tol_grad)
		printm('SQS+Momentum stopped at %d iterations', iter)
		break
	end
end
im(6, x, 'SQS+Nesterov', clim)
xlabelf('SNR = %4.1f dB', snr(x))
ylabelf('%d iter', iter)
end


figure(2)
cmin = min([min(cost_gp) min(cost_sqs) min(cost_mom)]);
cpl = @(c) (c - cmin) / (cost_gp(1) - cmin);
plot(1:niter, cpl(cost_gp), '-o', ...
	1:niter, cpl(cost_sqs), '-s', ...
	1:niter, cpl(cost_mom), '-+')
legend('GP', 'SQS', 'SQS+Nesterov')
xlabelf iteration
ylabelf '(relative) cost'
axis([0 100 0 1])
ytick([0 1])
