%CODISTRIBUTOR1D 1D codistributor object for codistributed arrays
%   The 1D codistributor distributes arrays along a single specified dimension,
%   the distribution dimension, in a non-cyclic, partitioned manner.
%   Conceptually, the overall global array could be reconstructed
%   by concatenating the various local parts along the distribution dimension.
%
%   Example: If an array is distributed across the second dimension, i.e., its
%   columns, each worker stores only some of the columns of the array.  The
%   Partition property of the codistributor specifies the number of columns
%   stored on each worker.  The Dimension property of 2 indicates distributed by
%   the second dimension, columns.
%
%   Example: Look at how a 999-by-999 matrix is distributed by its columns using
%   the default partition.
%       spmd
%           N = 999;
%           codistr = codistributor1d(2, codistributor1d.unsetPartition, [N,N]);
%           [e, f] = codistr.globalIndices(2);
%           fprintf('Worker %d stores columns %d:%d, a total of %d columns.\n', ...
%                    spmdIndex, e, f, codistr.Partition(spmdIndex) );
%       end
%   When run on 4 workers, this shows:
%       Worker 1 stores columns 1:250, a total of 250 columns.
%       Worker 2 stores columns 251:500, a total of 250 columns.
%       Worker 3 stores columns 501:750, a total of 250 columns.
%       Worker 4 stores columns 751:999, a total of 249 columns.
%
%   codistributor1d methods:
%     codistributor1d/codistributor1d  - Create 1D codistributor object
%     globalIndices    - Global indices for the local part corresponding to codistributor
%     isComplete       - Return true if codistributor has all of its information set
%     CELL    - Create codistributed cell array using codistributor
%     COLON   - Build codistributed arrays of the form j:d:k using codistributor
%     EYE     - Identity codistributed matrix using codistributor
%     FALSE   - False codistributed array using codistributor
%     INF     - Infinity codistributed array using codistributor
%     NAN     - Build codistributed array containing Not-a-Number using codistributor
%     ONES    - Ones codistributed array using codistributor
%     RAND    - codistributed array of uniformly distributed pseudorandom numbers using codistributor
%     RANDN   - codistributed array of normally distributed pseudorandom numbers using codistributor
%     SPALLOC - Allocate space for sparse codistributed matrix using codistributor
%     SPARSE  - Create sparse codistributed matrix using codistributor
%     SPEYE   - Overloaded to create a codistributed array
%     SPRAND  - Sparse uniformly distributed random codistributed matrix using codistributor
%     SPRANDN - Sparse normally distributed random codistributed matrix using codistributor
%     TRUE    - True codistributed array using codistributor
%     ZEROS   - Zeros codistributed array using codistributor
%
%   codistributor1d static methods and properties:
%     defaultPartition - Default partition across the workers
%     unsetDimension   - Value indicating unspecified distribution dimension
%     unsetPartition   - Value indicating unspecified distribution partition
%
%   codistributor1d object properties:
%     Dimension  - Distribution dimension of a codistributor
%     Partition  - Distribution partition of a codistributor
%
%   See also CODISTRIBUTOR, CODISTRIBUTED, CODISTRIBUTOR2DBC.

%   Copyright 2008-2022 The MathWorks, Inc.

