clear,close all,clc
% A comparision of NNLS using various iterative first order methods
%{ 
% Descriptions :
% NNLS model 
%      min_{ x \geq 0} (1/2) || Ax - b ||_2^2
%
% Equivalent NNQP model
%      min_{ x \geq 0} (1/2) x'Qx - p'x 
% where Q = A'*A, p = A'*b
%
% For THETA and Theta,
%
% THETA = I_n - Q/norm(Q,2)
% Theta = p/norm(Q,2) 
%
% Algorithms :
% PGD : projected gradient descent
% APGD : accelerated gradient descent, with variations :
%     1. FISTA parameter
%     2. FISTA parameter + adaptive restart
%     3. Optimal beta (using conditional number)
%     4. Optimal beta + adaptive restart 
%     5. Paul Tseng parameter
%     6. Paul Tseng parameter + adaptive restart
%
% Also see https://angms.science/doc/NMF/nnls_pgd.pdf
% Created : 2018-11-3, 
% Last update : 2018-11-7
% Andersen Ang @ UMONS, BE, angms.science
%}
%%
% dimensions
m = 100;  % row 
n = 20;   % column
% random generate matrix A
 A = rand(m,n);
% random generate the ground true nonnegative vector x
 xtrue = rand(n,1);
 xtrue(xtrue<0) = 0;
% form the vector b
 b = A*xtrue;
%% Some useful variables
% Initial guess
 x0 = rand(n,1);
% iteration max
 kmax = 600;
%% PGD 
 % Parameter matrix and vector
 THETA = eye(n) - A'*A/norm(A'*A,2);
 theta = A'*b/norm(A'*A,2);
 
 x = x0;  % initialize
 for k = 1 : kmax  % Loop
    x = max(THETA*x + theta, 0);  % update
    e_pgd(k) = norm(A*x - b);     % compute error
 end
%% APGD-1 (Beck and Teboulle alpha beta)
 % Acceleration parameter
 alpha = 0.9; % alpha in (0, 1)
 % initialize the variable and the second variable 
 x = x0;
 y = x;
 % Loop
 for k = 1 : kmax
    % store previous variable
    x_ = x; 
    % store previous alpha
    alpha_ = alpha;
    % computation of acceleration parameter
    alpha = 0.5*(sqrt(alpha^4+4*alpha^2) - alpha^2);
    beta = alpha_*(1-alpha_)/(alpha_^2+alpha);
    % update
    x = max(THETA*y + theta, 0);
    % extrapolation
    y = x + beta*(x - x_);
    % compute error
    e_apgd(k) = norm(A*x - b);
 end
%% APGD-1 with restart
 % Acceleration parameter
 alpha = 0.9; % alpha in (0, 1)
  % store initial alpha value for restart
  alpha0 = alpha; 
 % initialize the variable and the second variable 
 x = x0;
 y = x;
 % Loop
for k = 1 : kmax 
    % store previous variable
    x_ = x;
    % store previous alpha
    alpha_ = alpha;
    % computation of acceleration parameter
    alpha = 0.5*(sqrt(alpha^4+4*alpha^2) - alpha^2);
    beta = alpha_*(1-alpha_)/(alpha_^2+alpha);
    % update
    x = max(THETA*y + theta, 0);
    % extrapolation
    y = x + beta*(x - x_);
    % compute error
    e_apgd_r(k) = norm(A*x - b);
    % check does error increase
    if (k > 1) && (e_apgd_r(k) > e_apgd_r(k-1))
      % restart, just consider a gradient step update
      x  = max(THETA*x_ + theta, 0);
      e_apgd_r(k) = norm(A*x - b);
      % restart 
      y = x;
      alpha = alpha0;
    end 
end
%% APGD-2 with optimal beta
 Q = A'*A;
 condQm1 = 1/cond(Q); %  smallest / largest  singular value of A 
 beta = (1-sqrt(condQm1)) / (1 + sqrt(condQm1)); 
 % initialize the variable and the second variable 
 x = x0;
 y = x;
 % Loop
 for k = 1 : kmax
    % store previous variable
    x_ = x; 
    % update
    x = max(THETA*y + theta, 0);
    % extrapolation
    y = x + beta*(x - x_);
    % compute error
    e_apgd_b(k) = norm(A*x - b);
 end
%% APGD-2 with optimal beta with restart
 % beta calculated above
 
 % initialize the variable and the second variable 
 x = x0;
 y = x;
 % Loop
for k = 1 : kmax 
    % store previous variable
    x_ = x;
    % update
    x = max(THETA*y + theta, 0);
    % extrapolation
    y = x + beta*(x - x_);
    % compute error
    e_apgd_b_r(k) = norm(A*x - b);
    % check does error increase
    if (k > 1) && (e_apgd_b_r(k) > e_apgd_b_r(k-1))
      % restart, just consider a gradient step update
      x  = max(THETA*x_ + theta, 0);
      e_apgd_b_r(k) = norm(A*x - b);
      % restart 
      y = x;
    end 
end 
%% APGD-3 Paul Tseng parameters 
 % initialize the variable and the second variable 
 x = x0;
 y = x;
 % Loop
 for k = 1 : kmax
    % store previous variable
    x_ = x; 
    % update
    x = max(THETA*y + theta, 0);
    % extrapolation
      % computation of acceleration parameter
      beta = (k-1)/(k+2);
    y = x + beta*(x - x_);
    % compute error
    e_apgd_p(k) = norm(A*x - b);
 end
%% APGD-3 Paul Tseng parameters with restart
 % store restart k counter
  kr = 0;
 % initialize the variable and the second variable 
 x = x0;
 y = x;
 % Loop
for k = 1 : kmax   
    % store previous variable
    x_ = x;
    % computation of acceleration parameter
    beta = (k-kr-1)/(k-kr+2);
    % update
    x = max(THETA*y + theta, 0);
    % extrapolation
    y = x + beta*(x - x_);
    % compute error
    e_apgd_p_r(k) = norm(A*x - b);
    % check does error increase
    if (k > 1) && (e_apgd_p_r(k) > e_apgd_p_r(k-1))
      % restart, just consider a gradient step update
      x  = max(THETA*x_ + theta, 0);
      e_apgd_p_r(k) = norm(A*x - b);
      % restart 
      y = x;
      kr = kr + k;
    end 
end
%% Plot error vs iterations
semilogy(e_pgd,'k--','linewidth',2),hold on,
semilogy(e_apgd,'b','linewidth',8),hold on,
semilogy(e_apgd_r,'b--','linewidth',2),hold on,
semilogy(e_apgd_b,'r','linewidth',4),hold on,
semilogy(e_apgd_b_r,'y--','linewidth',3),hold on,
semilogy(e_apgd_p,'g','linewidth',4),hold on,
semilogy(e_apgd_p_r,'g--','linewidth',2),hold on,
legend('normal PGD','A','A + r','A(best beta)','A(best beta) + r','A (Paul Tseng)','A (Paul Tseng) + r')
title('NNLS','interpreter','latex','fontsize',20)
xlabel('Iteration/$k$','interpreter','latex','fontsize',15)
ylabel('$\| Ax - b\|$','interpreter','latex','fontsize',15)
axis tight,grid on