classdef ReconstructionICAImpl < featlearn.impl.FeatureTransformImpl
%

%   Copyright 2016 The MathWorks, Inc.
    
    properties
        NumPredictors;              % Original number of predictors P.
        NumLearnedFeatures;         % Number of transformed predictors Q.
        Mu;                         % Predictor means if data was standardized - a P-by-1 vector or [].
        Sigma;                      % Predictor standard deviations if data was standardized - a P-by-1 vector or [].
        FitInfo;                    % Fitting information - a structure containing fitting information.
        TransformWeights;           % Matrix of learned transformation weights - a P-by-Q matrix.
        InitialTransformWeights;    % Matrix of initial transformation weights - a P-by-Q matrix.
        NonGaussianityIndicator;    % Q-by-1 vector indicating type of non-Gaussianity (elements +1 means super-Gaussian and elements -1 means sub-Gaussian)
    end
    
    properties(Constant)
        Delta = 1e-8;
    end
    
    methods
        function this = ReconstructionICAImpl(X,Q,transformWeights0,nonGaussianityIndicator,modelParams)
%ReconstructionICAImpl - Make implementation for reconstruction ICA (RICA).
%   this = ReconstructionICAImpl(X,Q,transformWeights0,nonGaussianityIndicator,modelParams)
%   takes a N-by-P matrix of predictors X, an integer Q specifying the
%   number of new features to extract, a P-by-Q matrix transformWeights0
%   specifying the initial transformation weights, a vector
%   nonGaussianityIndicator indicating non-Gaussianity type and an object
%   modelParams of type ReconstructionICAParams and creates an
%   implementation for RICA. N is the number of observations and P is the
%   number of predictors.

            % 1. How many inputs and how many outputs?
            this.NumPredictors      = size(X,2);
            this.NumLearnedFeatures = Q;

            % 2. Initial transformation weights.
            this.InitialTransformWeights = transformWeights0;

            % 3. Set the type of non-Gaussianity.
            this.NonGaussianityIndicator = nonGaussianityIndicator;

            % 4. Set model parameters.
            this.ModelParams = modelParams;

            % 5. Train - sets Mu, Sigma, TransformWeights and FitInfo.
            this = doRICA(this,X,Q);
        end
    end
    
    methods
        function Fhat = transform(this,X)
