%%%%%%%%%%%%%%%%%% Demo for fig. 2 
%%% "Nonconvex Robust Low-rank Matrix Recovery,"
%%%  by X. Li, Z. Zhu, A.M.-C. So, and R. Vidal

clear all; close all;
rand('seed',2018);randn('seed',2018);
%% setup
n=50;  
r = 5;
m = 5*n*r;
outlier_ratio = 0.3;

Uo=randn(n,r);
X=Uo*Uo';
%sensing
y=zeros(m,1); Acell=cell(m,1);
for it=1:m
    Acell{it}=randn(n,n);
    y(it)=trace(X'*Acell{it});
end
%add outliers
s = zeros(m,1);
indx = randperm(m);
s(indx(1:floor(outlier_ratio*m))) = 10*randn(floor(outlier_ratio*m),1);
%observation
y = y + s;
Aop=zeros(m,n^2);
for it=1:m
    Aop(it,:)=vec(Acell{it})';
end
%% geometrically diminishing stepsizes
maxiter = 3e3;
%random initialization
U_0 = randn(n,r); 
mu_0 = 1; 

%%% choice 1 for rho
rho = 0.9;
U = U_0;
dist1 = [];dist1(1) = norm(U*U' - X,'fro');
for i = 2:maxiter+1
    dev=Aop*vec(U*U')-y;
    gradU = (A_adj(sign(dev),Acell,m)*U) + (A_adj(sign(dev),Acell,m)'*U);
    mu = (mu_0*rho^i)/m;
    U = U - mu*gradU;
    [P S Q] = svd(U'*Uo);
    dist1(i) = norm(U - Uo*Q*P','fro');
end

%%% choice 2 for rho
rho = 0.93;
U = U_0;
dist2 = [];dist2(1) = norm(U*U' - X,'fro');
for i = 2:maxiter+1
    dev=Aop*vec(U*U')-y;
    gradU = (A_adj(sign(dev),Acell,m)*U) + (A_adj(sign(dev),Acell,m)'*U);
    mu = (mu_0*rho^i)/m;
    U = U - mu*gradU;
    [P S Q] = svd(U'*Uo);
    dist2(i) = norm(U - Uo*Q*P','fro');
end
  
%%% choice 2 for rho
rho = 0.96;
U = U_0;
dist3 = [];dist3(1) = norm(U*U' - X,'fro');
for i = 2:maxiter+1
    dev=Aop*vec(U*U')-y;
    gradU = (A_adj(sign(dev),Acell,m)*U) + (A_adj(sign(dev),Acell,m)'*U);
    mu = (mu_0*rho^i)/m;
    U = U - mu*gradU;
    [P S Q] = svd(U'*Uo);
    dist3(i) = norm(U - Uo*Q*P','fro');
end

  
%%% choice 4 for rho
rho = 0.99;
U = U_0;
dist4 = [];dist4(1) = norm(U*U' - X,'fro');
for i = 2:maxiter+1
    dev=Aop*vec(U*U')-y;
    gradU = (A_adj(sign(dev),Acell,m)*U) + (A_adj(sign(dev),Acell,m)'*U);
    mu = (mu_0*rho^i)/m;
    U = U - mu*gradU;
    [P S Q] = svd(U'*Uo);
    dist4(i) = norm(U - Uo*Q*P','fro');
end

%% plot 
figure
semilogy(dist1,'k:','linewidth',2);hold on
semilogy(dist2,'r--','linewidth',2);semilogy(dist3,'b-.','linewidth',2);semilogy(dist4,'g-','linewidth',2);
legend('\rho = 0.9','\rho = 0.93','\rho = 0.96','\rho = 0.99','location','best')
xlabel('iteration','FontSize',18,'FontName','Times New Roman');
set(gca, 'LineWidth' , 1.8,'FontSize',20);
set(gcf, 'Color', 'w');
xlabel('Iteration','FontSize',22);
ylabel('dist$({\bf U}_k, {\cal U})$','FontSize',22,'FontName', 'Times New Roman','Interpreter','LaTex');


%% function used
function X=A_adj(z,Acell,m)
[n1,n2] = size(Acell{1});
X=zeros(n1,n2);
for itt=1:m
    X=X+z(itt)*Acell{itt};
end 
end

