%
% An example for solving the least squares 
% by stochastic gradient descent (SGD)
%
% Reference: 
% D. Needell et al, Stochastic gradient descent, weighted
% sampling and the randomized Kaczmarz algorithm, 
% Math. Program. Ser. A (2016) 155:549-573
%
% -----------------------------------------------------------
%
% generate an n-by-m matrix with uniform leavage scores. 
n = 1000; 
m = 10;  
kappa = 1e+3; 
U = orth(randn(n,m)); 
S = diag(linspace(1,1/kappa,m)); 
V = orth(randn(m,m)); 
A = U*S*V'; 
x = randn(m,1); 
b = A*x; 
err = randn(n,1); 
b = b + 0.25*norm(b)*err/norm(err); 

% compute the probability distribution prob(i)  
AnormF = norm(A,'fro');
for i = 1:n
   prob(i) = norm(A(i,:))^2/AnormF^2;
end

% set up parameters 
kmax = 100;    % max number of sgd iterations 
gamma = 0.01;  % learning rate 

x0 = zeros(m,1); 
for k = 1:kmax 

    s = rand;
    i = sum(s >= cumsum([0, prob]));    % with probability prob(i)

    dx = n*(A(i,:)*x0 - b(i))*A(i,:); 
    x = x0 - (gamma/(n*prob(i)))*dx';   % weighted SGD
    x0 = x; 

end 
ResSGD = norm(A*x0 - b)

xexact = A\b; 
ResExact = norm(A*xexact - b) 

ResRatio = ResSGD/ResExact