classdef codistributor1d < TensorProductCodistributor
% This class maintains the following private invariants:
% 1) Sparsity of the local parts must match across all the workers
% 2) size(LocalPart, codistr.Dimension) must equal
%    codistr.Partition(spmdIndex).
% 3) In all dimensions other than codistr.Dimension, size(LocalPart, dim) must
%    be the same across all the workers.
% 4) The local parts must be the panels of the codistributed array:
%    The sum of size(LocalPart, codistr.Dimension) across the workers must equal
%    the global size of the codistributed array in that dimension.
% 5) All of the above must be consistent with the MATLAB convention that
%    for an array X, size(X, n) equals 1 if n > ndims(X).
%
% Consequences of the above:
%
% a) LocalPart must be of size 0 in the distribution dimension when
%    codistr.Partition(spmdIndex) == 0.
% b) Take a particular example: spmdSize is 2, global size is [5, 5],
%    codistr.Dimension is 3.  Then codistr.Partition is either [1, 0] or [0, 1].
%    If codistr.Partition is [1, 0], worker 1 has LocalPart of size [5, 5], and worker
%    2 has LocalPart of size [5, 5, 0].
% c) To generalize b), if codistr.Dimension exceeds the number of dimensions
%    of the codistributed array (i.e. length(codistr.Cached.GlobalSize),
%    the following must hold true:
%    i)   codistr.Partition is 1 on exactly one worker, 0 on all other workers.
%    ii)  On the worker where codistr.Partition is 1, ndims(LocalPart) equals the
%         number of dimensions of the codistributed array.
%    iii) On the worker where codistr.Partition is 0, ndims(LocalPart) equals
%         codistr.Dimension.
    properties(SetAccess = private)
        %Dimension  Distribution dimension of a codistributor
        %   codistr.Dimension is the distribution dimension of a codistributor1d object
        %   codistr.
        %
        %   Example: If the distribution dimension is 2, the corresponding
        %   codistributed array is distributed by columns (dimension 2), and
        %   each worker only stores some of the columns of the array.  The number
        %   of columns stored on each worker is given by codistr.Partition.
        %
        %   See also codistributor1d/Partition, codistributor1d/unsetDimension.
        Dimension

        %Partition  Distribution partition of a codistributor
        %   codistr.Partition is the distribution partition of a codistributor1d object
        %   codistr.
        %
        %   See also codistributor1d/Dimension, codistributor1d/unsetPartition,
        %   codistributor1d/defaultPartition.
        Partition
    end

    methods(Access = private)
        function localK = pConvertGlobalDiagToLocal(codistr, globalK, workerIdx)
        % This function computes the local diagonal number for worker number
        % workerIdx given the global diagonal number and the codistributor
        % of the matrix of interest.  The variable workerIdx can be a scalar
        % or a vector of worker indices.  If it is a vector, then localK will be
        % a column vector with the same length as workerIdx and globalK is
        % equivalent to localK(i) on worker workerIdx(i).
        %
        % This function is valid for any natural number distribution dimension
        % (codistr.Dimension == {1, 2, 3, 4, ...}) so long as
        % numel(codistr.Cached.GlobalSize) == 2.

            cumPart = [0; cumsum(codistr.Partition')];
            if codistr.Dimension == 2
                localK = globalK - cumPart(workerIdx);
            else
                localK = globalK + cumPart(workerIdx);
            end
        end % End of pConvertGlobalDiagToLocal
    end

    methods(Access = protected)
        function codistr = pTransposeCodistributor(codistr)
            sz = codistr.Cached.GlobalSize;
            dim = codistr.Dimension;

            if dim < 3
                % For dim = 1 or 2, (c)transpose modifies the
                % codistr dimension.
                dim = 3 - dim;
            end

            codistr = codistributor1d(dim, codistr.Partition, [sz(2), sz(1)]);
        end % End pTransposeCodistributor
        
        function attrs = pGetCodistAttributes(codistr)
            attrs = {class(codistr), codistr.Dimension, codistr.Partition};
        end
    end

    % Implementation of abstract instance methods.
    methods(Hidden = true)
        [LP, varargout] = hBicgImpl(codistr, varargin)
        [LP, varargout] = hBicgstabImpl(codistr, varargin)
        [LP, varargout] = hBicgstablImpl(codistr, varargin)
        [LP, codistr] = hBuildFromReplicatedImpl(codistr, srcWorker, X)
        tf = hCatCheck(codistr, catdim, cellOfCodistr, cellOfLPs)
        [LP, codistr] = hCatImpl(codistr, catdim, codistrs, ...
                                 LPs, templ, targetSize)

        function tf = hCell2MatCheck(codistr, LP) %#ok<INUSD>
            tf = true;
        end

        [LP, codistr] = hCell2MatImpl(codistr, LP)
        [LP, codistr] = hCell2StructImpl(codistr, LP, fields, dim)
        [LP, varargout] = hCgsImpl(codistr, varargin)

        [LP, codistr] = hColonImpl(codistr, a, d, b)
        
        function tf = hDoGatherBsxfunInput(codistr, szInput)
            % The second test below on the size of the output codistributor
            % in the distribution dimension is necessary because if both inputs
            % are singleton in the distribution dimension they should not be
            % gathered.  An example that requires gathering is if the first 
            % input array has size [2 1] distributed on two workers on the first 
            % dimension and the second input array has size [1 1].  In that 
            % case, the output codistributor will be [2 1] and the second input
            % array will be gathered so that it will be available for 
            % calculating one output from bsxfun on each worker.  On the other 
            % hand, if the first input has size [1 1], there will be only one 
            % output and neither input should be gathered.  This test is what 
            % prevents the gathering in that case. In this example the size of 
            % the second dimension of each input is arbitrary as long as they 
            % agree if both are ~= 1.
            szOut = codistr.Cached.GlobalSize;
            outputDistDim = codistr.Dimension;
            tf = outputDistDim <= length(szOut) && ...
                 szOut(outputDistDim) ~= 1 && ...
                 (outputDistDim > length(szInput) || szInput(outputDistDim) == 1);
        end

        function tf = hDiagCheck(codistr) %#ok<MANU>
            tf = true;
        end

        function tf = hDiffCheck(codistr) %#ok<MANU>
            tf = true;
        end
        [localX, xDist] = hDiffImpl(xDist, localX, dim)
        [vecLP, vecDist] = hDiagMatToVecImpl(codistr, matLP, k)
        [matLP, matDist] = hDiagVecToMatImpl(codistr, vecLP, k)
        function tf = hFftnCheck(~, ~)
             tf = true;
        end
        [LP, codistr] = hFftnImpl(codistr, LP, siz, sym, isFFT)
        localLinInd = hFindDiagElementsInLocalPart(codistr)
        X = hGatherImpl(codistr, LP, destWorker)

        function LP = hGetBsxfunLocalPartFromReplicated(~, replicatedInput) 
            LP = replicatedInput;
        end

        % hGetCodistForTopK Return codistributors for output of maxk, mink, topkrows
        [codistrL, codistrI] = hGetCodistForTopK(codistr, k, sortDim, distDim)

        % hGetColumn Return the column of a 2D matrix with global index 
        % globalColIndex
        column = hGetColumn(codistr, LP, globalColIndex)

        function codistr = hGetCompleteForSize(codistr, wantedGlobalSize)
        % Return a codistributor that is based on the input codistributor and matches
        % the desired global size.
            if ~codistr.isComplete() ...
                    || ~isequal(codistr.Cached.GlobalSize, wantedGlobalSize)
                codistr = codistributor1d(codistr.Dimension, ...
                                          codistr.Partition, ...
                                          wantedGlobalSize);
            end
        end

        function codistr = hGetCompleteSparseBalancedForSize(...
                codistr, wantedGlobalSize, nnz1D)
        % Return a codistributor that is based on codistr.Dimension 
        % and wantedGlobalSize. If the partition is not set yet, calculate 
        % a balanced partition based on nnz1D, where nnz1D is a vector 
        % containing the NNZ of each row (codistr.Dimension==1) or each 
        % column (codistr.Dimension==2). Input codistr is either not 
        % complete yet or its global size does not match the wantedGlobalSize
            part = codistr.Partition;
            if isequal(part, codistributor1d.unsetPartition)
                part = codistr.sparseBalancedPartition(...
                    wantedGlobalSize(codistr.Dimension), nnz1D);
            end
            codistr = codistributor1d(codistr.Dimension, part, ...
                wantedGlobalSize);
        end
                
        function codistr = hGetNewForSize(codistr, wantedGlobalSize)
        % Return a codistributor that is based on the input codistributor
        % and matches the desired global size.
            if ~codistr.isComplete() ...
                    || ~isequal(codistr.Cached.GlobalSize, wantedGlobalSize)
                dim = codistr.Dimension;
                codistr = codistributor1d(dim, ...
                                          codistributor1d.unsetPartition, ...
                                          wantedGlobalSize);
            end
        end

        function dims = hGetDimensions(codistr)
            dims = codistr.Dimension;
        end

        function [LP, codistr] = hGetRepmatCatDimensionOrdering(codistr, LP, sz, replFcn)
            distDim = codistr.Dimension;
            order = length(sz):-1:1;
            if distDim < length(sz)
                order = [distDim order(order~=distDim)];
            end
            for d = order
                if sz(d) ~= 1
                    LP = replFcn(LP, d, sz(d));
                end
            end
        end

        % hGetRow Return the row of a 2D matrix with global index 
        % globalRowIndex
        row = hGetRow(codistr, LP, globalRowIndex)

        varargout = hGlobalIndicesImpl(codistr, dim, workerIdx)
        localIdx = hGlobalToLocalIndices(codistr, globalIdx, dim, workerIdx)
        [LP, varargout] = hGmresImpl(codistr, varargin)
        varargout = hIluColImpl(codistr, LP, setup)
        varargout = hIluRowImpl(codistr, LP, setup) 
        tf = hIsGlobalIndexOnWorker(codistr, dim, gIndexInDim, workerIdx)
        szs = hLocalSize(codistr, workerIdx)
        globalIdx = hLocalToGlobalIndices(codistr, localIdx, dim, workerIdx)
        globalIdx = hLocalToGlobalLinearIndices(codistr, localIdx, workerIdx)
        [LP, varargout] = hLsqrImpl(codistr, varargin)

        [LPY, LPI, codistr] = hMinMaxImpl(codistr, fcnMinMax, LP, dim, wantI, linearI)
        [LP, varargout] = hMinresImpl(codistr, varargin)

        function tf = hMtimesCheck(codistr, areBothMatricesDense) %#ok<INUSD>
            tf = true;
        end

        [LP, codistr] = hMtimesImpl(codistr, codistrA, LPA, transA, ...
            codistrB, LPB, transB)

        function tf = hMtimesReplicatedCheck(codistr, areBothMatricesDense) %#ok<INUSD>
            tf = true;
        end

        [LP, codistr] = hMtimesReplicatedImpl(codistr, codistrA, LPA, transA, ...
            codistrB, LPB, transB)

        function tf = hNonzerosCheck(codistr) %#ok<MANU>
            tf = true;
        end

        [LP, codistr] = hNonzerosImpl(codistr, LP)

        function tf = hNum2CellWithDimCheck(codistr, LP, dim) %#ok<INUSD>
            tf = true;
        end

        [LP, codistr] = hNum2CellWithDimImpl(codistr, LP, dims)

        function numWorkers = hNumWorkers(codistr)
            numWorkers = numel(codistr.Partition);
        end

        [LP, varargout] = hPcgImpl(codistr, varargin)

        function tf = hPermuteCheck(codistr, lenOrder) %#ok<INUSD>
            tf = true;
        end

        [LP, codistr] = hPermuteImpl(codistr, A, order)
        
        [LP, varargout] = hQmrImpl(codistr, varargin)

        [LP, codistr] = hReductionOpAlongDimImpl(codistr, fcn, LP, dim)

        function tf = hReshapeCheck(codistr, newSz)
            [newDimension, newPartition] = pGetReshapeTarget(codistr, newSz);
            tf = isValidDistributionDimension(newDimension) && ...
                 isValidDistributionPartition(newPartition);
        end

        function [LP, codistr] = hReshapeImpl(codistr, LP, newSz)
            [newDimension, newPartition] = pGetReshapeTarget(codistr, newSz);
            codistr = codistributor1d(newDimension, newPartition, newSz);
            LP = reshape(LP, codistr.hLocalSize());
        end

        function tf = hSortArrayCheck(codistr) %#ok<MANU>
            tf = true;
        end % End of hSortArrayCheck

        [LP, LPI, codistr] = hSortArrayImpl(codistr, LP, trailingArgs, wantI)

        function tf = hSortCellStrCheck(codistr) %#ok<MANU>
            tf = true;
        end

        [LP, LPI, codistr] = hSortCellStrImpl(codistr, LP, trailingArgs, wantI)

        function tf = hSortVectorCheck(codistr) %#ok<MANU>
            tf = true;
        end

        [LP, LPI, codistr] = hSortVectorImpl(codistr, LP, trailingArgs, wantI)

        function tf = hSortrowsCheck(codistr)
            tf = codistr.Dimension == 2;
        end

        [LP, I, codistr] = hSortrowsDefaultImpl(codistr, LP, trailingArgs)
        [LP, I, codistr] = hSortrowsGeneralImpl(codistr, LP, localCols, trailingArgs)
        [LP, I, codistr] = hSortrowsOneColImpl(codistr, LP, localCols, trailingArgs)

        [LP, codistr] = hSpallocImpl(codistr, m, n, nzmx)
        [LP, codistr] = hSparsifyImpl(codistr, fcn, LP)

        [LP, codistr] = hStruct2CellImpl(codistr, LP)

        function tf = hSupportsDimensionality(codistr, ndims) %#ok<INUSD>
            tf = true;
        end

        function tf = hSupportsStruct2Cell(codistr) %#ok<MANU>
            tf = true;
        end
        
        B = hSubsrefParensImpl(codistr, A, aidx, varargin)
        
        [LP, varargout] = hSymmlqImpl(codistr, varargin)
        [LP, varargout] = hTfqmrImpl(codistr, varargin)

        [LP, codistr] = hTypecastImpl(codistr, LP, datatype)

        hVerifySupportsSparse(codistr)
    end

    properties (Constant = true)
        %unsetDimension Value indicating unspecified distribution dimension
        %   DIM = codistributor1d.unsetDimension returns the special
        %   value used by codistributor1d to indicate that a distribution
        %   dimension has not been specified.
        %
        %   Example:
        %       codistr = codistributor1d();
        %    returns a distribution scheme such that the value of
        %    codistr.Dimension equals codistributor1d.unsetDimension
        %
        %   See also codistributor1d, codistributor1d/unsetPartition,
        %   codistributor1d/defaultPartition.
        unsetDimension = 0;

        %unsetPartition Value indicating unspecified distribution partition
        %   PART = codistributor1d.unsetPartition returns the special
        %   value used by codistributor1d to indicate that a distribution
        %   partition has not been specified.
        %
        %   Example:
        %       codistr = codistributor1d();
        %    returns a distribution scheme such that the value of
        %    codistr.Partition equals codistributor1d.unsetPartition
        %
        %   See also codistributor1d, codistributor1d/defaultPartition,
        %   codistributor1d/unsetDimension.
        unsetPartition = [];
    end % End of constant properties.

    % Public, static functions.
    methods(Static = true)
        par = defaultPartition(len);
    end % End of public, static methods.

    % Sealed, hidden, static functions.
    methods(Static = true, Sealed = true, Hidden = true)
        % Balanced partition based on nnz per row/column as row vector
        par = sparseBalancedPartition(sizeInDistDimension, nnz1D)
    end % End of sealed, hidden, static methods.
    
    methods
        function dist = codistributor1d(varargin)
        %CODISTRIBUTOR1D   Create 1D codistributor object for codistributed arrays
        %   The 1D codistributor distributes arrays along a single, specified
        %   distribution dimension, in a non-cyclic, partitioned manner.
        %
        %   Each of the following forms a 1D codistributor object with its global
        %   size unspecified.  When provided, the distribution dimension and
        %   partition are set to DIM and PART, respectively, otherwise they are
        %   unspecified.  The resulting codistributor is incomplete because its
        %   global size is not specified.
        %   DIST = CODISTRIBUTOR1D()
        %   DIST = CODISTRIBUTOR1D(DIM)
        %   DIST = CODISTRIBUTOR1D(DIM, PART)
        %   A codistributor constructed in this manner can then be used as an argument to
        %   other functions as a template codistributor when creating codistributed
        %   arrays.
        %
        %   DIST = CODISTRIBUTOR1D(DIM, PART, GSIZE) forms a codistributor object with
        %   distribution dimension DIM, distribution partition PART, and global size
        %   of its codistributed arrays GSIZE.  The resulting codistributor object is
        %   complete and can therefore be used to build a codistributed array from its
        %   local parts with codistributed.build.
        %   If DIM equals codistributor1d.unsetDimension, the distribution dimension of
        %   DIST is derived from GSIZE and is set to be the last non-singleton
        %   dimension.  Similarly, if PART equals codistributor1d.unsetPartition, the
        %   distribution partition of DIST is set to be the default partition for that
        %   global size and distribution dimension.
        %
        %   The local part on worker IDX of a codistributed array using such a
        %   codistributor is of size GSIZE in all dimensions except DIM, where the
        %   size is PART(IDX). The local part has the same class and attributes
        %   as the overall codistributed array. Conceptually, the overall global
        %   array could be reconstructed by concatenating the various local parts
        %   along dimension DIM.
        %
        %   Example: Use a codistributor1d object to create an N-by-N matrix of ones,
        %   distributed by rows.
        %   N = 1000;
        %   spmd
        %       codistr = codistributor1d(1);  % 1 specifies first dimension (rows).
        %       D = codistributed.ones(N, codistr);
        %   end
        %
        %   Example: Use a fully specified codistributor1d object to create a trivial
        %   N-by-N codistributed matrix from its local parts.  Then visualize which
        %   elements are stored on worker 2.
        %   N = 1000;
        %   spmd
        %       codistr = codistributor1d(codistributor1d.unsetDimension, ...
        %                                 codistributor1d.unsetPartition, ...
        %                                 [N, N]);
        %       localSize = [N, N];
        %       localSize(codistr.Dimension) = codistr.Partition(spmdIndex);
        %       localPart = spmdIndex*ones(localSize);
        %       D = codistributed.build(localPart, codistr);
        %   end
        %   spy(D == 2);
        %
        %   See also CODISTRIBUTOR, CODISTRIBUTED, CODISTRIBUTOR2DBC.

            narginchk(0, 3);
            try
                [dim, part, gsize] = iParseConstructorArgs(varargin{:});
            catch e
                throw(e);
            end

            % Set the property so that we can benefit from all the error checking performed
            % in the parent class.
            dist.Cached.GlobalSize = gsize;
            % Get it again.  We then have a valid MATLAB size vector.
            gsize = dist.Cached.GlobalSize;
            % At this point, we have values for dim, part and gsize.  The values
            % are either the user-provided values or the raw default
            % values.  We can do better because if gsize is specified,
            % and either dimension or partition are unspecified, we
            % fill in those values.
            if ~isempty(gsize)
                if dim == codistributor1d.unsetDimension
                    dim = distributedutil.Sizes.lastNonSingletonDimension(gsize);
                end
                % Calculate the expected size in the distribution dimension.
                if dim <= length(gsize)
                    expSizeInDim = gsize(dim);
                else
                    % The distribution dimension exceeds the number of
                    % dimensions of the array.
                    expSizeInDim = 1;
                end
                if isequal(part, codistributor1d.unsetPartition)
                    % Set partition to be the default for this length in the
                    % distribution dimension.
                    part = codistributor1d.defaultPartition(expSizeInDim);
                else
                    % Error check partition against the global size.
                    if sum(part) ~= expSizeInDim
                        error(message('parallel:codistributors:c1dParMatchSizeDim'));
                    end
                end
            end
            dist.Dimension = dim;
            dist.Partition = part;
        end
    end
end

function [dim, part, gsize] = iParseConstructorArgs(dim, part, gsize)
% Parse the arguments passed to the codistributor1d constructor, perform basic
% error checking on them individually, and provide the default values.

% Error check dim, or fill in default value.
    if nargin >= 1
        dim = distributedutil.CodistParser.gatherIfCodistributed(dim);
        if ~isequal(dim,  codistributor1d.unsetDimension) ...
                && ~isValidDistributionDimension(dim)
            error(message('parallel:codistributors:c1dDistributionDimensionInput'))
        end
    else
        dim = codistributor1d.unsetDimension;
    end

    % Error check part, or fill in default value.
    if nargin >= 2
        part = distributedutil.CodistParser.gatherIfCodistributed(part);
        if ~isequal(part, codistributor1d.unsetPartition) ...
                && ~isValidDistributionPartition(part)
            error(message('parallel:codistributors:c1dDistributionPartitionInput'))
        end
    else
        part = codistributor1d.unsetPartition;
    end

    % Error check gsize, or fill in default value.
    if nargin >= 3
        gsize = distributedutil.CodistParser.gatherIfCodistributed(gsize);
        if ~isempty(gsize) && ~isPositiveIntegerValuedNumeric(gsize, true)
            error(message('parallel:codistributors:c1dSizeInput'))
        end
    else
        gsize = [];
    end
end % End of iParseConstructorArgs.
