hamilt_lcao.cpp

#include "source_lcao/hamilt_lcao.h"

#include "source_base/global_variable.h"
#include "source_base/memory_recorder.h"
#include "source_base/timer.h"
#include "source_lcao/module_dftu/dftu_lcao.h"
#include "source_lcao/setup_exx.h"
#include "source_lcao/setup_deepks.h"
#include "source_estate/module_dm/density_matrix.h"
#include "source_estate/module_pot/potential_new.h"
#include "source_io/module_parameter/parameter.h"

#include <vector>

#ifdef __MLALGO
#include "source_lcao/module_deepks/lcao_deepks.h"
#include "module_operator_lcao/deepks_lcao.h"
#endif

#ifdef __EXX
#include "source_lcao/module_ri/exx_lri_interface.h"
#include "module_operator_lcao/op_exx_lcao.h"
#endif

#ifdef __ELPA
#include "source_hsolver/diago_elpa.h"
#endif

#include "source_estate/module_pot/h_tddft_pw.h"
#include "source_hamilt/module_xc/xc_functional.h"
#include "source_lcao/module_deltaspin/spin_constrain.h"
#include "source_hamilt/module_hcontainer/hcontainer_funcs.h"
#include "source_hsolver/hsolver_lcao.h"
#include "module_dftu/dftu_lcao_op.h"
#include "module_operator_lcao/dspin_lcao.h"
#include "module_operator_lcao/ekinetic.h"
#include "module_operator_lcao/meta_lcao.h"
#include "module_operator_lcao/nonlocal.h"
#include "module_dftu/dftu_lcao_op_legacy.h"
#include "module_operator_lcao/op_exx_lcao.h"
#include "module_operator_lcao/overlap.h"
#include "module_operator_lcao/td_ekinetic_lcao.h"
#include "module_operator_lcao/td_nonlocal_lcao.h"
#include "module_operator_lcao/td_pot_hybrid.h"
#include "module_operator_lcao/veff_lcao.h"


