function [x,fval,exitflag,output] = sgd(fun, x0, trainNdx, options)
% SGD Stochastic gradient descent optimizer
%
% [x,fval,exitflag,output] = SGD(fun, x0, trainNdx, options) finds
% parameter value x that minimizes the value of the function handle
% fun, starting from x0. fun must accept two arguments: the parameter
% vector and an index vector indicating which elements to include in a
% batch (i.e., a subset of vector trainNdx). fun produces the
% objective function value to be minimized. options is a struct to set
% parameters of the optimization. fval is the final value of the
% function accumulated in the last epoch. exitflag indicates the
% termination condition, and output is a struct with information about
% the optimization.
%
%
% Valid fields for options struct include:
%
% MaxIter	Maximum number of optimization iterations [1000]
% MaxFunEvals	Maximum number of function evaluations [100*numel(x0)]
% TolX	Tolerance on the parameter vector [1e-6]
% TolFun	Tolerance on the function value [1e-4]
% TolGrad	Tolerance on gradient L2 norm [1e-4]
% GradObj 
%     [on]: fun returns a gradient value as the second output
%     off:  Finite differences are used to estimate the gradient
% Display	
%     off: 	    No display information
%     batch:    Display information at every batch
%			iter: 	  Display information at every iteration
%			[final]: 	Display information at convergence [default]
% Shuffle    Shuffle batches after each iteration/epoch [true]
% BatchSize  Number of elements of fun to use in each batch [5]
% Lambda     Convex weight for long-term (second-order) weight updates [0.99]
% Mu         Meta learning rate for Eta [0.02]
% Eta        Initial learning rate [1e-4]
% SaveFreq   How often (in epochs) to snapshot parameters [0]
% SaveFile   Base path for parameter snapshots (epoch iteration number appended)
% DiffMinChange Smallest step for finite difference gradient calculation [1e-8]
% DiffMaxChange Largest step for finite difference gradient calculation [1e-1]
% 
% 
% Values of the exitflag indicate
%
% -2     Inf in function evaluation (reduce Eta or Mu)
% -1     NaN in function evaluation (reduce Eta or Mu)
%  0     Iteration limit exceeded
%  1     Converged on gradient norm
%  2     Change in x was smaller than the TolX tolerance.
%  3     Converged on value function difference
%
%
%  Learning rate/momentum updates are given by Equation (10) in Orr et
%  al. (1999). Another article by Schraudolph (1999) uses the
%  technique, crediting its origin as Almeida et al. (1999).
%
%  Note: if fun uses the Matlab Parallel/Distributed computing
%  toolbox, this function works transparently, but the values
%  returned by fun must be on the CPU (achieved by using GATHER).
%
%
%  References:
%    G. B. Orr, N. Schraudolph, and F. Cummins. (1999). Momentum and
%    Learning Rate Adaptation.
%    http://www.willamette.edu/~gorr/classes/cs449/momrate.html
%    (Accessed 17 June 2016).
%
%    N. N. Schraudolph. (1999). "Local gain adaptation in stochastic
%    gradient descent", in Proc. Intl. Conf. Artificial Neural
%    Networks, 569--574.
%
%    L. L. B. Almeida, T. Langlois, J. D. Amaral, and A. Plakhov,
%    "Parameter adaptation in stochastic optimization", in On-Line
%    Learning in Neural Networks (D. Saad, Ed.), 111--134.
                                                                  
  
  

% Jerod Weinman
% jerod@acm.org
% Copyright (c) 2012-2016

% $Id: sgd.m 98 2016-06-19 02:26:53Z weinman $
 
% Jerod Weinman
% jerod@acm.org
% Copyright 2012, 2014, 2016
 
%    This file is part of Matlab SGD.
%
%    Matlab SGD is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.
%
%    Matlab SGD is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with Matlab SGD.  If not, see <http://www.gnu.org/licenses/>.

%% Process arguments
if nargin<4
  options = struct;
end;

if isfield(options,'MaxIter') && ~isempty(options.MaxIter)
  maxIter = options.MaxIter;
else
  maxIter = 1000;
end;

if isfield(options,'MaxFunEvals') && ~isempty(options.MaxFunEvals)
  maxFunEvals = options.MaxFunEvals;
else
  maxFunEvals = 100*numel(x0);
end;

if isfield(options,'TolX') && ~isempty(options.TolX)
  xtol = options.TolX;
else
  xtol = 1e-6;
end;

if isfield(options,'TolFun') && ~isempty(options.TolFun)
  tol = options.TolFun;
else
  tol = 1e-4;
end;

if isfield(options,'TolGrad') && ~isempty(options.TolGrad)
  gtol = options.TolGrad;
