classdef hnswSearcher < NeighborSearcher
%hnswSearcher Neighbor search object using a hnsw network.
%   A hnswSearcher object performs approximate KNN (K-nearest-neighbor)
%   search using Hierarchical Navigable Small World (HNSW) graphs.  
%   You can create a hnswSearcher object based on X using either of  
%   the following syntaxes:
%
%   CREATENS function:
%      NS = CREATENS(X,'NSMethod','hnsw')
%   hnswSearcher constructor:
%      NS = hnswSearcher(X)
%
%   Rows of X correspond to observations and columns correspond to
%   variables. When one of the above methods creates a hnswSearcher
%   object, it creates and saves HNSW graphs based on X. The KNNSEARCH
%   method uses the HNSW graph to find requested points in X to the query
%   points.
%
%   hnswSearcher properties:
%       X               - Data used to create the hnswSearcher object
%       Distance        - The distance metric
%       DistParameter   - The additional parameter for the distance
%                         metric
%       MaxNumLinksPerNode
%                       - The maximum number of connections or links that a
%                         single node can establish within one layer of
%                         HNSW graphs
%       TrainSetSize    - The size of the candidate list for potential
%                         nearest neighbors when constructing the HNSW
%                         graphs
%   hnswSearcher methods:
%       hnswSearcher    - Construct a hnswSearcher object
%       knnsearch       - Find approximate nearest neighbors for query points
%
%
%   Example:
%      % Create a hnswSearcher object for data X with the 'euclidean'
%      % distance.
%      X = randn(100,5);
%      ns = createns(X,'nsmethod','hnsw');
%
%      % Find 5 approximate nearest neighbors in X and the corresponding
%      % distance values for each point in Y.
%      Y = randn(25, 5);
%      [idx, dist] = knnsearch(ns,Y,'k',5);
%
%   See also CREATENS, ExhaustiveSearcher, KDTreeSearcher, STATS/KNNSEARCH.

