%% NMF by APGD
% Andersen Ang @ MARO, FPM, UMONS, Belgium
% Last update 2018-March-14
%%
clear,close all,clc
% load an data, for example the rice image
X = double(imread('rice.png'));
% get the size of X
[m n] = size(X);
% set the rank r
r = 64;
%% Initialization
  % Initialize W and H
  Wini = rand(m,r); W = Wini; W2 = W;
  Hini = rand(r,n); H = Hini; H2 = H;
  % Set the epsilon 
  epsilon = 0;
  epsW = epsilon*ones(size(W));
  epsH = epsilon*ones(size(H));
  % inital error value
  err0 = norm(X-W*H,'fro')^2/norm(X,'fro')^2*100;
  % max iteration
  maxiter = 1500;
  errval = zeros(1,maxiter);
  errval2 = errval;
%% The main loop
ii = 1;
while ii < maxiter
  % On varaible W
      % gradient 
      gradW = 2*(W*H-X)*H';
      % step size 
      stepW = 1/(2*norm(H*H','fro')); 
      % PGD step
      W = max(W - stepW * gradW, epsW );
      % compute error value
      errval(ii) = norm(X-W*H,'fro')^2/norm(X,'fro')^2*100;
      ii = ii + 1;

  % On variable H
      % gradient 
      gradH = 2*W'*(W*H-X);
      % step size
      stepH = 1/(2*norm(W'*W,'fro'));
      % PGD step
      H = max(H - stepH * gradH, epsH);
      % compute error value
      errval(ii) = norm(X-W*H,'fro')^2/norm(X,'fro')^2*100;  
      ii = ii + 1;
end
%% APGD
ii = 1;
maxiterinner = 15;
while ii < maxiter
    %On W
    V = W2;
    lambda = 0;  
    for i = 1 : maxiterinner % or other stopping conditions
        lambda_ = lambda;
        lambda = 0.5*(  1 + sqrt(1 + 4*lambda^2) );
        gamma = (1-lambda_) / lambda;
        V_ = V;
          % gradient 
          gradW2 = 2*(W2*H2-X)*H2';
          % step size 
          stepW2 = 1/(2*norm(H2*H2')); 
          % PGD step
          V = max(W2 - stepW2 * gradW2, epsW );
          % Extrpolation
          W2 = (1-gamma)*V + gamma*V_;
          % compute error value
          errval2(ii) = norm(X-W2*H2,'fro')^2/norm(X,'fro')^2*100;
          ii = ii + 1;
    end;        
    % On H
    G = H2;
    lambda = 0;
    for i = 1 : maxiterinner % or other stopping conditions
        lambda_ = lambda;
        lambda = 0.5*(  1 + sqrt(1 + 4*lambda^2) );
        gamma = (1-lambda_) / lambda;
        G_ = G;
          % gradient 
          gradH2 = 2*W2'*(W2*H2-X);
          % step size
          stepH2 = 1/(2*norm(W2'*W2));
          % PGD step
          G = max(H2 - stepH2 * gradH2, epsH);
          % Extrapolation
          H2 = (1-gamma)*G + gamma*G_;
          % compute error value
          errval2(ii) = norm(X-W2*H2,'fro')^2/norm(X,'fro')^2*100  ;
          ii = ii + 1;
    end
end
%% Plot the error curve
figure,
subplot(2,3,[1 2 3])
semilogy([err0 errval],'b'),hold on,
semilogy([err0 errval2],'r'),hold on
title('Relative Percentage Error vs iterations','interpreter','latex','fontsize',10)
xlabel('Iteration','interpreter','latex','fontsize',10)
ylabel('RPE = $\frac{\| X - WH \|_F^2}{\| X \|_F^2} \times 100\% $','interpreter','latex','fontsize',10)
legend('PGD','APGD'), grid on, axis tight

subplot(234)
imagesc(X),colormap gray
title('Original image','interpreter','latex','fontsize',10);

subplot(235)
imagesc(W*H),colormap gray
title('Reconstrcted using PGD','interpreter','latex','fontsize',10);

subplot(236)
imagesc(W2*H2),colormap gray
title('Reconstrcted using APGD','interpreter','latex','fontsize',10);