else
  gtol = 1e-4;
end;

if isfield(options,'Display') && ~isempty(options.Display)
  switch options.Display
    case 'off'
      dspOpt = 0;
    case 'iter'
      dspOpt = 1;
    case 'final'
      dspOpt = 2;
    case 'batch'
      dspOpt=3;
    otherwise
      error('Unknown value for Display option');
  end;
else
  dspOpt = 2;
end;

if isfield(options,'GradObj') && ~isempty(options.GradObj)
  switch options.GradObj
    case 'off'
      gradObj = 0;
    case 'on'
      gradObj = 1;
    otherwise
      error('Unknown value for GradObj option');
  end;
else
  gradObj = 1;
end;

if isfield(options,'Shuffle')
  if isscalar(options.Shuffle);
    shuffle = logical(options.Shuffle);
  else
    error('Shuffle option must be a scalar (used as logical)')
  end
else
  shuffle = true;
end

if isfield(options,'BatchSize') && ~isempty(options.BatchSize)
  batchSize = options.BatchSize;
else
  batchSize = 5;
end;

if isfield(options,'Lambda') && ~isempty(options.Lambda)
  lambda = options.Lambda;
else
  lambda = 0.99;
end;

if isfield(options,'Mu') && ~isempty(options.Mu)
  mu = options.Mu;
else
  mu = 0.02;
end;


if isfield(options,'Eta') && ~isempty(options.Eta)
  eta = options.Eta;
else
  eta = 1e-4;
end;

if isfield(options,'DiffMinChange') && ~isempty(options.DiffMinChange)
  diffMin = options.DiffMinChange;
else
  diffMin = 1e-8;
end;

if isfield(options,'DiffMaxChange') && ~isempty(options.DiffMaxChange)
  diffMax = options.DiffMaxChange;
else
  diffMax = 1e-1;
end;


if isfield(options,'SaveFile') && ~isempty(options.SaveFile)
  fname = options.SaveFile;
  sfreq = 1;
else
  fname = [];
end;

if isfield(options,'SaveFreq') && ~isempty(options.SaveFreq)
  if isempty(fname)
    error('Must specify SaveFile if SaveFreq is given');
  else
    sfreq = options.SaveFreq;
  end;
else
  sfreq = 0;
end;

dspHeader = sprintf('\titer\tbatch\tf(x)\t||g||\t|eta|\t||v||\n');

maxIter=maxIter+1; % ???
batchIndices = mkBatches(trainNdx, batchSize);
numBatches = length(batchIndices);

% Initialization
v = zeros(size(x0));      % auxiliary vector for meta learning rate
vn = zeros(size(x0));     % auxiliary vector for gradient's exponential average
vd = ones(size(x0));      % auxiliary vector for gradient^2 exponential average
grad = zeros(size(x0));   % epoch gradient accumulator
eta = eta*ones(size(x0)); % per parameter learning rate 

x = x0;                   % current parameter setting
fvalold = NaN;
fval = NaN;
f = zeros(1,numBatches);

% Note: The resetMean test is not enabled so that with little history
% (whether at beginning or after value-increasing learning rate adjustment)
%  the damped multiplicative updates slowly move away from unity.
resetMean = 1;

funcCount = 0;

iter = 0; % Note: iter means epoch (complete pass over training data)


if (dspOpt==1 | dspOpt==3)
  fprintf(dspHeader);
end

