DMRGKron.hpp

#ifndef __DMRG_KRON_HPP
#define __DMRG_KRON_HPP

#include <petscsys.h>
#include <slepceps.h>

#include "DMRGBlock.hpp"
#include "Hamiltonians.hpp"

typedef std::tuple<PetscReal, PetscInt, PetscInt, PetscInt> KronBlock_t;

PETSC_EXTERN PetscErrorCode MatMult_KronSumShell(Mat A, Vec x, Vec y);

struct KronSumTerm {
    PetscScalar a;
    Op_t OpTypeA;
    Mat A;
    Op_t OpTypeB;
    Mat B;
};

struct KronSumCtx {
    PetscInt rstart=0;  
    PetscInt rend=0;    
    PetscInt lrows=0;   
    PetscInt cstart=0;  
    PetscInt cend=0;    
    PetscInt lcols=0;   
    PetscInt Nrows=0;   
    PetscInt Ncols=0;   
    PetscInt NReqRowsL, NReqRowsR;

    std::vector< PetscInt > ReqRowsL, ReqRowsR;

    std::unordered_map< PetscInt, PetscInt > MapRowsL, MapRowsR;

    std::vector< KronSumTerm > Terms;

    std::vector< Mat* > LocalSubMats;

    PetscInt *Dnnz;

    PetscInt *Onnz;

    PetscInt MinIdx=0;

    PetscInt MaxIdx=0;

    std::vector< PetscInt > Maxnnz;

    PetscInt Nfiltered=0;

    PetscInt Nnz=0;

    PetscInt Nelts=0;
};

struct KronSumTermRow {
    PetscInt    nz_L, nz_R, bks_L, col_NStatesR, fws_O;
    PetscInt    *idx_L, *idx_R;
    PetscScalar *v_L, *v_R;
};

struct KronSumShellCtx {

    KronSumCtx      *p_ctx;

    KronSumTermRow  *kstr;
    PetscInt        Nterms;
    PetscInt        Nrowterms;

    PetscInt        *Rows_L;
    PetscInt        *Rows_R;
    PetscScalar     *term_a;

    PetscScalar     one;

    /* Mat-Vec multiplication */
    VecScatter      vsctx;
    Vec             x_seq;
};

PetscErrorCode MatDestroy_KronSumShell(Mat *p_mat);

class KronBlocks_t
{

friend class KronBlocksIterator;

public:

    KronBlocks_t(
        Block::SpinBase& LeftBlock,
        Block::SpinBase& RightBlock,
        const std::vector<PetscReal>& QNSectors, 
        FILE *fp_prealloc,
        const PetscInt& GlobIdx
        ):
        GlobIdx(GlobIdx),
        LeftBlock(LeftBlock),
        RightBlock(RightBlock),
        fp_prealloc(fp_prealloc)
    {
        /* Require blocks to be initialized */
        if(!LeftBlock.Initialized()) throw std::runtime_error("Left input block not initialized.");
        if(!RightBlock.Initialized()) throw std::runtime_error("Right input block not initialized.");
        /* Fill in mpi information */
        mpi_comm = LeftBlock.MPIComm();
        if(mpi_comm != RightBlock.MPIComm())
            throw std::runtime_error("Left and right blocks must have the same communicator.");
        if(MPI_Comm_rank(mpi_comm, &mpi_rank)) throw std::runtime_error("MPI Error.");
        if(MPI_Comm_size(mpi_comm, &mpi_size)) throw std::runtime_error("MPI Error.");

        if(QNSectors.size() == 0)
        {
            for (size_t IL = 0; IL < LeftBlock.Magnetization.List().size(); ++IL){
                for (size_t IR = 0; IR < RightBlock.Magnetization.List().size(); ++IR){
                    KronBlocks.push_back(std::make_tuple(
                        LeftBlock.Magnetization.List()[IL] + RightBlock.Magnetization.List()[IR], IL, IR,
                        LeftBlock.Magnetization.Sizes()[IL] * RightBlock.Magnetization.Sizes()[IR]));
                }
            }
            /*  Sort by descending quantum numbers */
            std::stable_sort(KronBlocks.begin(), KronBlocks.end(), DescendingQN);
        }
        else if (QNSectors.size() == 1)
        {
            for (size_t IL = 0; IL < LeftBlock.Magnetization.List().size(); ++IL){
                for (size_t IR = 0; IR < RightBlock.Magnetization.List().size(); ++IR){
                    PetscReal QN = LeftBlock.Magnetization.List()[IL] + RightBlock.Magnetization.List()[IR];
                    if(QN == QNSectors[0])
                        KronBlocks.push_back(std::make_tuple(
                            QN, IL, IR,
                            LeftBlock.Magnetization.Sizes()[IL] * RightBlock.Magnetization.Sizes()[IR]));
                }
            }
        }
        else
        {
            /* Convert QNSectors into a set */
            std::set< PetscReal > QNSectorsSet(QNSectors.begin(), QNSectors.end());
            for (size_t IL = 0; IL < LeftBlock.Magnetization.List().size(); ++IL){
                for (size_t IR = 0; IR < RightBlock.Magnetization.List().size(); ++IR){
                    PetscReal QN = LeftBlock.Magnetization.List()[IL] + RightBlock.Magnetization.List()[IR];
                    if(QNSectorsSet.find(QN) != QNSectorsSet.end())
                        KronBlocks.push_back(std::make_tuple(
                            QN, IL, IR,
                            LeftBlock.Magnetization.Sizes()[IL] * RightBlock.Magnetization.Sizes()[IR]));
                }
            }
        }

        /* Fill-in size and offset data from the sorted blocks */
        num_blocks = KronBlocks.size();
        kb_list.reserve(num_blocks);
        kb_size.reserve(num_blocks);
        kb_offset.reserve(num_blocks+1);

        /* Unload the data into separate vectors */
        for(KronBlock_t kb: KronBlocks) kb_list.push_back(std::get<0>(kb));
        for(KronBlock_t kb: KronBlocks) kb_size.push_back(std::get<3>(kb));

        PetscInt idx = 0;
        for(KronBlock_t kb: KronBlocks)
            kb_map[std::make_tuple( std::get<1>(kb), std::get<2>(kb) )] = idx++;

        PetscInt sum = 0;
        for(KronBlock_t kb: KronBlocks)
        {
            kb_offset.push_back(sum);
            sum += std::get<3>(kb);
        }
        kb_offset.push_back(sum);
        num_states = sum;

        do_saveprealloc = PetscBool(fp_prealloc!=NULL);
        MPIU_Allreduce(MPI_IN_PLACE, &do_saveprealloc, 1, MPI_INT, MPI_SUM, PETSC_COMM_WORLD);
    };

    PetscInt size() const { return PetscInt(KronBlocks.size()); }

    const std::vector<KronBlock_t>& data() const { return KronBlocks; }

    KronBlock_t data(size_t idx) const { return KronBlocks[idx]; }

    KronBlock_t operator[](size_t idx) const { return KronBlocks[idx]; }

    std::vector<PetscReal> List() const { return kb_list; }

    std::vector<PetscInt> Offsets() const { return kb_offset; }

    PetscInt Offsets(const PetscInt& idx ) const {
        assert(idx >= 0 && idx < num_blocks + 1);
        return kb_offset[idx];
    }

    PetscReal QN(const PetscInt& idx) const {
        return std::get<0>(KronBlocks[idx]);
    }

    PetscInt LeftIdx(const PetscInt& idx) const {
        return std::get<1>(KronBlocks[idx]);
    }

    PetscInt RightIdx(const PetscInt& idx) const {
        return std::get<2>(KronBlocks[idx]);
    }

    PetscInt Sizes(const PetscInt& idx) const {
        return std::get<3>(KronBlocks[idx]);
    }

    const Block::SpinBase& LeftBlockRef() const { return LeftBlock; }

    const Block::SpinBase& RightBlockRef() const { return RightBlock; }

    Block::SpinBase& LeftBlockRefMod(){ return LeftBlock; }

    Block::SpinBase& RightBlockRefMod(){ return RightBlock; }

    PetscInt Offsets(const PetscInt& lidx, const PetscInt& ridx) const {
        PetscInt idx = Map(lidx, ridx);
        if(idx >= 0) return kb_offset[idx];
        else return -1;
    }

    std::vector<PetscInt> Sizes() const{ return kb_size; }

    PetscInt Map(const PetscInt& lidx, const PetscInt& ridx) const
    {
        auto tup = std::make_tuple(lidx,ridx);
        auto kb_find = kb_map.find(tup);
        if(kb_find != kb_map.end())
        {
            return kb_find->second;
        } else
        {
            return -1;
        }
    }