%   References:
%     Yury Malkov and D. A. Yashunin. "Efficient and robust approximate
%     nearest neighbor search using Hierarchical Navigable Small World graphs."
%     In Proceedings of the 25th International Conference on World Wide Web
%     (WWW '16), pages 427-437, 2016.

%   Copyright 2023-2024 The MathWorks, Inc.


    properties (GetAccess='private', SetAccess='private')
        % The constructed search structure of HNSW graphs
        SearchStructure = [];

        % The number of layers in the constructed HNSW graphs
        NumLvs = 0;

        % The node ID of the entry point in the constructed HNSW graphs
        NodeIDEntryPoint = 0;

        % Number of nodes in each layer of the constructed HNSW graphs
        NumNodesPerLv = [];

        % Flag variable to control modification of properties Distance and
        % DistParameter
        IsDistModifiable = true;

        % Precomputed information for various distance metrics, used to
        % enhance search efficiency
        PrecomputedData
    end

    properties (Constant, Access = private)
        AllowedMethods = {'minkowski';'seuclidean'; 'mahalanobis'; ...
                     'euclidean';'cityblock'; 'chebychev'; ...
                     'cosine'; 'correlation'; ...
                    'spearman'; 'hamming'; 'jaccard' };
    end

    properties (GetAccess='public', SetAccess='private')
        %MaxNumLinksPerNode Maximum number of links per node.
        %   The MaxNumLinksPerNode property is a positive integer,
        %   indicating the maximum number of links or connections that a
        %   single node can have within one layer in the HNSW graphs.
        %
        %   See also hnswSearcher.
        MaxNumLinksPerNode = 16;

        %TrainSetSize The size of candidate list when constructing HNSW.
        %   The TrainSetSize property is a positive integer, indicating
        %   the size of candidate list of nearest neighbors for a single
        %   node when constructing HNSW.
        %
        %   See also hnswSearcher.
        TrainSetSize = 200;
    end

    methods (Access = protected)
        function this = setDistance(this, distMetric)
            if ~this.IsDistModifiable
                error(message('stats:ApproximateSearcher:hnswSearcher:DistModificationNotAllowed'));
            end

            if ischar(distMetric)
                distMetricValid = validatestring(distMetric, this.AllowedMethods, mfilename, 'Distance');
                this.PrivDistance = distMetricValid;
            else
                error(message('stats:ApproximateSearcher:hnswSearcher:BadDistance'));
            end
        end

        function this = setDistParameter(this, para)
            if ~this.IsDistModifiable
                error(message('stats:ApproximateSearcher:hnswSearcher:DistModificationNotAllowed'));
            end

            % Set DistParameter with the provided parameters. 
            % Use default values if 'para' is empty.
            switch this.Distance
                case 'seuclidean'
                    if isempty(para)
                        this.PrivDistParameter = std(this.X,[],1,"omitnan");
                        if any(this.PrivDistParameter == 0)
                            warning(message('stats:pdist2:ConstantColumns'));
                        end
                    else
                        nDims = size(this.X,2);
                        if ~(isvector(para) && length(para) == nDims...
                                && all(para >= 0))
                            error(message('stats:ApproximateSearcher:hnswSearcher:BadScale'));
                        end
                        if any(para == 0)
                              warning(message('stats:pdist2:ZeroInverseWeights'));
                        end
                        this.PrivDistParameter = para;
                    end
                    this.PrecomputedData = 1 ./ (this.PrivDistParameter .^ 2);
                    
                case 'minkowski'
                    if isempty(para)
                        this.PrivDistParameter = 2;
                    else
                        if ~(isscalar(para) && isnumeric(para) && para > 0 )
                            error(message('stats:ApproximateSearcher:hnswSearcher:BadMinExp'));
                        end
                        this.PrivDistParameter= para;
                    end
                case 'mahalanobis'
                    nDims = size(this.X,2);

                    if isempty(para)
                        if size(this.X,1)==1
                            error(message('stats:ApproximateSearcher:hnswSearcher:TooFewXRowsForMah'));
                        end
                        this.PrivDistParameter = cov(this.X,"omitrows");
                        [T,flag] = chol(this.PrivDistParameter);
                    else
                        if ~isequal(size(para),[nDims,nDims])
                            error(message('stats:ApproximateSearcher:hnswSearcher:BadCov'));
                        end
                        % Check if the covariance is symmetric and positive
                        [T, flag] = statslib.internal.cholcov(para, 0);

                        this.PrivDistParameter= para;
                    end
                    
                    if flag ~= 0
                        error(message('stats:ApproximateSearcher:hnswSearcher:SingularCov'));
                    end
                    
                    this.PrecomputedData = full(T \ eye(nDims));  % inv(T)
                otherwise
                    error(message('stats:ApproximateSearcher:hnswSearcher:InvalidDistanceParam'));
            end
        end

    end

    methods (Static, Access=protected)
        function [Xnorm,flag] = normalizeData(X, DistanceMetric)
            % Preprocess before normalization
            switch DistanceMetric
                case 'correlation'
                    X = X - mean(X, 2);
                case 'spearman'
                    X = statslib.internal.tiedrank(X')'; % treat rows as a series
                    X = X - (size(X,2) + 1) / 2; % subtract off the (constant) mean
            end
            
            % Rescale each row by largest element to prevent over/underflow, and
            % compute the 2-norm.
            Xmax = max(abs(X),[],2);
            Xnorm = X./Xmax;
            Xnorm = sqrt(sum(Xnorm.^2, 2));
            
            % The norm will be NaN for rows that are all zeros, fix that for the test
            % below.
            Xnorm(Xmax==0) = 0;
            
            % The norm will be NaN for rows of X that have any +/-Inf. Those should be
            % Inf, but leave them as is so those rows will not affect the test below.
            % The points can't be normalized, so any distances from them will be NaN
            % anyway.
            
            % Find points that are effectively zero relative to the point with largest norm.
            flag = any(Xnorm <= eps(max(Xnorm)));
            
            % Back out the rescaling, and normalize rows of X to have unit 2-norm.
            % Rows can't be normalized become all NaN.
            Xnorm = Xnorm .* Xmax;
            Xnorm = X./Xnorm;

            if flag
                switch DistanceMetric
                    case 'cosine'
                        warning(message('stats:pdist2:ZeroPoints'));
                    case 'correlation'
                        warning(message('stats:pdist2:ConstantPoints'));
                    case 'spearman'
                        warning(message('stats:pdist2:TiedPoints'));
                end
            end
        end % normalizeData
    end

    methods

        function obj = hnswSearcher(X, varargin)
        %hnswSearcher Construct a hnswSearcher object.
        %   NS = hnswSearcher(X,'NAME1',VALUE1,...,'NAMEN',VALUEN) creates
        %   a hnswSearcher object for the N-by-D data matrix X with one 
        %   row per observation and one column per predictor, specifying
        %   the following optional argument name/value pairs:
        %
        %     Name          Value
        %     'Distance'    A string specifying the default distance metric used
        %                   when you call the KNNSEARCH method. It can be one of
        %                   the following:
        %                     'euclidean'   - Euclidean distance (default)
        %                     'cityblock'   - City Block distance
        %                     'chebychev'   - Chebychev distance (maximum
        %                                     coordinate  difference)
        %                     'minkowski'   - Minkowski distance
        %                     'seuclidean'  - Standardized Euclidean 
        %                                     distance. 
        %                     'mahalanobis' - Mahalanobis distance.
        %                     'cosine'      - One minus the cosine of 
        %                                     the included angle between
        %                                     observations.
        %                     'correlation' - One minus the sample linear
        %                                     correlation between observations.
        %                     'spearman'    - One minus the sample Spearman's
        %                                     rank correlation between 
        %                                     observations.
        %                     'hamming'     - Hamming distance, percentage
        %                                     of coordinates that differ.
        %                     'jaccard'     - One minus the Jaccard 
        %                                     coefficient, the percentage 
        %                                     of nonzero coordinates
        %                                     that differ.
        %
        %    'P'            A positive scalar, indicating the exponent of
        %                   Minkowski distance. This argument is only
        %                   valid when 'Distance' is 'minkowski'. Default
        %                   is 2.
        %
        %    'Cov'          A positive definite matrix indicating the 
        %                   covariance matrix when computing the 
        %                   Mahalanobis distance. This argument is 
        %                   only valid when 'Distance' is 'mahalanobis'.
        %                   Default is COV(X,"omitrows").
        %
        %    'Scale'        A vector S containing non-negative values,
        %                   with length equal to the number of 
        %                   columns in X. Each coordinate difference 
        %                   between X and a query point is scaled by 
        %                   the corresponding element of S when computing 
        %                   the standardized Euclidean distance. This 
        %                   argument is only valid when 'Distance' is 
        %                   'seuclidean'. Default is STD(X,"omitnan").
        %
        %    'MaxNumLinksPerNode'
        %                   The maximum number of links or connections that
        %                   a single node can have within one layer in
        %                   the HNSW graphs, a positive integer. It must 
        %                   not exceed N. Default is min(16,N).
        %
        %    'TrainSetSize' The size of candidate list of nearest neighbors
        %                   for a single node when constructing the HNSW 
        %                   graphs, a positive integer. Larger values lead
        %                   to better graph quality, but slows the
        %                   construction. It must be at least
        %                   MaxNumLinksPerNode, and must not exceed N. 
        %                   Default is min(200,N).
        %
        %   See also  hnswSearcher, CREATENS.

            % 1. Sanity check for input data
            if nargin == 0
                error(message('stats:ApproximateSearcher:hnswSearcher:NoDataInput'));
            end

            if isempty(X)
                error(message('stats:ApproximateSearcher:hnswSearcher:EmptyDataNotSupported'));
            end

            % Tall data is not supported for hnswSearcher
            if ~isempty(find(cellfun(@istall, [{X},varargin(:)']), 1, 'first'))
                error(message('MATLAB:bigdata:array:FcnNotSupported', upper(mfilename)))
            end

            % gpuArrays are not supported
            if isa(X,'gpuArray')
                error(message('stats:ApproximateSearcher:hnswSearcher:GpuArrayData'));
            end

            % Non-float data, such as Integer/logical/char/anything data
            % will be converted double.
            if ~isfloat(X)
                warning(message('stats:ApproximateSearcher:hnswSearcher:DataConversion', class( X ), 'double'));
                X = double(X);
            end
            %  Complex floating point data can't be handled
            if ~isreal(X)
                error(message('stats:ApproximateSearcher:hnswSearcher:ComplexData'));
            end

            %convert sparse data to full, hnswSearcher can't handle
            %sparse data directly
            if issparse(X)
                warning(message('stats:ApproximateSearcher:hnswSearcher:DataConversion', 'sparse', 'full'));
                X = full(X);
            end

            obj.X = X;
            nX = size(X,1);

            % 2. Parse varargin
            pnames = {'distance' 'p' 'cov' 'scale' ...
                'MaxNumLinksPerNode' 'TrainSetSize'};
            dflts =  {'euclidean' [] []    [] ...
                min(16,nX)           min(200,nX)};
            [varargin{:}] = convertStringsToChars(varargin{:});
            [dist,minExp, mahaCov, seuScale, MaxNumLinksPerNode,TrainSetSize] = ...
                statslib.internal.parseArgs(pnames, dflts, varargin{:});

            % Set distance
            if isempty(dist)
                dist = 'euclidean' ;
                obj.PrivDistance = dist;
            else
                obj.Distance = dist;
            end

            % Set distance parameters
            if ~isempty(minExp) && ~strncmp(obj.Distance, 'minkowski', 3)
                error(message('stats:ApproximateSearcher:hnswSearcher:InvalidMinExp'));
            end
            if ~isempty(seuScale) && ~strncmp(obj.Distance, 'seuclidean', 3)
                error(message('stats:ApproximateSearcher:hnswSearcher:InvalidScale'));
            end
            if ~isempty(mahaCov) && ~strncmp(obj.Distance, 'mahalanobis', 3)
                error(message('stats:ApproximateSearcher:hnswSearcher:InvalidCov'));
            end

            switch obj.Distance
                case 'seuclidean'
                    obj.DistParameter = seuScale;
                case 'minkowski'
                    obj.DistParameter = minExp;
                case 'mahalanobis'
                    obj.DistParameter = mahaCov;
                otherwise
                    
            end

            % Sanity check for MaxNumLinksPerNode and TrainSetSize
            validateattributes(MaxNumLinksPerNode,{'numeric'},...
                {'scalar','real','integer','positive'},mfilename,'MaxNumLinksPerNode');
            validateattributes(TrainSetSize,{'numeric'},...
                {'scalar','real','integer','positive'},mfilename,'TrainSetSize');
            if MaxNumLinksPerNode>nX
                error(message('stats:ApproximateSearcher:hnswSearcher:BadMaxNumLinksPerNode'));
            end
            if TrainSetSize>nX || TrainSetSize<MaxNumLinksPerNode
                error(message('stats:ApproximateSearcher:hnswSearcher:BadTrainSetSize'));
            end
            % Convert MaxNumLinksPerNode to the type of X after validation
            MaxNumLinksPerNode = cast(MaxNumLinksPerNode, class(X));
            
            % Normalize input data for cosine, correlation, or spearman 
            % distance.
            if ismember(obj.Distance, {'cosine', 'correlation', 'spearman'})
                obj.PrecomputedData = hnswSearcher.normalizeData(X, obj.Distance);
            end

            % 3. Construct HNSW graphs
            normFactor = 1/log(max(MaxNumLinksPerNode,2)); % avoid inf level by setting a cutoff
            LevelEachSample=double(floor(-log(rand(nX,1))*normFactor));
            verbose = 0;

            switch obj.Distance
                case {'cosine', 'correlation', 'spearman'}
                    % Use normalized input data stored in PrecomputedData
                    % for 'cosine', 'correlation', and 'spearman' distances.
                    hnsw = internal.stats.constructHNSW(obj.PrecomputedData', LevelEachSample, ...
                        MaxNumLinksPerNode, obj.PrivDistance(1:3), TrainSetSize, ...
                        verbose, cast(obj.PrivDistParameter,'like',obj.PrecomputedData));
                case {'seuclidean', 'mahalanobis'}
                    % Use precomputed distance parameters stored in 
                    % PrecomputedData for 'seuclidean' and 'mahalanobis' 
                    % distances.
                    hnsw = internal.stats.constructHNSW(X', LevelEachSample, ...
                        MaxNumLinksPerNode, obj.PrivDistance(1:3), TrainSetSize, ...
                        verbose, cast(obj.PrecomputedData,'like',X));
                otherwise
                    % For other distance metrics, use default settings.
                    hnsw = internal.stats.constructHNSW(X', LevelEachSample, ...
                        MaxNumLinksPerNode, obj.PrivDistance(1:3), TrainSetSize, ...
                        verbose, cast(obj.PrivDistParameter,'like',X));
            end
            obj.SearchStructure = hnsw.Network;
            obj.NumLvs = hnsw.NumLvs;
            obj.NodeIDEntryPoint = hnsw.NodeIDEnterPoint;
            obj.NumNodesPerLv = hnsw.NumNodesPerLv;

            obj.MaxNumLinksPerNode = MaxNumLinksPerNode;
            obj.TrainSetSize = TrainSetSize;
            
            % Prohibit modification of Distance and DistanceParameter once 
            % search structure is constructed
            obj.IsDistModifiable = false;
            
        end % constructor

    end % methods block

    methods (Hidden)
        function rangesearch(~,~,varargin)
            error(message('stats:ApproximateSearcher:hnswSearcher:RangSearchNotAvailable'));
        end
    end

end % classdef