namespace hamilt
{

template <typename TK, typename TR>
HamiltLCAO<TK, TR>::HamiltLCAO(const UnitCell& ucell,
                               const Grid_Driver& grid_d,
                               const Parallel_Orbitals* paraV,
                               const K_Vectors& kv_in,
                               const TwoCenterIntegrator& intor_overlap_orb,
                               const std::vector<double>& orb_cutoff)
{
    this->classname = "HamiltLCAO";

    this->kv = &kv_in;

    // initialize the overlap matrix
    this->sR = new HContainer<TR>(paraV);

    this->getOperator() = new Overlap<OperatorLCAO<TK, TR>>(this->hsk,
                                                               this->kv->kvec_d,
                                                               this->hR,
                                                               this->sR,
                                                               &ucell,
                                                               orb_cutoff,
                                                               &grid_d,
                                                               &intor_overlap_orb);
}

template <typename TK, typename TR>
HamiltLCAO<TK, TR>::HamiltLCAO(const UnitCell& ucell,
                               const Grid_Driver& grid_d,
                               const Parallel_Orbitals* paraV,
                               elecstate::Potential* pot_in,
                               const K_Vectors& kv_in,
                               const TwoCenterBundle& two_center_bundle,
                               const LCAO_Orbitals& orb,
							   elecstate::DensityMatrix<TK, double>* DM_in,
							   Plus_U* p_dftu, // mohan add 2025-11-05
							   Setup_DeePKS<TK> &deepks,
							   const int istep,
							   Exx_NAO<TK> &exx_nao,
							   const Exx_Info& exx_info)
{
    this->classname = "HamiltLCAO";

    this->kv = &kv_in;

    // Real space Hamiltonian is inited with template TR
    this->hR = new HContainer<TR>(paraV);
    this->sR = new HContainer<TR>(paraV);
    this->hsk = new HS_Matrix_K<TK>(paraV);

    // Effective potential term (\sum_r <psi(r)|Veff(r)|psi(r)>) is registered without template
    std::vector<std::string> pot_register_in;
    if (PARAM.inp.vl_in_h)
    {
        if (PARAM.inp.vion_in_h)
        {
            pot_register_in.push_back("local");
        }
        if (PARAM.inp.vh_in_h)
        {
            pot_register_in.push_back("hartree");
        }
        pot_register_in.push_back("xc");
        if (PARAM.inp.imp_sol)
        {
            pot_register_in.push_back("surchem");
        }
        if (PARAM.inp.efield_flag)
        {
            pot_register_in.push_back("efield");
        }
        if (PARAM.inp.gate_flag)
        {
            pot_register_in.push_back("gatefield");
        }
        if (PARAM.inp.esolver_type == "tddft")
        {
            pot_register_in.push_back("tddft");
        }
        if (PARAM.inp.ml_exx) // sunliang
        {
            pot_register_in.push_back("ml_exx");
        }
    }

    // Gamma_only case to initialize HamiltLCAO
    //
    // code block to construct Operator Chains
    if (std::is_same<TK, double>::value)
    {
        // fix HR to gamma case, where SR will be fixed in Overlap Operator
        this->hR->fix_gamma();
        // initial operator for Gamma_only case
        // overlap term (<psi|psi>) is indispensable
        // in Gamma_only case, target SK is this->hsk->get_sk(), the target SR is this->sR
        this->getOperator() = new Overlap<OperatorLCAO<TK, TR>>(this->hsk,
                                                                   this->kv->kvec_d,
                                                                   this->hR,
                                                                   this->sR,
                                                                   &ucell,
                                                                   orb.cutoffs(),
                                                                   &grid_d,
                                                                   two_center_bundle.overlap_orb.get());

        // kinetic term (<psi|T|psi>)
        if (PARAM.inp.t_in_h)
        {
            Operator<TK>* ekinetic = new EKinetic<OperatorLCAO<TK, TR>>(this->hsk,
                                                                           this->kv->kvec_d,
                                                                           this->hR,
                                                                           &ucell,
                                                                           orb.cutoffs(),
                                                                           &grid_d,
                                                                           two_center_bundle.kinetic_orb.get());
            this->getOperator()->add(ekinetic);
        }

        // nonlocal term (<psi|beta>D<beta|psi>)
        // in general case, target HR is this->hR, while target HK is this->hsk->get_hk()
        if (PARAM.inp.vnl_in_h)
        {
            Operator<TK>* nonlocal = new Nonlocal<OperatorLCAO<TK, TR>>(this->hsk,
                                                                           this->kv->kvec_d,
                                                                           this->hR,
                                                                           &ucell,
                                                                           orb.cutoffs(),
                                                                           &grid_d,
                                                                           two_center_bundle.overlap_orb_beta.get());
            this->getOperator()->add(nonlocal);
        }

        // Effective potential term (\sum_r <psi(r)|Veff(r)|psi(r)>)
        // in general case, target HR is Gint::hRGint, while target HK is this->hsk->get_hk()
        if (PARAM.inp.vl_in_h)
        {
            // only Potential is not empty, Veff and Meta are available
            if (pot_register_in.size() > 0)
            {
                // register Potential by gathered operator
                pot_in->pot_register(pot_register_in);
                // effective potential term
                Operator<TK>* veff = new Veff<OperatorLCAO<TK, TR>>(this->hsk,
                                                                    this->kv->kvec_d,
                                                                    pot_in,
                                                                    this->hR, // no explicit call yet
                                                                    &ucell,
                                                                    orb.cutoffs(),
                                                                    &grid_d,
                                                                    PARAM.inp.nspin);
                this->getOperator()->add(veff);
            }
        }

#ifdef __MLALGO
        if (PARAM.inp.deepks_scf)
        {
            Operator<TK>* deepks_op = new DeePKS<OperatorLCAO<TK, TR>>(this->hsk,
                                                                    this->kv->kvec_d,
                                                                    this->hR, // no explicit call yet
                                                                    &ucell,
                                                                    &grid_d,
                                                                    two_center_bundle.overlap_orb_alpha.get(),
                                                                    &orb,
                                                                    this->kv->get_nks(),
                                                                    DM_in,
                                                                    &deepks.ld);
            this->getOperator()->add(deepks_op);
            this->V_delta_R = dynamic_cast<DeePKS<OperatorLCAO<TK, TR>>*>(deepks_op)->get_V_delta_R();
        }
#endif

        // end node should be OperatorDFTU
        if (PARAM.inp.dft_plus_u)
        {
            Operator<TK>* plus_u = nullptr;
            if (PARAM.inp.dft_plus_u == 2)
            {
                plus_u = new OperatorDFTU<OperatorLCAO<TK, TR>>(this->hsk,
                                                              this->kv->kvec_d,
															  this->hR,
															  p_dftu,
															  this->kv->isk,
															  PARAM.globalv.npol);
            }
            else
            {
                plus_u = new DFTU<OperatorLCAO<TK, TR>>(this->hsk,
                                                      this->kv->kvec_d,
                                                      this->hR,
                                                      ucell,
                                                      &grid_d,
                                                      two_center_bundle.overlap_orb_onsite.get(),
                                                      orb.cutoffs(),
                                                      p_dftu);
            }
            this->getOperator()->add(plus_u);
        }
    }
    // multi-k-points case to initialize HamiltLCAO, ops will be used
    else if (std::is_same<TK, std::complex<double>>::value)
    {
        // Effective potential term (\sum_r <psi(r)|Veff(r)|psi(r)>)
        // Meta potential term (\sum_r <psi(r)|tau(r)|psi(r)>)
        // in general case, target HR is Gint::pvpR_reduced, while target HK is this->hsk->get_hk()
        if (PARAM.inp.vl_in_h)
        {
            // only Potential is not empty, Veff and Meta are available
            if (pot_register_in.size() > 0)
            {
                // register Potential by gathered operator
                pot_in->pot_register(pot_register_in);
                // Veff term
                this->getOperator() = new Veff<OperatorLCAO<TK, TR>>(this->hsk,
                                                                     this->kv->kvec_d,
                                                                     pot_in,
                                                                     this->hR,
                                                                     &ucell,
                                                                     orb.cutoffs(),
                                                                     &grid_d,
                                                                     PARAM.inp.nspin);
            }
        }

        // initial operator for multi-k case
        // overlap term is indispensable
        Operator<TK>* overlap = new Overlap<OperatorLCAO<TK, TR>>(this->hsk,
                                                                     this->kv->kvec_d,
                                                                     this->hR,
                                                                     this->sR,
                                                                     &ucell,
                                                                     orb.cutoffs(),
                                                                     &grid_d,
                                                                     two_center_bundle.overlap_orb.get());
        if (this->getOperator() == nullptr)
        {
            this->getOperator() = overlap;
        }
        else
        {
            this->getOperator()->add(overlap);
        }

        // kinetic term (<psi|T|psi>),
        // in general case, target HR is this->hR, while target HK is this->hsk->get_hk()
        if (PARAM.inp.t_in_h)
        {
            Operator<TK>* ekinetic = new EKinetic<OperatorLCAO<TK, TR>>(this->hsk,
                                                                           this->kv->kvec_d,
                                                                           this->hR,
                                                                           &ucell,
                                                                           orb.cutoffs(),
                                                                           &grid_d,
                                                                           two_center_bundle.kinetic_orb.get());
            this->getOperator()->add(ekinetic);
        }

        // nonlocal term (<psi|beta>D<beta|psi>)
        // in general case, target HR is this->hR, while target HK is this->hsk->get_hk()
        if (PARAM.inp.vnl_in_h)
        {
            Operator<TK>* nonlocal = new Nonlocal<OperatorLCAO<TK, TR>>(this->hsk,
                                                                           this->kv->kvec_d,
                                                                           this->hR,
                                                                           &ucell,
                                                                           orb.cutoffs(),
                                                                           &grid_d,
                                                                           two_center_bundle.overlap_orb_beta.get());
            // TDDFT velocity gauge will calculate full non-local potential including the original one and the
            // correction on its own. So the original non-local potential term should be skipped
            if (PARAM.inp.esolver_type != "tddft" || elecstate::H_TDDFT_pw::stype != 1)
            {
                this->getOperator()->add(nonlocal);
            }
            else
            {
                delete nonlocal;
            }
        }

#ifdef __MLALGO
        if (PARAM.inp.deepks_scf)
        {
            Operator<TK>* deepks_op = new DeePKS<OperatorLCAO<TK, TR>>(this->hsk,
                                                                    this->kv->kvec_d,
                                                                    hR,
                                                                    &ucell,
                                                                    &grid_d,
                                                                    two_center_bundle.overlap_orb_alpha.get(),
                                                                    &orb,
                                                                    this->kv->get_nks(),
                                                                    DM_in,
                                                                    &deepks.ld);
            this->getOperator()->add(deepks_op);
            this->V_delta_R = dynamic_cast<DeePKS<OperatorLCAO<TK, TR>>*>(deepks_op)->get_V_delta_R();
        }
#endif
        // TDDFT_velocity_gauge
        if (PARAM.inp.esolver_type == "tddft" && PARAM.inp.td_stype == 1)
        {
            Operator<TK>* td_ekinetic = new TDEkinetic<OperatorLCAO<TK, TR>>(this->hsk,
                                                                             this->hR,
                                                                             this->kv,
                                                                             &ucell,
                                                                             orb.cutoffs(),
                                                                             &grid_d,
                                                                             two_center_bundle.overlap_orb.get());
            this->getOperator()->add(td_ekinetic);

            Operator<TK>* td_nonlocal = new TDNonlocal<OperatorLCAO<TK, TR>>(this->hsk,
                                                                             this->kv->kvec_d,
                                                                             this->hR,
                                                                             &ucell,
                                                                             orb,
                                                                             &grid_d);
            this->getOperator()->add(td_nonlocal);
        }
        if (PARAM.inp.esolver_type == "tddft" && PARAM.inp.td_stype == 2)
        {
            Operator<TK>* td_pot_hybrid = new TD_pot_hybrid<OperatorLCAO<TK, TR>>(this->hsk,
                                                                           this->kv,
                                                                           this->hR,
                                                                           this->sR,
                                                                           orb,
                                                                           &ucell,
                                                                           orb.cutoffs(),
                                                                           &grid_d,
                                                                           two_center_bundle.kinetic_orb.get());
            this->getOperator()->add(td_pot_hybrid);
        }
        if (PARAM.inp.dft_plus_u)
        {
            Operator<TK>* plus_u = nullptr;
            if (PARAM.inp.dft_plus_u == 2)
            {
                plus_u = new OperatorDFTU<OperatorLCAO<TK, TR>>(this->hsk,
                                                              this->kv->kvec_d,
															  this->hR,
															  p_dftu,
                                                              this->kv->isk,
                                                              PARAM.globalv.npol);
            }
            else
            {
                plus_u = new DFTU<OperatorLCAO<TK, TR>>(this->hsk,
                                                      this->kv->kvec_d,
                                                      this->hR,
                                                      ucell,
                                                      &grid_d,
                                                      two_center_bundle.overlap_orb_onsite.get(),
                                                      orb.cutoffs(),
                                                      p_dftu);
            }
            this->getOperator()->add(plus_u);
        }
        if (PARAM.inp.sc_mag_switch)
        {
            Operator<TK>* sc_lambda = new DeltaSpin<OperatorLCAO<TK, TR>>(this->hsk,
                                                                          this->kv->kvec_d,
                                                                          this->hR,
                                                                          ucell,
                                                                          &grid_d,
                                                                          two_center_bundle.overlap_orb_onsite.get(),
                                                                          orb.cutoffs());
            this->getOperator()->add(sc_lambda);
            spinconstrain::SpinConstrain<TK>& sc = spinconstrain::SpinConstrain<TK>::getScInstance();
            sc.set_operator(sc_lambda);
        }
    }

#ifdef __EXX
    if (exx_info.info_global.cal_exx)
    {
        // Peize Lin add 2016-12-03
        // set xc type before the first cal of xc in pelec->init_scf
        // and calculate Cs, Vs
        Operator<TK>* exx;
        if (PARAM.inp.esolver_type == "tddft")
        {
            exx = new OperatorEXX<OperatorLCAO<TK, TR>>(this->hsk,
                                                        this->hR,
                                                        ucell,
                                                        *this->kv,
                                                        exx_nao.exd.get(),
                                                        exx_nao.exc.get(),
                                                        exx_info,
                                                        Add_Hexx_Type::k,
                                                        istep,
                                                        !GlobalC::restart.info_load.restart_exx
                                                            && GlobalC::restart.info_load.load_H);
        }
        else
        {
            exx = new OperatorEXX<OperatorLCAO<TK, TR>>(this->hsk,
                                                        this->hR,
                                                        ucell,
                                                        *kv,
                                                        exx_nao.exd.get(),
                                                        exx_nao.exc.get(),
                                                        exx_info,
                                                        Add_Hexx_Type::R,
                                                        istep,
                                                        !GlobalC::restart.info_load.restart_exx
                                                            && GlobalC::restart.info_load.load_H);
        }
        this->getOperator()->add(exx);
    }
#endif

    // if NSPIN==2, HR should be separated into two parts, save HR into this->hRS2
    int memory_fold = 1;
    if (PARAM.inp.nspin == 2)
    {
        this->hRS2.resize(this->hR->get_nnr() * 2);
        this->hR->allocate(this->hRS2.data(), 0);
        memory_fold = 2;
    }

    ModuleBase::Memory::record("HamiltLCAO::hR", this->hR->get_memory_size() * memory_fold);
    ModuleBase::Memory::record("HamiltLCAO::sR", this->sR->get_memory_size());

    return;
}

template <typename TK, typename TR>
std::vector<HContainer<TR>*> HamiltLCAO<TK, TR>::getHR_vector()
{
    if (PARAM.inp.nspin == 2)
    {
        const int nnr = this->hRS2.size() / 2;
        this->hr_spin_up_.reset(new HContainer<TR>(*this->hR, this->hRS2.data()));
        this->hr_spin_dn_.reset(new HContainer<TR>(*this->hR, this->hRS2.data() + nnr));
        return {this->hr_spin_up_.get(), this->hr_spin_dn_.get()};
    }
    else
    {
        return {this->hR};
    }
}

// case for multi-k-points
template <typename TK, typename TR>
void HamiltLCAO<TK, TR>::matrix(MatrixBlock<TK>& hk_in, MatrixBlock<TK>& sk_in)
{
    auto op = dynamic_cast<OperatorLCAO<TK, TR>*>(this->getOperator());
    assert(op != nullptr);
    op->matrixHk(hk_in, sk_in);
}

template <typename TK, typename TR>
void HamiltLCAO<TK, TR>::updateHk(const int ik)
{
    ModuleBase::TITLE("HamiltLCAO", "updateHk");
    ModuleBase::timer::start("HamiltLCAO", "updateHk");

    // update global spin index
    if (PARAM.inp.nspin == 2)
    {
        // if Veff is added and current_spin is changed, refresh HR
        if (PARAM.inp.vl_in_h && this->kv->isk[ik] != this->current_spin)
        {
            // change data pointer of HR
            this->hR->allocate(this->hRS2.data() + this->hRS2.size() / 2 * this->kv->isk[ik], 0);
            if (this->refresh_times > 0)
            {
                this->refresh_times--;
                dynamic_cast<hamilt::OperatorLCAO<TK, TR>*>(this->ops)->set_hr_done(false);
            }
        }
        this->current_spin = this->kv->isk[ik];
        dynamic_cast<hamilt::OperatorLCAO<TK, TR>*>(this->ops)->set_current_spin(this->kv->isk[ik]);
    }
    this->getOperator()->init(ik);
    ModuleBase::timer::end("HamiltLCAO", "updateHk");
}

template <typename TK, typename TR>
void HamiltLCAO<TK, TR>::refresh(bool yes)
{
    ModuleBase::TITLE("HamiltLCAO", "refresh");
    if(yes)
    {
        dynamic_cast<hamilt::OperatorLCAO<TK, TR>*>(this->ops)->set_hr_done(false);
        if (PARAM.inp.nspin == 2)
        {
            this->refresh_times = 1;
            this->current_spin = 0;
            if (this->hR->get_nnr() != this->hRS2.size() / 2)
            {
                // operator has changed, resize hRS2
                this->hRS2.resize(this->hR->get_nnr() * 2);
            }
            this->hR->allocate(this->hRS2.data(), 0);
        }
    }
    else {
        dynamic_cast<hamilt::OperatorLCAO<TK, TR>*>(this->ops)->set_hr_done(true);
        this->refresh_times = 0;
        if (PARAM.inp.nspin == 2)
        {
            // HR has been loaded from file into both halves of hRS2.
            // Reset to spin-up; updateHk will switch pointers as needed.
            this->current_spin = 0;
            this->hR->allocate(this->hRS2.data(), 0);
        }
    }
}

// get Operator base class pointer
template <typename TK, typename TR>
Operator<TK>*& HamiltLCAO<TK, TR>::getOperator()
{
    return this->ops;
}

template <typename TK, typename TR>
void HamiltLCAO<TK, TR>::updateSk(
		const int ik, 
		const int hk_type)
{
    ModuleBase::TITLE("HamiltLCAO", "updateSk");
    ModuleBase::timer::start("HamiltLCAO", "updateSk");

    ModuleBase::GlobalFunc::ZEROS(this->getSk(), this->get_size_hsk());

    if (hk_type == 1) // collumn-major matrix for SK
    {
        const int nrow = this->hsk->get_pv()->get_row_size();
		hamilt::folding_HR(*this->sR, this->getSk(), this->kv->kvec_d[ik], nrow, 1);
	}
	else if (hk_type == 0) // row-major matrix for SK
	{
        const int ncol = this->hsk->get_pv()->get_col_size();
        hamilt::folding_HR(*this->sR, this->getSk(), this->kv->kvec_d[ik], ncol, 0);
    }
	else
	{
        ModuleBase::WARNING_QUIT("updateSk","the value of hk_type is incorrect.");
	}

    ModuleBase::timer::end("HamiltLCAO", "updateSk");
}

// case for nspin<4, gamma-only k-point
template class HamiltLCAO<double, double>;
// case for nspin<4, multi-k-points
template class HamiltLCAO<std::complex<double>, double>;
// case for nspin == 4, non-collinear spin case
template class HamiltLCAO<std::complex<double>, std::complex<double>>;
} // namespace hamilt