% Script to perform super-resolution of point sources that share a common
% support. More specifically, we aim to estimate the common support of m
% signals X_1, ..., X_m from measurements of the form Y_k = F X_k where F 
% is the low-pass operator that maps a function in the unit interval to its 
% 2fc+1 lower Fourier series coefficients. Recovery is carried out by 
% solving
% 
% max_C sum_k Re(Y_k'C_k) 
% subject to max_t sum_k |(F*C_k)_j|^2 <= 1,
% 
% using an equivalent SDP formulation. 
% The script uses CVX (http://cvxr.com/cvx/).
%
% For more information see the paper "Super-resolution of point sources via 
% convex programming" by C. Fernandez-Granda. 

clear all
close all

rand('state',1234) % random seed 
randn('state',1234)

m=3;
nspikes=15;
% cutoff frequency
fc = 40;
distance = 0.7/fc;
jitter_ratio=0.75;
jitter_factor=10;
jitter = distance/jitter_factor;
linespectra_aux = ( (distance/2) : distance : (nspikes+1)*distance )';
offsets = cumsum(jitter.*rand(nspikes,1).*(rand(nspikes,1)<jitter_ratio));
% spike locations
tspikes = 0.4 + linespectra_aux(1:nspikes) + offsets;
% amplitudes 
X=randn(nspikes,m);

% data 
k = -fc:1:fc;
F = exp(-1i*2*pi*k'*tspikes'); % Fourier matrix
Y = F*X; 
n = length(Y);

% Solve SDP
cvx_precision best
cvx_solver sdpt3
cvx_begin sdp 
    variable M(n,n) hermitian;
    variable C(n,m) complex;
    %M >= B*B';
    [M C; C' eye(m)]>=0
    trace(M) == 1;
    for j = 1:n-1,
        sum(diag(M,j)) == 0;
    end
    maximize(real(trace(C'*Y)))
cvx_end

dualpoly_mm = @(t)  sum( conj(exp(1i*2*pi*t*k)*C).* (exp(1i*2*pi*t*k)*C) ,2);

for ind_col =1:m

    cvx_precision best
    cvx_solver sdpt3
    cvx_begin sdp 
        variable Q(n+1,n+1) hermitian;
        Q >= 0;
        Q(n+1,n+1) == 1;
        trace(Q) == 2;
        for j = 1:n-1,
            sum(diag(Q,j)) == Q(n+1-j,n+1);
        end
        maximize(real(Q(1:n,n+1)'*Y(:,ind_col)))
    cvx_end
    u = Q(1:n,n+1);
    dualpoly{ind_col} = @(t)real(exp(1i*2*pi*t*k) * u);
    
end

t = ([0.35:1e-4:0.725])';
F_adjoint = @(u)  exp(1i*2*pi*t*k) * u;

x_lowpass_1 = real(F_adjoint(Y(:,1)./(n)));
x_lowpass_2 = real(F_adjoint(Y(:,2)./(n)));
x_lowpass_3 = real(F_adjoint(Y(:,3)./(n)));

figure
plot(t,x_lowpass_1,'b')
hold on
stem(tspikes, X(:, 1), 'r')
legend('Data', 'Signal')
title('Signal 1')
figure
plot(t,x_lowpass_2,'b')
hold on
stem(tspikes, X(:, 2), 'r')
legend('Data', 'Signal')
title('Signal 2')
figure
plot(t,x_lowpass_3,'b')
hold on
stem(tspikes, X(:, 3), 'r')
legend('Data', 'Signal')
title('Signal 3')

figure, 
plot(t,abs(dualpoly_mm(t)))
hold on,
stem(tspikes, ones(1,nspikes),'r','Marker','none','linewidth',1)
title('Support-locating polynomial (gTV on all the data)')


figure, 
plot(t,real(dualpoly{1}(t)),'k')
hold on,
stem(tspikes, sign(X(:,1)),'r','Marker','none','linewidth',1)
set(gca,'xtick',[])
set(gca,'xticklabel',[])
title('Support-locating polynomial (TV on signal 1)')
figure, 
plot(t,real(dualpoly{2}(t)),'g')
hold on,
stem(tspikes, sign(X(:,2)),'r','Marker','none','linewidth',1)
set(gca,'xtick',[])
set(gca,'xticklabel',[])
title('Support-locating polynomial (TV on signal 2)')
figure, 
plot(t,real(dualpoly{3}(t)),'c')
hold on,
stem(tspikes, sign(X(:,3)),'r','Marker','none','linewidth',1)
set(gca,'xtick',[])
set(gca,'xticklabel',[])
title('Support-locating polynomial (TV on signal 3)')

tol=1e-4;
aux_u=0;
for ind=1:m
    aux_u =aux_u - conv(C(:,ind),flipud(conj(C(:,ind))));
end
aux_u(2*fc+1)=1+aux_u(2*fc+1);
k2 = (-2*fc):(2*fc);
roots_pol = roots(flipud(aux_u));

% Isolate roots on the unit circle
roots_detected = roots_pol(abs(1-abs(roots_pol)) < tol);
[auxsort,ind_sort]=sort(real(roots_detected));
roots_detected = roots_detected(ind_sort);
% Roots are double so take 1 and out of 2 and compute argument
t_rec = angle(roots_detected(1:2:end))/2/pi;
% Argument is between -1/2 and 1/2 so convert angles
t_rec(t_rec < 0)= t_rec(t_rec < 0) + 1;  
t_rec = sort(t_rec)
% Accuracy 
delta = abs(t_rec - tspikes);

F_est = exp(-1i*2*pi*k'*t_rec'); % estimated Fourier matrix
X_est = real(F_est\Y);

figure, 
plot(tspikes, X(:,1), 'kx')
hold on,
plot(tspikes, X(:,2), 'bx')
plot(tspikes, X(:,3), 'rx')
plot(t_rec, X_est(:,1), 'ko')
plot(t_rec, X_est(:,2), 'bo')
plot(t_rec, X_est(:,3), 'ro')
legend('Signal 1', 'Signal 2', 'Signal 3', 'Estimate 1', 'Estimate 2', ... 
       'Estimate 3')