%transform - Apply learned transformation to new data.
%   Fhat = transform(this,X) takes object of type ReconstructionICAImpl and
%   a N-by-P matrix of predictors X (unstandardized) where N is the number
%   of observations and P is the number of predictors and returns a N-by-Q
%   matrix Fhat of new features. Q is the number of new features.
            
            % 1. If standardization was used, apply that transformation to
            % X as well.
            mu          = this.Mu(:);
            sigma       = this.Sigma(:);
            standardize = this.ModelParams.Standardize;

            if ( standardize && ~isempty(mu) && ~isempty(sigma) )
                X = bsxfun(@rdivide,bsxfun(@minus,X,mu'),sigma');
            end

            % 2. Get learned transformation weights. W is P-by-Q.
            W = this.TransformWeights;

            % 3. Compute new features Fhat - this will be N-by-Q.
            Fhat = featlearn.impl.ReconstructionICAImpl.computeFHat(X,W);
        end
    end
    
    methods
        function this = doRICA(this,X,Q)
            % 1. Standardize X and set Mu and Sigma.
            if this.ModelParams.Standardize
                [X,this.Mu,this.Sigma] = classreg.learning.gputils.standardizeData(X);
            else
                this.Mu    = [];
                this.Sigma = [];
            end
            
            % 2. X is a N-by-P matrix where N is the number of observations
            % and P is the number of predictors. Transpose X to make it
            % P-by-N for efficiency purposes.
            X = X';
            P = size(X,1);
            
            % 3. Make RICA objective function for minimization.
            lambda      = this.ModelParams.Lambda;
            theta       = this.ModelParams.Theta;
            contrastfcn = this.ModelParams.ContrastFcn;
            delta       = this.Delta;
            ng          = this.NonGaussianityIndicator;
            
            switch this.ModelParams.ObjectiveVersion
                case 1
                    obj = featlearn.impl.ReconstructionICAImpl.makeRICAObjectiveReference(X,Q,lambda,theta,contrastfcn,delta,ng);
                case 2
                    obj = featlearn.impl.ReconstructionICAImpl.makeRICAObjective(X,Q,lambda,theta,contrastfcn,delta,ng);
                otherwise
                    error(message('stats:featlearn:impl:ReconstructionICAImpl:UnsupportedObjectiveFunctionVersion'));
            end
            
            % 4. Initialize solution.
            initialTransformWeights = this.InitialTransformWeights;
            
            if isempty(initialTransformWeights)
                v0 = randn(P*Q,1);
            else
                v0 = initialTransformWeights(:);
            end
            
            % 5. Do minimization of obj starting at v0.
            iterationLimit    = this.ModelParams.IterationLimit;
            verbose           = this.ModelParams.VerbosityLevel;
            gradientTolerance = this.ModelParams.GradientTolerance;
            stepTolerance     = this.ModelParams.StepTolerance;
            
            [vHat,fitInfo] = featlearn.utils.Solver.doMinimization(obj,v0,iterationLimit,verbose,gradientTolerance,stepTolerance);
            
            % 6. Reshape vHat into a matrix.
            V = reshape(vHat,P,Q);
            
            % 7. Get unit norm weights. 
            this.TransformWeights = featlearn.impl.ReconstructionICAImpl.getUnitNormTransformWeights(V,delta);
            
            % 8. Create a structure to hold fitting information.
            this.FitInfo = struct('Iteration',fitInfo.iter,...
                                  'Objective',fitInfo.fval);
        end
    end
    
    methods(Static)
        function Fhat = computeFHat(X,W)
        % X    = N-by-P matrix.
        % W    = P-by-Q matrix.
        % Fhat = N-by-Q matrix.
        %
        % N = number of observations.
        % P = number of original predictors.
        % Q = number of new features.
            
            % Compute N-by-Q matrix of new features.
            Fhat = X*W;
        end
        
        function Fhat = computeFHatReference(X,W)
        % X    = P-by-N matrix.
        % W    = P-by-Q matrix.
        % Fhat = Q-by-N matrix.
        %
        % N = number of observations.
        % P = number of original predictors.
        % Q = number of new features.
            
            % Compute Q-by-N matrix of new features.
            Fhat = W'*X;
        end
        
        function [W,V_col_norms] = getUnitNormTransformWeights(V,delta)
            V_col_norms = sqrt(sum(V.^2,1) + delta);
            W           = bsxfun(@rdivide,V,V_col_norms);
        end
        
        function z = makeRICAObjectiveReference(X,Q,lambda,theta,contrastName,delta,ng)
        % X      = P-by-N matrix of predictors.
        % Q      = number of transformed features required.
        % lambda = Penalty for violating orthogonality constraint on columns of W.
        % theta  = Penalty for violating unit length constraint on columns of W.
        %
        % contrastName is the name of a contrast function.
        %
        % delta is usually set to 1e-8.
        %
        % ng is a Q-by-1 non-Gaussianity indicator vector. +1 means
        % super-Gaussian and -1 means sub-Gaussian.
        
        [P,N] = size(X);
        ng    = -ng;
        
        contrastFcn = featlearn.utils.ICAContrastFunctions.makeICAContrastFcn(contrastName);
            
            z = @myricaobjref;
            function [h,gh] = myricaobjref(vvec)
            % Let W = P-by-Q weight matrix such that:
            %
            % W = [w1,w2,...,wQ]
            %
            % where each wi is a P-by-1 vector. We require each column of W to have
            % unit norm. To do this, suppose,
            %
            % V = [v1,v2,...,vQ] is a P-by-Q matrix such that:
            %
            % wk = vk/sqrt(vk'*vk + delta) where delta = 1e-8.
            %
            % The input vvec to this function is a P*Q-by-1 vector created using
            % V(:). vvec is used to create W.
                
                % 1. Form P-by-Q matrix V.
                V = reshape(vvec,P,Q);
                
                % 2. Get W from V.
                [W,V_col_norms] = featlearn.impl.ReconstructionICAImpl.getUnitNormTransformWeights(V,delta);
                
                % 3. Form Z = WTX = Q-by-N matrix.
                Z = W'*X;
                
                % 4. Apply ICA non-linearity g(Z) = Q-by-N matrix. Also get
                % g'(Z) if gradient needs to be computed. Multiply gZ by ng
                % to get super or sub Gaussian contrast functions.
                if ( nargout > 1 )
                    [gZ,gpZ] = contrastFcn(Z);
                    gZ       = bsxfun(@times,gZ,ng);
                    gpZ      = bsxfun(@times,gpZ,ng);
                else
                    gZ  = contrastFcn(Z);
                    gZ  = bsxfun(@times,gZ,ng);
                end
                
                % 5. Compute W*Z = P-by-N matrix.
                S = W*Z;
                
                % 6. Compute objective.
                sumW2minus1 = sum(W.^2,1)-1;
                h           = (lambda)*sum(sum((S - X).^2)) + sum(sum(gZ)) + N*theta*sum(sumW2minus1.^2);
                
                % 7. Compute gradient if needed.
                if ( nargout > 1 )
                    % 7.1 Get the gradient.
                    gh  = (2*lambda)*(X*(S'*W - 2*Z') + S*Z') + (X*gpZ') + 4*N*theta*bsxfun(@times,W,sumW2minus1);
                    
                    % 7.2 We want gradient w.r.t elements of V not W.
                    gh = bsxfun(@rdivide,gh - bsxfun(@times,W,sum(gh.*W,1)),V_col_norms);
                    
                    % 7.3 Make gh into column vector.
                    gh = gh(:);
                end
            end
        end
    
        function z = makeRICAObjective(X,Q,lambda,theta,contrastName,delta,ng)
        % X      = P-by-N matrix of predictors.
        % Q      = number of transformed features required.
        % lambda = Penalty for violating orthogonality constraint on columns of W.
        % theta  = Penalty for violating unit length constraint on columns of W.
        %
        % contrastName is the name of a contrast function.
        %
        % delta is usually set to 1e-8.
        %
        % ng is a Q-by-1 non-Gaussianity indicator vector. +1 means
        % super-Gaussian and -1 means sub-Gaussian.
        
        [P,N] = size(X);
        C     = X*X';
        trC   = full(sum(diag(C)));
        ng    = -ng;
        
        contrastFcn = featlearn.utils.ICAContrastFunctions.makeICAContrastFcn(contrastName);
            
            z = @myricaobj;
            function [h,gh] = myricaobj(vvec)
            % Let W = P-by-Q weight matrix such that:
            %
            % W = [w1,w2,...,wQ]
            %
            % where each wi is a P-by-1 vector. We require each column of W to have
            % unit norm. To do this, suppose,
            %
            % V = [v1,v2,...,vQ] is a P-by-Q matrix such that:
            %
            % wk = vk/sqrt(vk'*vk + delta) where delta = 1e-8.
            %
            % The input vvec to this function is a P*Q-by-1 vector created using
            % V(:). vvec is used to create W.

                % 1. Form P-by-Q matrix V.
                V = reshape(vvec,P,Q);

                % 2. Get W from V.
                [W,V_col_norms] = featlearn.impl.ReconstructionICAImpl.getUnitNormTransformWeights(V,delta);

                % 3. Apply ICA non-linearity g(Z) = Q-by-N matrix. Also get
                % g'(Z) if gradient needs to be computed. Multiply gZ by ng
                % to get super or sub Gaussian contrast functions.
                if ( nargout > 1 )
                    [gZ,gpZ] = contrastFcn(W'*X);
                    gZ       = bsxfun(@times,gZ,ng);
                    gpZ      = bsxfun(@times,gpZ,ng);
                else
                    gZ  = contrastFcn(W'*X);
                    gZ  = bsxfun(@times,gZ,ng);
                end

                % 4. Compute objective.
                CW          = C*W;
                CWWt        = CW*W';
                WWtCW       = CWWt'*W;
                sumW2minus1 = sum(W.^2,1)-1;
                h           = (lambda)*(trC - 2*sum(diag(CWWt)) + sum(sum(W.*WWtCW,1))) + sum(sum(gZ)) + N*theta*sum(sumW2minus1.^2);

                % 5. Compute gradient if needed.
                if ( nargout > 1 )
                    % 5.1 Get the gradient.
                    gh  = (2*lambda)*(-2*CW + CWWt*W + WWtCW) + (X*gpZ') + 4*N*theta*bsxfun(@times,W,sumW2minus1);

                    % 5.2 We want gradient w.r.t elements of V not W.
                    gh = bsxfun(@rdivide,gh - bsxfun(@times,W,sum(gh.*W,1)),V_col_norms);

                    % 5.3 Make gh into column vector.
                    gh = gh(:);
                end
            end
        end
    end
end