    PetscInt NumStates() const { return num_states; }

    PetscErrorCode KronSumConstruct(
        const std::vector< Hamiltonians::Term >& Terms, 
        Mat& MatOut                                     
        );

    PetscErrorCode KronConstruct(
        const Mat& Mat_L,
        const Op_t& OpType_L,
        const Mat& Mat_R,
        const Op_t& OpType_R,
        Mat& MatOut
        );

    PetscErrorCode KronSumSetShellMatrix(const PetscBool& do_shell_in)
    {
        do_shell = do_shell_in;
        return(0);
    }

    PetscErrorCode KronSumSetRedistribute(
        const PetscBool& do_redistribute_in = PETSC_TRUE
        )
    {
        do_redistribute = do_redistribute_in;
        return(0);
    }

    PetscErrorCode KronSumSetToleranceFromOptions()
    {
        PetscErrorCode ierr;
        ierr = PetscOptionsGetReal(NULL,NULL,"-ks_tol",&ks_tol,NULL); CHKERRQ(ierr);
        return(0);
    }

private:

    MPI_Comm mpi_comm = PETSC_COMM_SELF;
    PetscMPIInt mpi_rank, mpi_size;

    const PetscInt GlobIdx;

    std::vector<KronBlock_t> KronBlocks;

    std::vector<PetscReal> kb_list;

    std::vector<PetscInt> kb_size;

    std::vector<PetscInt> kb_offset;

    std::map< std::tuple<PetscInt,PetscInt>, PetscInt > kb_map;

    PetscInt num_blocks = 0;

    PetscInt num_states = 0;

    Block::SpinBase& LeftBlock;

    Block::SpinBase& RightBlock;

    FILE *fp_prealloc;

    PetscBool do_saveprealloc = PETSC_FALSE;

    PetscBool do_redistribute = PETSC_FALSE;

    PetscBool do_shell = PETSC_FALSE;

    PetscBool called_KronSumShellSplitOwnership = PETSC_FALSE;

    PetscInt  prev_lrows = 0;

    PetscInt  prev_rstart = 0;

    #if defined(PETSC_USE_REAL_DOUBLE)
    PetscReal ks_tol = 1.0e-16;
    #else
    #error Only double precision real numbers supported.
    #endif

    static bool DescendingQN(const KronBlock_t& a, const KronBlock_t& b)
    {
        return (std::get<0>(a)) > (std::get<0>(b));
    }

    PetscErrorCode MatKronEyeAdd(
        const std::vector< Mat >& Matrices,
        const Side_t& SideType,
        Mat& MatOut
        );

    PetscErrorCode VerifySzAssumption(
        const std::vector< Mat >& Matrices,
        const Side_t& SideType
        );

    PetscErrorCode KronGetSubmatrices(
        const Mat& Mat_L,
        const Op_t& OpType_L,
        const Mat& Mat_R,
        const Op_t& OpType_R,
        KronSumCtx& ctx
        );

    PetscErrorCode KronSumConstructExplicit(
        const std::vector< Hamiltonians::Term >& TermsLR,
        Mat& MatOut
        );

    PetscErrorCode KronSumConstructShell(
        const std::vector< Hamiltonians::Term >& TermsLR,
        Mat& MatOut
        );

    PetscErrorCode KronSumSetUpShellTerms(
        KronSumShellCtx *shellctx
        );

    PetscErrorCode KronSumGetSubmatrices(
        const Mat& OpProdSumLL,
        const Mat& OpProdSumRR,
        const std::vector< Hamiltonians::Term >& TermsLR,
        KronSumCtx& SubMat
        );

    PetscErrorCode KronSumCalcPreallocation(
        KronSumCtx& ctx
        );

    PetscErrorCode KronSumShellSplitOwnership(
        const Mat& OpProdSumLL,
        const Mat& OpProdSumRR,
        const std::vector< Hamiltonians::Term >& TermsLR,
        const PetscInt Nrows,
        PetscInt& lrows,
        PetscInt& rstart
        );

    PetscErrorCode KronSumRedistribute(
        KronSumCtx& ctx,
        PetscBool& flg
        );

    PetscErrorCode KronSumPreallocate(
        KronSumCtx& ctx,
        Mat& MatOut
        );

