% Demonstration of the interest of l1 regularization compared to linear
% solutions in the case of an ill-posed problem with a good a priori on the
% original data.
%
% The considered linear problem is 1D magnetic resonance imaging of a
% piecewise constant signal. The measurements correspond to samples in the
% continuous Fourier domain. The reconstructed data is parametrized by its
% coefficients in the basis spanned by an adapted generating function and
% its shifts. Noise can be added on the measurements.
%
% The problem is underdetermined (and consequently ill-posed) since there
% is more coefficients to restore than measurements. The consequence is
% that several sets of coefficients can lead to the observed measurements.
% By determining the null space of the forward model, we generate signals
%that are different from the original data but lead to the same
% measurements. They are admissible solutions of the inverse problem.
%
% To make the solution unique, we introduce regularization functionals
% that have to be minimized. This leads to pseudo-inverse, Tikhonov and
% Total Variation (TV) solutions.
%
% By considering the signal-to-error ratio of the reconstructions and
% looking at the extrapolated k-space samples, we demonstrate that the
% adapted l1 regularization constraints (TV in the case of piecewise
% constant signals) allow a better reconstruction, where the missing
% samples are partly recovered (almost perfectly when there is no noise).
%
% Matthieu Guerquin-Kern, Biomedical Imaging Group - EPFL, 13-04-2009 (DD-MM-YYYY)
% updated 15-06-2009

clear all;close all;clc;

noisy = true; % TRUE for noisy measurements and FALSE for noiseless
SNR = 15;                           % noise on the measurements (SNR in dB)

saveplot = true;

%% generating the spatial dimensions and signals

% continuous
dr = 1e-2;                          % discretization of the continuous space (the finer, the better)
FOV = 30;                           % Field of view
r = -FOV/2:dr:FOV/2-dr;             % spatial axis

rect = @(r) double((abs(r)<1/2)|(r==1/2));

% parameters of the original data made of shifted and dilated rect
% functions (simulates a line profile of the Shepp-Logan brain phantom)
L = 2*FOV/3+1;                      % width of the signal
a = [0.2 -0.2 -0.2 0.8 0.8];        % weights
b = [L 5 3 2 2];                    % dilation
c = [0 -3 +4 +9.5 -9.5];            % shifts
f = @(r) a(:).'*rect((repmat(r,length(c),1)+repmat(c(:),1,length(r)))./repmat(b(:),1,length(r)));

% discrete
Delta = 1;                          % sampling step for the discrete reprensentation
n = -floor(FOV/2/Delta):1:ceil(FOV/2/Delta);    % discrete spatial axis

c_0 = f(n)';                         % coefficients of the discretized original data (only works if phi is interpolating!)

phi = @(r) rect(r/Delta);           % generating function in space domain

%% generating the frequency dimensions and signals
W = 10*pi/4;
dw = W/length(r);
w = -W/2:dw:W/2-dw;                 % frequency axis

R = 1.5;                            % undersampling factor
dk = 1/FOV;                         % Nyquist sampling
k = -1/Delta/2/R:dk:1/Delta/2/R;    % k-space trajectory
k_miss = [-1/Delta/2:dk:-1/Delta/2/R-dk, 1/Delta/2/R+dk:dk:1/Delta/2]; % missing k-space samples
undersampling = length(n)/length(k)

f_hat = @(w) (a(:).*b(:)).'*(exp(1j*repmat(w,length(c),1).*repmat(c(:),1,length(w))).*sinc((repmat(w,length(c),1).*repmat(b(:),1,length(w)))/2/pi));

phi_hat = @(w) sinc(w*Delta/2/pi)*Delta;    % generating function in Fourier domain

m = f_hat(-2*pi*k);

% for noisy simulations
if noisy
    randn('seed',0);
    noise = randn(size(m))+1j*randn(size(m));
    noise = norm(m)*10^(-SNR/20)*noise/norm(noise);
    m = m + noise;
end

m_miss = f_hat(-2*pi*k_miss);

%% Estimate
% continuous function represented by a linear combination of shifted
% versions of the generating function

estimate = @(c,r) c(:)'*phi(repmat(r,length(n),1)-repmat(n(:),1,length(r))*Delta);
estimate_hat = @(c,w) phi_hat(w).*(c(:)'*exp(-1j*n(:)*Delta*w(:)'));

%% Linear forward model

% construction of the forward matrix E such that m = E*c_0
E = zeros(length(k),length(n));
for i = 1:length(n)
    vec = zeros(size(c_0));
    vec(i) = 1;
    E(:,i) = estimate_hat(vec,-2*pi*k);
end

%% Random solution that is real

c_random = pinv(E)*m(:);

% nullspace
Z = null(E);
Z_tmp = pinv(imag(Z))*real(Z);
ZR = real(-Z*Z_tmp + 1j*Z);%ZR = ZR/norm(ZR,2);
ZI = 1j*real(Z + 1j*Z*Z_tmp);%ZI = ZI/norm(ZI,2);

