clear,close all,clc
m = 20; % no. of rows
n = 14; % no. of cols
r = 5; % factoization rank

Wtrue = randn(m,r); % ground truth W 
Htrue = randn(r,n); % ground truth H
sigma = 0.01; % noise level of additive noise
M     = Wtrue * Htrue + sigma * randn(m,n); % observed data matrix
normM = norm(M,'fro'); % norm of M

Wini = randn(m,r); % initial W 
Hini = randn(r,n); % inital H
kmax = 200;  % max number of iteration
%% Matrix way
W = Wini;
H = Hini;
t_start = tic;
for k = 1 : kmax
 HHt   = H*H';
 MHt   = M*H';
 alpha = 1/norm(HHt,2);   % step size 
 gradW = W*HHt - MHt;     % gradient W
 W     = W - alpha * gradW;

 WtW   = W'*W;
 WtM   = W'*M;
 alpha = 1/norm(WtW,2); % step size 
 gradH = WtW*H - WtM;   % gradient H
 H     = H - alpha * gradH;
 
 e_mat(k)  = norm(M - W*H,'fro')/normM;   % error
 t_mat(k)  = toc(t_start); % time 
end
%% Vector way,  improved
W = Wini;
H = Hini;
Im = eye(m);
In = eye(n);
t_start = tic;
for k = 1 : kmax
 A     = kron( H', Im );
 x     = W(:); % x = vec(W)
 alpha = 1/norm(A'*A,2);
 Atb   = M*H';  Atb = Atb(:); 
 gradx = A'*A*x - Atb;
 x     = x - alpha * gradx;
 W     = reshape(x,m,r);  % inverse of vec
 clear x A
 
 A     = kron( In, W );
 x     = H(:); % x = vec(H)
 alpha = 1/norm(A'*A,2);
 Atb   = W'*M;  Atb = Atb(:);
 gradx = A'*A*x - Atb;
 x     = x - alpha * gradx;
 H     = reshape(x,r,n);  % inverse of vec 
 clear x A
 
 e_vec(k)  = norm(M - W*H,'fro')/normM;  
 t_vec(k)  = toc(t_start);
end
%% Plotting
subplot(121)
plot(e_mat,'c-','linewidth',2.5),hold on
plot(e_vec,'r:','linewidth',2.5),
axis tight,grid on
xlabel('Iteration','FontSize',22)
ylabel('$\frac{ \| M-WH\|_F} {\| M \|_F}$','FontSize',35,'interpreter','latex')
lgd = legend('Matrix-form','Vector-form');
lgd.FontSize = 22;

subplot(122)
plot(t_mat,e_mat,'c-','linewidth',2.5),hold on
plot(t_vec,e_vec,'r:','linewidth',2.5),
axis tight,grid on
xlabel('Time (sec)','FontSize',22)
lgd.FontSize = 22;