    PetscErrorCode KronSumFillMatrix(
        KronSumCtx& ctx,
        Mat& MatOut
        );

    PetscErrorCode SavePreallocData(const KronSumCtx& ctx);
};


PetscErrorCode KronEye_Explicit(
    Block::SpinBase& LeftBlock,          
    Block::SpinBase& RightBlock,         
    const std::vector< Hamiltonians::Term >& Terms, 
    Block::SpinBase& BlockOut            
    );

class KronBlocksIterator
{
public:

    typedef KronBlocksIterator Self_t;

    KronBlocksIterator(
        const KronBlocks_t& KronBlocks,
        const PetscInt& GlobIdxStart,
        const PetscInt& GlobIdxEnd
        ):
        KronBlocks(KronBlocks),
        istart_(GlobIdxStart),
        iend_(GlobIdxEnd),
        idx_(istart_)
    {
        if(istart_==iend_) {}
        else
        {
            kb_size.reserve(KronBlocks.size());
            kb_offset.reserve(KronBlocks.size()+1);

            for(KronBlock_t kb: KronBlocks.data()) kb_size.push_back(std::get<3>(kb));

            PetscInt sum = 0;
            for(KronBlock_t kb: KronBlocks.data())
            {
                kb_offset.push_back(sum);
                sum += std::get<3>(kb);
            }
            kb_offset.push_back(sum);
            num_states = sum;
            assert(istart_ < sum);

            while(idx_ >= kb_offset[blockidx_+1]) ++blockidx_;
        }
    }

    PetscInt IdxStart() const {return istart_;}

    PetscInt IdxEnd() const {return iend_;}

    PetscInt Idx() const {return idx_;}

    PetscInt BlockIdx() const {return blockidx_;}

    PetscInt LocIdx() const {return idx_ - kb_offset[blockidx_];}

    PetscInt BlockStartIdx(
        const PetscInt& BlockShift  
        ) const
    {
        PetscInt BlockIdx_out = blockidx_ + BlockShift;
        if(BlockIdx_out < 0 || BlockIdx_out >= num_states){
            return -1;
        }
        return kb_offset[BlockIdx_out];
    }

    PetscInt BlockSize(
        const PetscInt& BlockShift  
        ) const
    {
        PetscInt BlockIdx_out = blockidx_ + BlockShift;
        if(BlockIdx_out < 0 || BlockIdx_out >= num_states){
            return -1;
        }
        return kb_size[BlockIdx_out];
    }


    bool Loop() const {return idx_ < iend_;}

    PetscInt Steps() const {return idx_-istart_;}

    Self_t operator++()
    {
        ++idx_;
        if(idx_ >= kb_offset[blockidx_+1]){
            ++blockidx_;
            updated_block = PETSC_TRUE;
        } else {
            updated_block = PETSC_FALSE;
        }
        return *this;
    }

    PetscInt BlockIdxLeft() const {return std::get<1>(KronBlocks.data()[blockidx_]);}

    PetscInt LocIdxLeft() const { return LocIdx() / NumStatesRight(); }

    PetscInt NumStatesRight() const { return KronBlocks.RightBlock.Magnetization.Sizes()[BlockIdxRight()]; }

    PetscInt BlockIdxRight() const {return std::get<2>(KronBlocks.data()[blockidx_]);}

    PetscInt LocIdxRight() const { return LocIdx() % NumStatesRight(); }

    PetscInt GlobalIdxLeft() const {
        return KronBlocks.LeftBlock.Magnetization.BlockIdxToGlobalIdx(BlockIdxLeft(), LocIdxLeft());
    }

    PetscInt GlobalIdxRight() const {
        return KronBlocks.RightBlock.Magnetization.BlockIdxToGlobalIdx(BlockIdxRight(), LocIdxRight());
    }

    PetscBool UpdatedBlock() const {return updated_block; }

private:

    const KronBlocks_t& KronBlocks;

    PetscInt istart_ = 0;

    PetscInt iend_ = 0;

    PetscInt idx_ = 0;

    PetscInt blockidx_ = -1;

    PetscInt locidx_ = 0;

    std::vector<PetscInt> kb_size;

    std::vector<PetscInt> kb_offset;

    PetscInt num_states = 0;

    PetscBool updated_block = PETSC_TRUE;
};


#endif // __DMRG_KRON_HPP