cz = ZR*randn(size(Z,2),1);
cz = cz/norm(cz,2); % element of the nullspace of E (arbitrarily chosen real)

c_random = c_random + 1*cz;

norm(E*c_random(:)-m(:))/norm(m(:)) % the measurements and E*c_random should be the same

SNR_random = -20*log10(norm(c_random(:)-c_0(:),2)/norm(c_0(:),2));

%% pseudo inverse

c_pseudo = pinv(E)*m(:);

SNR_pseudo = -20*log10(norm(c_pseudo(:)-c_0(:),2)/norm(c_0(:),2));

%% Regularized linear inverse (Tikhonov)

ca = E'*m(:);

if noisy
    lambda = 7e-1;
else
    lambda = 1e0;
end

%M = @(x) (E'*E)*x+lambda*(R'*R)*x;
M = @(x) (E'*E)*x+lambda*x;

c_tikhonov = pcg(M,ca,1e-8,30);

SNR_tikhonov = -20*log10(norm(c_tikhonov(:)-c_0(:),2)/norm(c_0(:),2));

%% TV (Iteratively reweighted Least Squares Algorithm)

if noisy
    lambda = 1e0;
else
    lambda = 1e-4;
end
epsilon = 1e-12;

FD = eye(length(c_0));
FD = FD-circshift(FD,-1); % finite differences matrix

c_TV = c_tikhonov;
for it = 1:20
    D = diag(1./sqrt(abs(FD*c_TV).^2+eps));
    M = @(x) (E'*E)*x+lambda*(FD'*D*FD)*x;
    c_TV = pcg(M,ca,1e-8,300,[],[],c_TV);
    figure(99);plot(r,f(r),'g-','LineWidth',2);hold on;plot(r,estimate(abs(c_TV-c_0),r),'r-');
    legend('original data','current error (modulus)','Location','North');title(['TV (IRLS algo): iteration ' num2str(it)]);hold off;drawnow;
end
close(99);

SNR_TV = -20*log10(norm(c_TV(:)-c_0(:),2)/norm(c_0(:),2));

%% Plots in Space

m1 = min([min(abs(c_0)),min(abs(c_0-c_random)),min(abs(c_0-c_pseudo)),min(abs(c_0-c_tikhonov)),min(abs(c_0-c_TV))])-0.1;
M1 = max([max(abs(c_0)),max(abs(c_0-c_random)),max(abs(c_0-c_pseudo)),max(abs(c_0-c_tikhonov)),max(abs(c_0-c_TV))])+0.1;

figure(1);
plot(r,f(r),'g-','LineWidth',2);xlabel('space domain');ylabel('Intensity');
hold on;plot(r,phi(r),'c-','LineWidth',1);
stem(n,c_0,'b-','LineWidth',1);
legend('original data','generating function','coefficients of discretization','Location','North');
xlabel('space domain');ylabel('Intensity');hold off;

figure;
plot(r,f(r),'g-','LineWidth',2);
hold on;plot(r,estimate(abs(c_0-c_random),r),'k-.','LineWidth',2);hold off;
xlabel('space domain');ylabel('Intensity');axis([n(1) n(end) m1 M1]);grid on;
title(['Random solution: SNR= ' sprintf('%.1f',SNR_random) ' dB']);
if saveplot
    print(gcf,'-depsc2',['random_space_SNR_' sprintf('%d',round(SNR_random)) 'dB.eps']);
end

figure;
plot(r,f(r),'g-','LineWidth',2);
hold on;plot(r,estimate(abs(c_0-c_pseudo),r),'m-.','LineWidth',2);hold off;
xlabel('space domain');ylabel('Intensity');axis([n(1) n(end) m1 M1]);grid on;
title(['Pseudoinverse solution: SNR= ' sprintf('%.1f',SNR_pseudo) ' dB']);
if saveplot
    print(gcf,'-depsc2',['pseudo_space_SNR_' sprintf('%d',round(SNR_pseudo)) 'dB.eps']);
end

figure;
plot(r,f(r),'g-','LineWidth',2);
hold on;plot(r,estimate(abs(c_0-c_tikhonov),r),'c--','LineWidth',2);hold off;
xlabel('space domain');ylabel('Intensity');axis([n(1) n(end) m1 M1]);grid on;
title(['Tikhonov solution: SNR= ' sprintf('%.1f',SNR_tikhonov) 'dB']);
if saveplot
    print(gcf,'-depsc2',['tikhonov_space_SNR_' sprintf('%d',round(SNR_tikhonov)) 'dB.eps']);
end

figure;
plot(r,f(r),'g-','LineWidth',2);
hold on;plot(r,estimate(abs(c_0-c_TV),r),'y-.','LineWidth',2);hold off;
xlabel('space domain');ylabel('Intensity');axis([n(1) n(end) m1 M1]);grid on;
title(['TV solution: SNR= ' sprintf('%.1f',SNR_TV) ' dB']);
if saveplot
    print(gcf,'-depsc2',['TV_space_SNR_' sprintf('%d',round(SNR_TV)) 'dB.eps']);
end

%% Plots in Fourier

m2 = min([min(abs(f_hat(w))),min(abs(estimate_hat(c_random,w))),min(abs(estimate_hat(c_pseudo,w))),min(abs(estimate_hat(c_tikhonov,w))),min(abs(estimate_hat(c_TV,w)))]);
M2 = max([max(abs(f_hat(w))),max(abs(estimate_hat(c_random,w))),max(abs(estimate_hat(c_pseudo,w))),max(abs(estimate_hat(c_tikhonov,w))),max(abs(estimate_hat(c_TV,w)))]);

figure;
semilogy(w,abs(f_hat(w)),'g-','LineWidth',2);
hold on;semilogy(-2*pi*k,abs(m),'bo','MarkerSize',7,'LineWidth',2);
semilogy(-2*pi*k_miss,abs(f_hat(-2*pi*k_miss)),'rx','MarkerSize',7,'LineWidth',2);
semilogy(w,abs(estimate_hat(c_random,w)),'k-.','LineWidth',2);
xlabel('Fourier domain');ylabel('Modulus (log scale)');hold off;
grid on;set( gca, 'XTick', -pi:pi/2:pi,'XTickLabel', '-pi|-pi/2|0|pi/2|pi');
axis([w(1) w(end) m2 M2]);
title(['Random solution: SNR= ' sprintf('%.1f',SNR_random) ' dB']);
if saveplot
    print(gcf,'-depsc2',['random_fourier_SNR_' sprintf('%d',round(SNR_random)) 'dB.eps']);
end

figure;
semilogy(w,abs(f_hat(w)),'g-','LineWidth',2);
hold on;semilogy(-2*pi*k,abs(m),'bo','MarkerSize',7,'LineWidth',2);
semilogy(-2*pi*k_miss,abs(f_hat(-2*pi*k_miss)),'rx','MarkerSize',7,'LineWidth',2);
semilogy(w,abs(estimate_hat(c_pseudo,w)),'m-.','LineWidth',2);
xlabel('Fourier domain');ylabel('Modulus (log scale)');hold off;
grid on;set( gca, 'XTick', -pi:pi/2:pi,'XTickLabel', '-pi|-pi/2|0|pi/2|pi');
axis([w(1) w(end) m2 M2]);
title(['Pseudoinverse solution: SNR= ' sprintf('%.1f',SNR_pseudo) ' dB']);
if saveplot
    print(gcf,'-depsc2',['pseudo_fourier_SNR_' sprintf('%d',round(SNR_pseudo)) 'dB.eps']);
end

figure;
semilogy(w,abs(f_hat(w)),'g-','LineWidth',2);
hold on;semilogy(-2*pi*k,abs(m),'bo','MarkerSize',7,'LineWidth',2);
semilogy(-2*pi*k_miss,abs(f_hat(-2*pi*k_miss)),'rx','MarkerSize',7,'LineWidth',2);
semilogy(w,abs(estimate_hat(c_tikhonov,w)),'c--','LineWidth',2);
xlabel('Fourier domain');ylabel('Modulus (log scale)');hold off;
grid on;set( gca, 'XTick', -pi:pi/2:pi,'XTickLabel', '-pi|-pi/2|0|pi/2|pi');
axis([w(1) w(end) m2 M2]);
title(['Tikhonov solution: SNR= ' sprintf('%.1f',SNR_tikhonov) ' dB']);
if saveplot
    print(gcf,'-depsc2',['tikhonov_fourier_SNR_' sprintf('%d',round(SNR_tikhonov)) 'dB.eps']);
end

figure;semilogy(w,abs(f_hat(w)),'g-','LineWidth',2);
hold on;semilogy(-2*pi*k,abs(m),'bo','MarkerSize',7,'LineWidth',2);
semilogy(-2*pi*k_miss,abs(f_hat(-2*pi*k_miss)),'rx','MarkerSize',7,'LineWidth',2);
semilogy(w,abs(estimate_hat(c_TV,w)),'y-.','LineWidth',2);
xlabel('Fourier domain');ylabel('Modulus (log scale)');hold off;
grid on;set( gca, 'XTick', -pi:pi/2:pi,'XTickLabel', '-pi|-pi/2|0|pi/2|pi');
axis([w(1) w(end) m2 M2]);
title(['TV solution: SNR= ' sprintf('%.1f',SNR_TV) ' dB']);if saveplot
    print(gcf,'-depsc2',['TV_fourier_SNR_' sprintf('%d',round(SNR_TV)) 'dB.eps']);
end