%% NMF by PGD
% Andersen Ang @ MARO, FPM, UMONS, Belgium
% Last update 2018-March-10
%%
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 = 86;
%% Initialization
  % Initialize W and H
  W0 = rand(m,r);
  H0 = rand(r,n);
  % Set the epsilon 
  epsilon = 0;
  epsW = epsilon*ones(size(W0));
  epsH = epsilon*ones(size(H0));
  % inital error value
  err0 = norm(X-W0*H0,'fro')^2/norm(X,'fro')^2*100;
  % max iteration
  maxiter = 30;
  
%% The main loop - PGD with constant step size
W = W0;
H = H0;
for i = 1 : maxiter % or other stopping conditions
  % On varaible W
      % gradient 
      gradW = (W*H-X)*H';
      % step size 
      stepW = 1/(norm(H*H','fro')); 
      % PGD step
      W = max(W - stepW * gradW, epsW );
      % compute error value
      errval(1,i) = norm(X-W*H,'fro')^2/norm(X,'fro')^2*100;
      % store W
      sWc(:,:,i) = W;
      
  % On variable H
      % gradient 
      gradH = W'*(W*H-X);
      % step size
      stepH = 1/(norm(W'*W,'fro'));
      % PGD step
      H = max(H - stepH * gradH, epsH);
      % compute error value
      errval(2,i) = norm(X-W*H,'fro')^2/norm(X,'fro')^2*100;
      % store H
      sHc(:,:,i) = H;
end
%% Plot the error curve
figure,
botherrval = errval(:);
plot(botherrval,'k'),hold on
plot(1:2:numel(botherrval), botherrval(1:2:end),'b.','markersize',16),hold on
plot(2:2:numel(botherrval), botherrval(2:2:end),'r.','markersize',16),hold on
title('Relative Percentage Error verus iterations','interpreter','latex','fontsize',20)
xlabel('Iteration','interpreter','latex','fontsize',15)
ylabel('RPE = $\frac{\| X - WH \|_F^2}{\| X \|_F^2} \times 100\% $','interpreter','latex','fontsize',20)
legend('Error curve (normal gradient, constant size size)','After W is updated (normal gradient, constant size size)','After H is updated (normal gradient, constant size size)')
grid on, axis tight

hold on

%% The main loop - PGD with exact line search
W = W0;
H = H0;
for i = 1 : maxiter % or other stopping conditions
  % On varaible W
      % gradient 
      gradW = (W*H-X)*H';
      % step size 
      stepW = trace  ( -(X-W*H)'*gradW*H  )  / trace( H'*gradW'*gradW*H );
      % PGD step
      W = max(W - stepW * gradW, epsW );
      % compute error value
      errval(1,i) = norm(X-W*H,'fro')^2/norm(X,'fro')^2*100;
      % store W
      sWb(:,:,i) = W;
  
  % On variable H
      % gradient 
      gradH = W'*(W*H-X);
      % step size
      stepH = trace( -(X-W*H)'*W*gradH ) / trace( gradH'*W'*W*gradH );
      % PGD step
      H = max(H - stepH * gradH, epsH);
      % compute error value
      errval(2,i) = norm(X-W*H,'fro')^2/norm(X,'fro')^2*100;
      % store H
      sHb(:,:,i) = H;
end
%% Plot the error curve

botherrval = errval(:);
plot(botherrval,'k--'),hold on
plot(1:2:numel(botherrval), botherrval(1:2:end),'c.','markersize',16),hold on
plot(2:2:numel(botherrval), botherrval(2:2:end),'m.','markersize',16),hold on
title('Relative Percentage Error verus iterations','interpreter','latex','fontsize',20)
xlabel('Iteration','interpreter','latex','fontsize',15)
ylabel('RPE = $\frac{\| X - WH \|_F^2}{\| X \|_F^2} \times 100\% $','interpreter','latex','fontsize',20)
legend('Error curve (normal gradient, constant size size)','After W update (normal gradient, const. size size)','After H update (normal gradient, const. size size)','Error curve (normal gradient, optimal size size)','After W update (normal gradient, optimal size size)','After H update (normal gradient, optimal size size)')
grid on, axis tight


%%
figure
subplot(5,2,1), i = 1;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,2),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,3), i = 2;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,4),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,5), i = 3;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,6),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,7), i = 4;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,8),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,9), i = 5;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,10),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)




figure
subplot(5,2,1), i = 20;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,2),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,3), i = 40;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,4),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,5), i = 100;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,6),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,7), i = 250;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,8),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)
subplot(5,2,9), i = 500;
imagesc(sWc(:,:,i)*sHc(:,:,i)),colormap gray,title(i)
subplot(5,2,10),
imagesc(sWb(:,:,i)*sHb(:,:,i)),colormap gray,title(i)




