function [p, vel] = sgdmupdate(p, g, vel, lr, m)
%SGDMUPDATE Update parameters via stochastic gradient descent with momentum
%
%   [NET,VEL] = SGDMUPDATE(NET,GRAD,VEL) updates the learnable parameters
%   of the dlnetwork NET using the stochastic gradient descent with
%   momentum algorithm, with a default global learning rate of 0.01 and a
%   momentum of 0.9. Use SGDMUPDATE to iteratively update learnable
%   parameters during network training.
%
%   Input GRAD contains the gradients of the loss with respect to each of
%   the network parameters. Input VEL contains the parameter velocities
%   from the previous call to SGDMUPDATE. GRAD and VEL must be tables with
%   the same structure as NET.Learnables, with a Value variable containing
%   a cell array of parameter gradients or velocities.  Input GRAD can be
%   obtained using the dlgradient and dlfeval functions.  The global
%   learning rate is multiplied by the corresponding learning rate factor
%   for each parameter in each layer of the dlnetwork.
%
%   If input VEL is empty, the function assumes no previous velocities and
%   executes like the first update in a series of iterations.
%
%   Outputs NET and VEL are the updated dlnetwork and parameter velocities,
%   respectively.
%
%   [PARAMS,VEL] = SGDMUPDATE(PARAMS,GRAD,VEL) updates the deep learning
%   parameters in PARAMS using the stochastic gradient descent with
%   momentum algorithm, with a default learning rate of 0.01 and a momentum
%   of 0.9. Input PARAMS can be a dlarray, a numeric array, a cell array, a
%   structure, or a table with a Value variable containing the learnable
%   parameters of the network. GRAD and VEL must have the same datatype and
%   ordering as PARAMS. Input GRAD can be obtained using the dlgradient and
%   dlfeval functions.  All parameter values are updated using the global
%   learning rate.
%
%   Outputs PARAMS and VEL are the updated parameters and parameter
%   velocities, respectively.
%
%   [___] = SGDMUPDATE(___,LEARNRATE,MOMENTUM) also specifies values to use
%   for the global learning rate and momentum. LEARNRATE must be a positive
%   scalar. MOMENTUM must be a scalar between 0 and 1.
%
%   Example 1:
%      % Perform an SGDM update step with learning rate of 0.01 and
%      % momentum of 0.95.
%      p = rand(3,3,4); 
%      vel = ones(3,3,4); 
%      g = ones(3,3,4); 
%      [p,vel] = sgdmupdate(p,g,vel,0.01,0.95);
%
%   Example 2:
%      % Perform a single epoch of training on a network that classifies 
%      % handwritten digits.
%
%      % Load training data and encode.
%      [XTrain,YTrain] = digitTrain4DArrayData;
%      XTrain = single(XTrain);
%      YEncoded = onehotencode(YTrain,2)';
%      
%      % Define the network and specify the average image using the 'Mean' option 
%      % in the image input layer.
%      layers = [
%          imageInputLayer([28 28 1],'Name','input','Mean',mean(XTrain,4))
%          convolution2dLayer(5,20,'Name','conv1')
%          reluLayer('Name','relu1')
%          convolution2dLayer(3,20,'Padding',1,'Name','conv2')
%          reluLayer('Name','relu2')
%          fullyConnectedLayer(size(YEncoded,1),'Name','fc')];
%      net = dlnetwork(layerGraph(layers));
%      
%      % Initialize parameter velocities.
%      velocity = [];
%      
%      % Loop over one epoch of mini-batches.
%      miniBatchSize = 128;
%      for i = 1:miniBatchSize:numel(YTrain)
%          % Read mini-batch of data and convert the labels to dummy variables.
%          idx = i:min(i+miniBatchSize-1,numel(YTrain));
%          
%          % Convert mini-batch of data to dlarray.
%          X = dlarray(XTrain(:,:,:,idx),'SSCB');
%          
%          % Evaluate the model gradients using dlfeval and the
%          % modelGradients function defined below.
%          gradients = dlfeval(@modelGradients,net,X,YEncoded(:,idx));
%          
%          % Update the network parameters using the SGDM algorithm.
%          [net,velocity] = sgdmupdate(net,gradients,velocity);
%      end
%
%      function gradients = modelGradients(net,X,Y)
%          loss = crossentropy(softmax(forward(net,X)),Y);
%          gradients = dlgradient(loss,net.Learnables);
%      end
%
%   See also ADAMUPDATE, DLFEVAL, DLGRADIENT, DLNETWORK, DLUPDATE, RMSPROPUPDATE

%   Copyright 2019-2023 The MathWorks, Inc.

arguments
    p
    g
    vel
   
    lr(1,1) {mustBeNumeric, mustBeFinite, mustBeNonnegative} = 0.01;
    m(1,1) {mustBeNumeric, mustBeGreaterThanOrEqual(m, 0), mustBeLessThanOrEqual(m, 1)} = 0.9;
end

persistent func
if isempty(func)
    func = deep.internal.LearnableUpdateFunction( ...
        @iSingleStepValue, ...
        @iSingleStepParameter );
end

if isempty(vel)
    % Execute a first-step update with vel set to 0.  The step will create an
    % array for vel that is the correct size.
    paramArgs = {g};
    fixedArgs = {0, lr, m};
else
    % Execute the normal update
    paramArgs = {g, matlab.lang.internal.move(vel)};
    fixedArgs = {lr, m};
end

[p, vel] = deep.internal.networkContainerFixedArgsFun(func, ...
    p, matlab.lang.internal.move(paramArgs), fixedArgs);
end


function [p, vel] = iSingleStepParameter(p, g, vel, lr, m)
% Apply per-parameter learn-rate factor
lr = lr .* p.LearnRateFactor;

if lr ~= 0
    [step, vel] = nnet.internal.cnn.solver.sgdmstep(g, vel, lr, m);
    v = p.Value;
    p.Value = [];
    v = v + step;
    p.Value = v;
end
end


function [p, vel] = iSingleStepValue(p, g, vel, lr, m)
% Short-circuit for 0 learn-rate.
if lr ~= 0
    [step, vel] = nnet.internal.cnn.solver.sgdmstep(g, vel, lr, m);
    p = p + step;
end
end
