QuantumNumbers.hpp

#ifndef __QUANTUM_NUMBERS_HPP__
#define __QUANTUM_NUMBERS_HPP__

#include <petscmat.h>
#include <vector>
#include <cassert>

/* For back-compatibility with versions that use single precision integers */
#if !defined(LLD)
    #if defined(PETSC_USE_64BIT_INDICES)
        #define LLD(INT) (INT)
    #else
        #define LLD(INT) ((long long)(INT))
    #endif
#endif

class QuantumNumbers
{

private:

    MPI_Comm mpi_comm = PETSC_COMM_SELF;

    PetscInt num_sectors;

    std::vector<PetscReal> qn_list;

    std::vector<PetscInt> qn_offset;

    std::vector<PetscInt> qn_size;

    PetscInt num_states;

    PetscBool initialized = PETSC_FALSE;

public:

    PetscErrorCode Initialize(
        const MPI_Comm& mpi_comm_in,                
        const std::vector<PetscReal>& qn_list_in,   
        const std::vector<PetscInt>& qn_size_in     
        );

    PetscErrorCode CheckInitialized() const
    {
        if(PetscUnlikely(!initialized))
            SETERRQ(mpi_comm, PETSC_ERR_ARG_CORRUPT, "QuantumNumbers object not yet initialized.");
        else
            return 0;
    }

    MPI_Comm MPIComm() const { return mpi_comm; }

    PetscInt NumSectors() const
    {
        assert(initialized);
        return num_sectors;
    }

    std::vector<PetscReal> List() const
    {
        assert(initialized);
        return qn_list;
    }

    const std::vector<PetscReal>& ListRef() const
    {
        assert(initialized);
        return qn_list;
    }

    PetscReal List(const PetscInt& idx) const
    {
        assert(initialized);
        if(0 <= idx && idx < num_sectors)
            return qn_list[idx];
        else
            return -1;
    }

    std::vector<PetscInt> Offsets() const
    {
        assert(initialized);
        return qn_offset;
    }

    PetscInt Offsets(const PetscInt& idx) const
    {
        assert(initialized);
        if(0 <= idx && idx < num_sectors)
            return qn_offset[idx];
        else
            return -1;
    }

    std::vector<PetscInt> Sizes() const
    {
        assert(initialized);
        return qn_size;
    }

    PetscInt Sizes(const PetscInt& idx) const
    {
        assert(initialized);
        if(0 <= idx && idx < num_sectors)
            return qn_size[idx];
        else
            return -1;
    }

    PetscInt NumStates() const
    {
        assert(initialized);
        return num_states;
    }

    PetscErrorCode BlockIdxToGlobalRange(
        const PetscInt& BlockIdx,   
        PetscInt& GlobIdxStart,     
        PetscInt& GlobIdxEnd        
        ) const;

    PetscErrorCode OpBlockToGlobalRange(
        const PetscInt& BlockIdx,   
        const PetscInt& BlockShift, 
        PetscInt& GlobIdxStart,     
        PetscInt& GlobIdxEnd,       
        PetscBool& flg              
        ) const;

    PetscInt OpBlockToGlobalRangeStart(
        const PetscInt& BlockIdx,   
        const PetscInt& BlockShift, 
        PetscBool& flg              
        ) const
    {
        PetscErrorCode ierr;
        PetscInt GlobIdxStart, GlobIdxEnd;
        ierr = OpBlockToGlobalRange(BlockIdx, BlockShift, GlobIdxStart, GlobIdxEnd, flg);
        assert(!ierr);
        return GlobIdxStart;
    }

    PetscErrorCode QNToGlobalRange(
        const PetscReal& QNValue,   
        PetscInt& GlobIdxStart,     
        PetscInt& GlobIdxEnd        
        ) const;

    PetscErrorCode GlobalIdxToBlockIdx(
        const PetscInt& GlobIdx,    
        PetscInt& BlockIdx          
        ) const;

    PetscErrorCode GlobalIdxToBlockIdx(
        const PetscInt& GlobIdx,    
        PetscInt& BlockIdx,         
        PetscInt& LocIdx            
        ) const;

    PetscErrorCode GlobalIdxToQN(
        const PetscInt& GlobIdx,    
        PetscReal& QNValue          
        ) const;

    PetscErrorCode BlockIdxToGlobalIdx(
        const PetscInt& BlockIdx,   
        const PetscInt& LocIdx,     
        PetscInt& GlobIdx           
        ) const;

    PetscInt BlockIdxToGlobalIdx(
        const PetscInt& BlockIdx,   
        const PetscInt& LocIdx      
        ) const;

    PetscErrorCode PrintQNs()
    {
        PetscErrorCode ierr;
        PetscMPIInt mpi_rank;
        ierr = MPI_Comm_rank(mpi_comm, &mpi_rank); CHKERRQ(ierr);
        if(!mpi_rank)
        {
            printf("[ ");
            for(PetscInt i = 0; i < num_sectors; ++i)
            {
                for(PetscInt j = 0; j < qn_size[i]; ++j) printf("%g ", qn_list[i]);
            }
            printf(" ]\n");
        }
        return(0);
    }

};


class QuantumNumbersIterator
{
private:

    const QuantumNumbers& QN;

    PetscInt istart_ = 0;

    PetscInt iend_ = 0;

    /* Stores the value of the current index */
    PetscInt idx_ = 0;

    PetscInt blockidx_ = 0;

    PetscInt locidx_ = 0;


public:

    typedef QuantumNumbersIterator Self_t;

    QuantumNumbersIterator(
        const QuantumNumbers& QN_in     
        ):
        QN(QN_in),
        iend_(QN.NumStates())
    {}

    QuantumNumbersIterator(
        const QuantumNumbers& QN_in,  
        const PetscInt& GlobIdxStart, 
        const PetscInt& GlobIdxEnd    
        ):
        QN(QN_in),
        istart_(GlobIdxStart),
        iend_(GlobIdxEnd),
        idx_(istart_)
    {
        PetscErrorCode ierr;
        if(istart_==iend_) {}
        else
        {
            ierr = QN.GlobalIdxToBlockIdx(istart_, blockidx_);
            assert(!ierr);
        }
    }

    /* TODO: Initialize an iterator through a range of quantum number blocks */

    PetscInt Idx() const {return idx_;}

    PetscInt BlockIdx() const {return blockidx_;}

    PetscInt IdxStart() const {return istart_;}

    PetscInt IdxEnd() const {return iend_;}

    PetscInt LocIdx() const {return idx_ - QN.Offsets()[blockidx_];}

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

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

    Self_t operator++()
    {
        ++idx_;
        if(idx_ >= QN.Offsets()[blockidx_+1]) ++blockidx_;
        return *this;
    }

    PetscErrorCode OpBlockToGlobalRange(
        const PetscInt& BlockShift,
        PetscInt& GlobIdxStart,
        PetscInt& GlobIdxEnd,
        PetscBool& flg
        ) const
    {
        PetscInt ierr;
        ierr = QN.OpBlockToGlobalRange(blockidx_, BlockShift, GlobIdxStart, GlobIdxEnd, flg); CHKERRQ(ierr);
        return(0);
    }

};


#endif // __QUANTUM_NUMBERS_HPP__