%% Main iteration loop
while (iter<=maxIter)
  
  grad(:) = 0;                                % Reset gradient accumulator
  xdiff = -Inf;                               % Reset parameter step
  
  %% Batch loop
  for batch=1:numBatches
    
    batchNdx = batchIndices{batch};

    if gradObj
      [fbval,g] = feval(fun, x, batchNdx );   % Evaluate function on batch
    else
      [fbval] = feval(fun, x, batchNdx );     % Calculate grad by 
      g = fd(fun,fbval,x,batchNdx,diffMin,diffMax);  % finite differences
    end      

    % Check for NaN/Inf
    [done,msg,exitflag,output] = checkBatchValue(fbval,iter,funcCount,dspOpt);
    if done, return; end
    
    f(batch) = fbval;                         % Cache batch fun/grad values
    grad = grad + g;                          % Tally for total gradient
    
    if 0 %resetMean %|| (iter==0 && batch==1) % (See note above on resetMean)
      vd = g.^2;                              % Initial condition for series vd
    else
      vd = lambda*vd + (1-lambda)*g.^2;       % Exponential mean of gradient^2
    end

    % Use eps as lower bound avoiding division by a too-small number
    v = vn./max(vd,eps);                      % Calculate linear normalizer

    % (Meta) Gradient descent update in for learning rate in log space:
    % 1+mu*g*v is an approximation to exp(mu*g*v). See Eq. (7) in Orr et
    % al. (1999). Bounded below at 0.5 "to guard against very small (or
    % even negative) factors"  (Orr et al.).
    eta = eta.*max(0.5, 1+mu*g.*v);           % Update learning rate
    x = x - eta.*g;                           % Update parameters

    if 0 %resetMean %|| (iter==0 && batch==1) % (See note above on resetMean)
      vn = g;                                 % Initial condition for series vn
      resetMean = 0;
    else
      vn = lambda*vn + (1-lambda)*g;          % Exponential mean of gradient
    end

    xdiff = max(xdiff, norm(eta.*g,'inf'));   % Fold: accumulate maximum step

    batchDisplay(dspOpt,dspHeader,iter,batch,fbval,g,eta,v);    
  end; % Inner optimization loop (for each batch)
  
  funcCount = funcCount + 1;                  % *Total* function evaluation
  fval = sum(f);                              % Tally values over epoch

  if gradObj==0                               % If finite differencing, there is
    funcCount = funcCount + numel(x);         % one evaluation per parameter
  end
  
  snapshot(iter,fname,sfreq,maxIter,x,eta,vn,vd);     % Save history
  iterDisplay(dspOpt,dspHeader,iter,fval,grad,eta,v); % Print epoch info
  
  [done,msg,exitflag,output] = ...
      checkTermConditions(iter,fval,tol, fvalold,funcCount,grad,gtol,xdiff, ...
                          xtol,maxIter,maxFunEvals,dspOpt);
  
  if done, return; end

  
  if (fvalold < fval)      % Function (epoch total) increased over prior epoch
    eta = 0.5 * eta;       % Halve learning rate to take smaller steps
    vn = zeros(size(vn));  % Clear exponential gradient mean
    vd = ones(size(vd));   % Clear exponential gradient^2 mean
    resetMean = 1;         % (If enabled) Start new mean series
  end

  fvalold = fval;
  iter = iter + 1;
  
  if shuffle % Shuffle examples and rebatch
    trainNdx = trainNdx(randperm(length(trainNdx)));
    batchIndices = mkBatches(trainNdx, batchSize);
  end
  
end  % Outer optimization loop (epoch)

end % SGD
%-------------------------------------------------------------------------------
function batchIndices = mkBatches(trainNdx, batchSize)
% MKBATCHES Divide trainNdx into batchSize chunks (except for last)

% NB: Assume trainNdx is row vector

if mod(length(trainNdx),batchSize)==0 % All equal-sized batches
  batchIndices = mat2cell(trainNdx, 1, ...
                          batchSize*ones(1,length(trainNdx)/batchSize) );
else                                  % Last batch will have leftovers
  lastBatchSize = mod( length(trainNdx), batchSize );
  numBatches = ceil( length(trainNdx) / batchSize );
  batchIndices = mat2cell( trainNdx, 1, ...
                           [batchSize*ones(1, numBatches-1), lastBatchSize] );
end
end %MKBATCHES
%-------------------------------------------------------------------------------
function snapshot(iter,fname,sfreq,maxIter,x,eta,vn,vd)
% SNAPSHOT Save parameters and learning rates when called for
  if sfreq && ~mod(iter,sfreq)
    %%% Intermediate save of parameters
    iterStr = sprintf( ['%0' num2str(floor(log10(maxIter)+1)) 'd'], iter );
    save([fname '-' iterStr] ,'x','eta','vn','vd');
  end;
end % SNAPSHOT
%-------------------------------------------------------------------------------
function grad = fd(fun, fval, x, batchIndices, diffMin, diffMax, fwd)
% FD Calculate simple finite difference gradient

% Calculate step magnitude: base is sqrt(eps), scaled by magnitude of x
del= sqrt(eps).*x;
del(del==0) = sqrt(eps); % Ensure no non-zero steps                         
% Bound step to [diffMin,diffMax]
del = sign(del).*min( max( abs(del), diffMin),diffMax); 

xdel = x + del;                         % Take forward step from x
grad = zeros(size(x));                  % Pre-allocate gradient
xfd = x;                                % Initialize stepped vector 

for k=1:numel(x)
  xfd(k) = xdel(k);                     % Step k'th parameter
  fdel = feval(fun,xfd,batchIndices);   % Evaluate function
  grad(k) = (fdel-fval)./del(k);        % Forward difference
  xfd(k) = x(k);                        % Unstep k'th parameter
end;
end % FD
%-------------------------------------------------------------------------------
function [done,msg,exitflag,output] = checkBatchValue(fbval,iter,funcCount,dspOpt)
    
