% denoise_pgd1.m
% illustrate preconditioning for 2D denoising

%% noiseless original image
nx = 32; ny = 20; % image size (must be a multiple of 4 for next line)
xtrue = zeros(nx,ny);
xtrue(nx/4+[0:4],ny/4+[0:1]) = 1; xtrue(nx/2:3*nx/4,ny/2:3*ny/4) = 1;

%% ideal noiseless data (just the true image in a denoising problem)
rng(0)
%samp = rand(nx,ny) > 0.5; % sample only this fraction of the pixels
samp = true(nx,ny); % denoising
A = Gdiag(samp);
ytrue = A * xtrue;

%% add noise
rng(0)
snr = 12; % specify desired SNR (of sampled values!) in dB
sig = 10^(-snr/20) * norm(ytrue(samp)) / sqrt(sum(samp(:)));
y = ytrue + sig * randn(size(ytrue));

im plc 2 2
clim = [-0.1 1.1]; % display all images on same gray scale
im(1, xtrue, clim)
im(3, ytrue, clim)
im(4, y, clim)
xlabelf('SNR = %4.1f dB', 20*log10(norm(xtrue(:)) / norm(y(:) - xtrue(:))))
drawnow
%prompt

C = Cfd2(nx,ny); % 2D finite differences

%% run gradient descent

reg = 1.1; % todo: experiment with this to find improved SNR
step = 1 / (1 + reg * 8);
niter = 200; % max # of iter

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

delta = 0.1; % small compared to max(xtrue) - min(xtrue)
dpot = @(t) t ./ sqrt(1 + (t / delta).^2); % derivative of potential
%dpot = @(t) t; % quadratic case

x = A' * y; % initial guess is just the noisy image itself
P = 1; % no preconditioner

if 0
	e0 = zeros(nx,ny); e0(nx/2+1,ny/2+1) = 1; % unit impulse in "middle"
	psf = (A' * A + reg * C' * C) * e0;
	im(3, psf)
	filter = 1 ./ reale(fft2(fftshift(psf)));
	im(2, fftshift(filter))
	prompt

	forw = @(arg, x) ifft2(fft2(x) .* filter);
	P = fatrix2('forw', forw, 'idim', [nx ny], 'odim', [nx ny]);
	step = 1; % todo: why?
end

% PGD
for iter=1:niter
	ticker(mfilename, iter, niter)
	grad = A' * (A * x - y) + reg * (C' * dpot(C * x)); % gradient
	x = x - step * (P * grad); % PGD update
	if (norm(grad(:)) < tol_grad)
		printm('stopped at %d iterations', iter)
		break
	end
end
im(2, x, clim)
xlabelf('SNR = %4.1f dB', 20*log10(norm(xtrue(:)) / norm(x(:) - xtrue(:))))
