Refactor&Feature: a new module Representation will be added for representation transform

[kirk0830]

Describe the Code Quality Issue

Overview

This issue describes a refactor task and provides possibility on general representation transform for future need.

• Feature On the need of postprocessing orbital information with methods like COHP (issue #2803 #1536 ) and wavefunction analysis (wannier, Quasi-atomic Orbital issue #3063), wavefunction is needed to transform from one representation to another,
• Refactor newly implemented psi_initializer provides possibility of initializing pw wavefunction with numerical atomic orbital, which is actually also a representation transform.

Expectations

These two above aspects can be sorted into one module, results would be:

1. representation transform will be available both before and after SCF calculation
2. before SCF the representation transform is for wavefunction initialization
3. after SCF the representation transform is for wavefunction postprocessing
4. psi_initializer class will still exist but transform from PAO (pseudo atomic-orbital read from pseudopotential) or NAO (numerical atomic orbital read from numerical orbital files) to pw will be moved to this new module, the random method will be kept in psi_initializer and initialization of wavefunction will still be done by psi_initializer.

Design

The design focuses on both workflow (developer use) and distribute develop aspects:

Workflow

Stage 1: wavefunction initialization

In ESolver_KS, create an instance of Representation:

    //esolver_ks.h
        ....
            Representation<T, Device> rep;
        ....

In ESolver_KS, configure Representation, linke Representation with UnitCell, Parallel_Kpoints, pseudopot_vnl, Structure_Factor, ModulePW::PW_Basis_K, etc.:

    //esolver_ks.cpp
        ESolver_KS<T, Device>::Init()
        {
            ....
            this->rep.configure(....);
            ....
        }

detailed implementation of configure() is simple:

template<typename T, typename Device>
void Representation<T, Device>::configure(Structure_Factor* sf_in, 
                                          ModulePW::PW_Basis_K* pw_wfc_in, 
                                          UnitCell* p_ucell_in, 
                                          Parallel_Kpoints* p_parakpts_in, 
                                          pseudopot_cell_vnl* p_pspot_nl_in)
{
    this->p_sf = sf_in;
    this->pw_wfc = pw_wfc_in;
    this->p_ucell = p_ucell_in;
    this->p_parakpts = p_parakpts_in;
    this->p_pspot_nl = p_pspot_nl_in;
}

When doing wavefunction initialization, set input representation according to GlobalV::init_wfc:

    //esolver_ks_pw.h (for pw, lcao wavefunction initialization is not supported yet)
        psi_initializer<T, Device>* psi_init = nullptr; // then allocate it according to init_wfc, 
                                                        // and, call initialize to psig, to check
                                                        // if complementary bands needed to add

    //psi_initiailizer_nao.cpp
        void psi_initializer<T, Device>::initialize(const psi::Psi<T, Device>& psi_in, psi::Psi<T, Device>& psi_out)
        {
            if(this->get_num_complem_band() > 0)
            {
                psi_out(iband, ibasis) = random(...)
            }
            if(GlobalV::init_wfc.find_first_of('+') != std::string::npos)
            {
                psi_out(iband, ibasis) = psi_in(iband, ibasis) + random(...)
                // NOTE the range of iband, no need to exceed the number of bands in psi_in, psi_in is actually psig of Representation instance
            }
            // the case random number is mixed with wavefunction
        }

after psi is initialized, clean Representation:

    //esolver_ks.cpp
        ESolver_KS<T, Device>::Run(...)
        {
            ....
            hamilt2density(...);
            this->rep.clean_representations();
            ....
        }

Stage 2: wavefunction transformation (postprocess)

Instance of Representation has been created and configured in ESolver_KS::Init(), therefore no need to do it again.
In ESolver_KS, transform psi:

    //esolver_ks_pw.cpp
        ESolver_KS_PW<T, Device>::postprocess(....)
        {
            ....
            if(GlobalV::out_qo)
            {
                this->rep.add_transform_pair("pw", "qo");
            }
            if(GlobalV::out_wannier)
            {
                this->rep.add_transform_pair("pw", "wannier");
            }
            ....
            //then output psi and refresh memory each time
            for(int iout = 0; iout < this->rep.representations["output"].size(); ++iout)
            {
                this->rep.transform(kspw_psi, psi_out);
                // or directly output to file?
            }
            this->rep.clean_representations();
        }

Distribute development

According to investigation, pw is feasible for COHP, QO, and easily to transform to GTO (analytical form very clean). Therefore all input psi will be transformed into pw representation first, called psig, this transform will be completed by class RepIn. Then psig is transformed to other representations (if needed) by class RepOut. Therefore RepIn and RepOut have functionalities from xxx to pw and from pw to xxx, can be implemented seperatedly. Example code:

//representation.h
template<typename T, typename Device>
class Representation
{
    public:
        /// @brief transform one psi in one representation to another representation
        /// @param psi_in input psi
        /// @param psi_out output psi
        void transform(const psi::Psi<T, Device>& psi_in, 
                             psi::Psi<T, Device>& psi_out);
        void set_repin(const std::string repin_name);
        void add_repout(const std::string repout_name);
        /// @brief add representation transformation pair
        /// @param repin_name the name of representation of input psi
        /// @param repout_name the name of representation of output psi
        /// @attention usage: add_transform_pair("pw", "wannier")
        void add_transform_pair(const std::string repin_name,
                                const std::string repout_name);
    private:
        /// @brief interface to functional class, RepIn, transforms psi from one representation to pw
        RepIn<T, Device>* repin;
        /// @brief interfaces to functional class, RepOut, transforms psi from pw to one representation
        /// @details there may be multiple output representations, so use vector. Multiple requirements may be from out_qo, out_wannier, out_cohp in INPUT file
        std::vector<RepOut<T, Device>*> repout;
}
//representation.cpp
template<typename T, typename Device>
void Representation<T, Device>::add_repout(const std::string repout_name)
{
    if (repout_name == "pw")
    {
        this->repout.push_back(new RepOut_PW<T, Device>());
    }
    else if (repout_name == "qo")
    {
        this->repout.push_back(new RepOut_QO<T, Device>());
    }
    else if (repout_name == "wannier")
    {
        this->repout.push_back(new RepOut_Wannier<T, Device>());
    }
    else
    {
        ....
    }
    this->repout.back()->set_rep_from(this->representations["input"][0]);
    this->representations["output"].push_back(repout_name);
}

Timeline

• [ ] (RepIn and RepOut) refactor psi_initializer
• [ ] investigate COHP and implement support for it
• [ ] investigate QO and implement support for it (before 15th Nov)
• [ ] investigate current implementation of Wannier function and refactor
• [ ] implement GTO support

Additional Context

No response

Task list for Issue attackers

• [x] Identify the specific code file or section with the code quality issue.
• [x] Investigate the issue and determine the root cause.
• [x] Research best practices and potential solutions for the identified issue.
• [ ] Refactor the code to improve code quality, following the suggested solution.
• [ ] Ensure the refactored code adheres to the project's coding standards.
• [ ] Test the refactored code to ensure it functions as expected.
• [ ] Update any relevant documentation, if necessary.
• [ ] Submit a pull request with the refactored code and a description of the changes made.

[QuantumMisaka]

can COHP part be lcao -> lcao and pw -> lcao? Theoretically, COHP is based by atomic orbital, the work of popular software LOBSTER is based by VASP so the COHP analysis seems to based by PW, but in LOBSTER, pw -> lcao transformation is still needed.

[kirk0830]

can COHP part be lcao -> lcao and pw -> lcao? Theoretically, COHP is based by atomic orbital, the work of popular software LOBSTER is based by VASP so the COHP analysis seems to based by PW, but in LOBSTER, pw -> lcao transformation is still needed.

@QuantumMisaka For rapid implementation and use, I plan to first support QE-like pw wavefunction output, so that one can directly use pw wavefunction to do such analysis. Other representations are, indeed in development plan. Let's see if present solution can really work. Well (on the other hand) I have to say COHP is just Hamiltonian-weighted DOS, the key reason is, present implementation of the way to calculate DOS is not compatible with COHP theoretical framework. Basically analysis of COHP and QO will be supported no later than 15th Nov.

[mohanchen]

We will close this issue since no more detailed plans will be made in the near future. We always welcome new issues to discuss this feature.