% Script to perform demixing of spikes and sines via convex programming. 
% More specifically, we aim to estimate the line spectra x and the spikes s
% from measurements of the form y=Fx + s 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 Re(y'c) 
% subject to max_t |(F*c)(t) | <= 1, max_j |c_j| <= gamma
% 
% for a regularization parameter gamma >0, 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

% Problem data 

% cutoff frequency
fc = 50;
nspikes = 5;
ncorruptions=5;
tol =1e-4;
% spike locations
distance=1.4./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));
linespectra = distance + linespectra_aux(1:nspikes) + offsets;
% For convenience, signal is real in the time domain
linespectra = [-flipud(linespectra); linespectra];
% line spectra amplitudes
x = [-0.9 -0.3 1.2 0.1 -0.7]'; % amplitudes can also be complex
% For convenience, signal is real in the time domain
x = [flipud(x); x]; 

% data 
t = linspace(0,0.184,2e3);
k = -fc:1:fc;
k_fine = linspace(0,2*fc,2e3);
F = exp(-1i*2*pi*k'*linespectra'); % Fourier matrix
F_adjoint = @(u)  exp(1i*2*pi*t'*k) * u;
F_fine = exp(-1i*2*pi*k_fine'*linespectra'); % Fourier matrix
n = 2*fc+1;
c= zeros(n,1);
aux_corr = randperm(fc);
ind_corr_orig=aux_corr(1:ncorruptions)+9;
c(fc+1+ ind_corr_orig )= -randn(ncorruptions,1).*(norm(F*x)./ncorruptions)*1.7;
c(fc+1- ind_corr_orig )= c(fc+1+ ind_corr_orig );
clean_ind = setdiff(1:n,ind_corr_orig);
y_clean = real(F*x);
y_fine = real(F_fine*x);
y = y_clean+c;

figure
plot(k_fine,y_fine)
hold on
stem(k((fc+1):end),y_clean((fc+1):end))
legend('Sines', 'Samples')
title('Sines')
figure
stem(k((fc+1):end),real(c((fc+1):end)))
title('Spikes')
figure
stem(k((fc+1):end),real(y((fc+1):end)))
title('Data')
figure
plot(t,real(F_adjoint(y)./n))
hold on
stem(linespectra((nspikes+1):end),x((nspikes+1):end),'r')
title('Spectrum of the data')
gamma=0.15;


% Solve SDP
cvx_precision best
cvx_solver sdpt3
cvx_begin sdp 
    variable X(n+1,n+1) hermitian;
    X >= 0;
    X(n+1,n+1) == 1;
    trace(X) == 2;
    for j = 1:n-1,
        sum(diag(X,j)) == X(n+1-j,n+1);
    end
    abs(X(1:n,n+1))<=gamma;
    maximize(real(X(1:n,n+1)'*y))
cvx_end

u = X(1:n,n+1);
ind_corr = find(abs((abs(u))-gamma)< tol );
clean_measurements = setdiff(1:n,ind_corr);
figure
plot(real(u),'.b')
hold on
plot(1:length(u),gamma*ones(length(u),1),'--r')
plot(1:length(u),-gamma*ones(length(u),1),'--r')
title('Support locating polynomial (coefficients)')

% Plot result

% k = 0:(2*fc);
dualpoly = @(t)  exp(1i*2*pi*t'*k) * u;

figure, plot(t,real(dualpoly(t)))
hold on,
plot(linespectra(nspikes+1:end), sign(x(nspikes+1:end)),'r*')
title('Support locating polynomial')
% Roots of dual polynomial

aux_u =- conv(u,flipud(conj(u)));
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 = sort(angle(roots_detected(1:2:end))/2/pi);
F_est = exp(-1i*2*pi*k(clean_measurements)'*t_rec'); % estimated Fourier matrix
x_est = F_est\y(clean_measurements); % estimated amplitudes
F_corr = exp(-1i*2*pi*k(ind_corr)'*t_rec'); 
corr_est =zeros(n,1);
corr_est(ind_corr) = y(ind_corr)-F_corr*x_est;

figure
plot(linespectra(nspikes+1:end), x(nspikes+1:end),'k*')
hold on,
stem(t_rec(nspikes+1:end),x_est(nspikes+1:end),'r')
hold off;
legend('Original signal','Estimated signal','Location','NW')
title('Line spectra of the sines')

figure
stem(k((fc+1):end),c((fc+1):end),'ob')
hold on
stem(k((fc+1):end),real(corr_est((fc+1):end)),'xr')
legend('Original signal','Estimated signal','Location','NW')
title('Spikes')