msg = []; exitflag = []; output = [];
done = 0;

if isnan(fbval)
  %%%%% NaN function value
  msg = 'SGD: NaN in batch function value';
  
  exitflag = -1;
  output = struct('iterations',{iter},...
                  'funcCount',{funcCount},...
                  'algorithm',{'stochastic gradient descent'},...
                  'message',{msg});
  
  if dspOpt
    display(msg);
  end;
  
  done = 1;
  return;
end

if isinf(fbval)
  %%%%% Inf function value
  msg = 'SGD: Inf batch function value';
  
  exitflag = -2;
  output = struct('iterations',{iter},...
                  'funcCount',{funcCount},...
                  'algorithm',{'stochastic gradient descent'},...
                  'message',{msg});
  
  if dspOpt
    display(msg);
  end;
  
  done = 1;
end
end % CHECKBATCHVALUE
%-------------------------------------------------------------------------------
function batchDisplay(dspOpt,dspHeader,iter,batch,fbval,g,eta,v)
% BATCHDISPLAY Print information about batch values

if (dspOpt==3)
  if mod(batch,20)==0
    fprintf(dspHeader);
  end
  fprintf('sgd\t%d\t%d\t%g\t%g\t%g\t%g\n', ...
                  iter, batch, fbval, ...
                  norm(g,2), norm(eta,inf), norm(v,2));
end
end % BATCHDISPLAY
%-------------------------------------------------------------------------------
function iterDisplay(dspOpt,dspHeader,iter,fval,grad,eta,v)
% ITERDISPLAY Print information about values at conclusion of epoch
  if (dspOpt==1 || dspOpt==3)
    if iter > 1 && mod(iter,20)==0
      fprintf(dspHeader);
    end
    fprintf('sgd\t%d\t-\t%g\t%g\t%g\t%g\n',iter,fval,...
                    norm(grad,2),norm(eta,inf), norm(v,2));
  end
end % ITERDISPLAY
%-------------------------------------------------------------------------------
function [done,msg,exitflag,output] = checkTermConditions(iter,fval, tol, ...
                                                  fvalold,funcCount,grad,...
                                                  gtol,xdiff,xtol, ...
                                                  maxIter,maxFunEvals,dspOpt)
% CHECKTERMCONDITIONS Check function difference, gradient
% magnitude, parameter change, iteration count, and function
% evaluations for termination.
  
msg = []; exitflag = []; output = [];
done = 0;

  if iter>0 &&  ...%fval<fvalold && ...
        (2*abs(fval-fvalold) <= tol*(abs(fval)+abs(fvalold)+eps))
    %%%%% Value tolerance convergence
    msg = 'SGD: value difference below tolerance';
    
    exitflag = 3;
    output = struct('iterations',{iter},...
                    'funcCount',{funcCount},...
                    'algorithm',{'stochastic gradient descent'},...
                    'message',{msg});
    if dspOpt
      display(msg);
    end;

    done = 1;
  end;
  
  if (norm(grad,2)<gtol)
    %%%%% Gradient tolerance convergance
    msg = 'SGD: gradient norm below tolerance';
    
    exitflag = 1;
    output = struct('iterations',{iter},...
                    'funcCount',{funcCount},...
                    'algorithm',{'stochastic gradient descent'},...
                    'message',{msg});
    if dspOpt
      display(msg);
    end;

    done = 1;
  end;
  
  if (xdiff<xtol)
    %%%%% Parameter tolerance convergence
    msg = 'SGD: parameter step below tolerance';
    
    exitflag = 2;
    output = struct('iterations',{iter},...
                    'funcCount',{funcCount},...
                    'algorithm',{'stochastic gradient descent'},...
                    'message',{msg});
    if dspOpt
      display(msg);
    end;
    
    done = 1;
  end;
  
  if (iter==maxIter)
    %%%%% Iteration count exceeded
    msg = 'SGD: iteration count exceeded';
    
    exitflag = 0;
    output = struct('iterations',{iter},...
                    'funcCount',{funcCount},...
                    'algorithm',{'stochastic gradient descent'},...
                    'message',{msg});
    if dspOpt
      display(msg);
    end;
    
    done = 1;
  end;
  
  if (funcCount>=maxFunEvals)
    %%%%% Function evaluation count exceeded
    msg = 'SGD: function evaluation count exceeded';
    
    exitflag = 0;
    output = struct('iterations',{iter},...
                    'funcCount',{funcCount},...
                    'algorithm',{'stochastic gradient descent'},...
                    'message',{msg});
    if dspOpt
      display(msg);
    end;
    
    done = 1;
  end;

end % CHECKTERMCONDITIONS
