classdef optimizationFunctionScaling
    % This class keeps the scale factors used to scale objective and
    % constraint functions (linear scaling) in nonlinear optimization.
    % 
    % The scale factors must be supplied via the method
    % setScaleFactors(options, factors)
    %
    %  - options must be a struct-style container with 'ScaleProblem' as a
    %    property. Valid values are 'none' (default) or 'obj-and-constr'
    %  - factors must be a struct with one or more of these fields
    %    (factors) and scalar values: 
    %     obj: scale factor for objective and objective gradients
    %     cEq: scale factor for nonlinear equalities and gradients
    %     cIneq: scale factor for nonlinear inequalities and gradients
    %
    %  After setting the scale factors, specific scaling functions can be
    %  accesses:
    %  - fval = scaleObjective(fval)
    %  - [cIneq, cEq] = scaleConstraints(cIneq, cEq)
    %  - grad = scaleObjectiveGrad (grad)
    %  - [cIneqGrad, cIneqGrad] = scaleConstraintGrads(cIneqGrad, cIneqGrad)

    %   Copyright 2023 The MathWorks, Inc.
    
    properties(SetAccess=private)
        % Default options for ScaleProblem
        defaultopt = struct('ScaleProblem','none');
        
        % Scale factor for objective and gradient
        obj = 1
        % Scale factor for nonlinear equalities and gradients
        cEq = 1
        % Scale factor for nonlinear inequalities and gradients
        cIneq = 1
        % factors used to scale ALL the constraints (including linear)
        constr = 1
    end

    properties(SetAccess=private)
        % Method for scaling objective
        scaleObjective = @noscale
        % Method for scaling constraint
        scaleConstraints = @noscale
        % Method for scaling objective gradient
        scaleObjectiveGrad = @noscale
        % Method for scaling constraint gradient
        scaleConstraintGrads = @noscale
        %scaleJacobian % lsq/fsolve use
    end

    properties(SetAccess=private,Hidden)
        Factors = {'obj', 'cEq' 'cIneq','constr'};
        % flags
        objIsScaled = false 
        constrIsScaled = false        
    end

    methods

        function self = setScaleFactors(self, options_, factors_)
            % Specify the scaling factors.
            %
            %  - options must be a struct-style container with
            %    'ScaleProblem' as a property. Valid values are 'none'
            %    (default) or 'obj-and-constr'.
            %  - factors must be a struct with one or more of these fields
            %    (factors) and scalar values:
            %    obj: scale factor for objective and objective gradients
            %    cEq: scale factor for nonlinear equalities and gradients
            %    cIneq: scale factor for nonlinear inequalities and gradients
            arguments
                self
                options_ = struct()
                factors_ struct = struct()
            end            
            fnames = intersect(self.Factors, fieldnames(factors_));

            for ii = 1:length(fnames)
                self.(fnames{ii}) = factors_.(fnames{ii});
            end
            
            scaleProblem = optimget(options_,'ScaleProblem',self.defaultopt,'fast');
            if strcmpi(scaleProblem,'obj-and-constr')
                self.constrIsScaled = true;
                self.objIsScaled = true;
                self.scaleObjective = @(x) scaleObjOrGrad(self, x);
                self.scaleObjectiveGrad = @(x) scaleObjOrGrad(self, x);

                self.scaleConstraints = @(cIneq, cEq) scaleConstr(self, cIneq, cEq);
                self.scaleConstraintGrads = @(cineqGrad, ceqGrad) ... 
                    scaleConstrGrad(self,numel(self.cIneq),numel(self.cEq), ...
                    cineqGrad, ceqGrad);
            end
        end
    end
end

function varargout = noscale(~, varargin)
% No scaling
    varargout = varargin;
end

function FvalOrGrad = scaleObjOrGrad(fscale, FvalOrGrad)
    FvalOrGrad = fscale.obj * FvalOrGrad;
end

function [cIneq, cEq] = scaleConstr(fscale, cIneq, cEq)
    cIneq = fscale.cIneq .* cIneq;
    cEq = fscale.cEq .* cEq;
end

function [cineqGrad, ceqGrad] = scaleConstrGrad(fscale,mNonlinIneq,mNonlinEq,cineqGrad,ceqGrad)

    % If Jacobian is a vector, do not scale by a sparse scalar.
    % Multiplication by a sparse scalar changes any matrix to sparse.
    if mNonlinEq == 1
        ceqGrad = ceqGrad*fscale.cEq;
    else
        ceqGrad = ceqGrad*spdiags(fscale.cEq,0,mNonlinEq,mNonlinEq);
    end
    if mNonlinIneq == 1
        cineqGrad = cineqGrad*fscale.cIneq;
    else
        cineqGrad = cineqGrad*spdiags(fscale.cIneq,0,mNonlinIneq,mNonlinIneq);
    end    
end
