Add support for INPUT deepks_v_delta>0 in multi-k points DeePKS calculations (deepmodeling#5700)

Author: ErjieWu

source/Makefile.Objects

@@ -208,6 +208,7 @@ OBJS_DEEPKS=LCAO_deepks.o\
         cal_gvx.o\
         cal_descriptor.o\
         v_delta_precalc.o\
+        v_delta_precalc_k.o\
 
 
 OBJS_ELECSTAT=elecstate.o\

source/module_esolver/esolver_ks_lcao.cpp

@@ -1002,7 +1002,7 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(UnitCell& ucell, const int istep)
 #ifdef __DEEPKS
             if (PARAM.inp.deepks_out_labels && PARAM.inp.deepks_v_delta)
             {
-                DeePKS_domain::save_h_mat(h_mat.p, this->pv.nloc);
+                DeePKS_domain::save_h_mat(h_mat.p, this->pv.nloc, ik);
             }
 #endif
         }

source/module_hamilt_lcao/hamilt_lcaodft/LCAO_allocate.cpp

@@ -27,6 +27,7 @@ void divide_HS_in_frag(const bool isGamma, const UnitCell& ucell, Parallel_Orbit
         GlobalC::ld.init(orb,
                          ucell.nat,
                          ucell.ntype,
+                         nks,
                          pv,
                          na);
 

source/module_hamilt_lcao/module_deepks/CMakeLists.txt

@@ -20,6 +20,7 @@ if(ENABLE_DEEPKS)
       cal_gvx.cpp      
       cal_descriptor.cpp
       v_delta_precalc.cpp
+      v_delta_precalc_k.cpp
   )
 
   add_library(

source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp

@@ -74,6 +74,7 @@ void LCAO_Deepks::init(
     const LCAO_Orbitals& orb,
     const int nat,
     const int ntype,
+    const int nks,
     const Parallel_Orbitals& pv_in,
     std::vector<int> na)
 {
@@ -97,13 +98,13 @@ void LCAO_Deepks::init(
     this->nmaxd = nm;
 
     GlobalV::ofs_running << " lmax of descriptor = " << this->lmaxd << std::endl;
-    GlobalV::ofs_running << " nmax of descriptor= " << nmaxd << std::endl;
+    GlobalV::ofs_running << " nmax of descriptor = " << nmaxd << std::endl;
 
     int pdm_size = 0;
     this->inlmax = tot_inl;
     if(!PARAM.inp.deepks_equiv)
     {
-        GlobalV::ofs_running << " total basis (all atoms) for descriptor= " << std::endl;
+        GlobalV::ofs_running << " total basis (all atoms) for descriptor = " << std::endl;
 
         //init pdm**
         pdm_size = (this->lmaxd * 2 + 1) * (this->lmaxd * 2 + 1);
@@ -150,6 +151,15 @@ void LCAO_Deepks::init(
             int nloc=this->pv->nloc;
             this->h_mat.resize(nloc,0.0);
         }
+        else
+        {
+            int nloc=this->pv->nloc;
+            this->h_mat_k.resize(nks);
+            for (int ik = 0; ik < nks; ik++)
+            {
+                this->h_mat_k[ik].resize(nloc,std::complex<double>(0.0,0.0));
+            }
+        }
     }
 
     return;
@@ -431,27 +441,51 @@ void LCAO_Deepks::del_orbital_pdm_shell(const int nks)
 
 void LCAO_Deepks::init_v_delta_pdm_shell(const int nks,const int nlocal)
 {
-    
-    this->v_delta_pdm_shell = new double**** [nks];
-
     const int mn_size=(2 * this->lmaxd + 1) * (2 * this->lmaxd + 1);
-    for (int iks=0; iks<nks; iks++)
-    {
-        this->v_delta_pdm_shell[iks] = new double*** [nlocal];
+    if (nks==1){
+        this->v_delta_pdm_shell = new double**** [nks];
+        for (int iks=0; iks<nks; iks++)
+        {
+            this->v_delta_pdm_shell[iks] = new double*** [nlocal];
 
-        for (int mu=0; mu<nlocal; mu++)
+            for (int mu=0; mu<nlocal; mu++)
+            {
+                this->v_delta_pdm_shell[iks][mu] = new double** [nlocal];
+
+                for (int nu=0; nu<nlocal; nu++)
+                {
+                    this->v_delta_pdm_shell[iks][mu][nu] = new double* [this->inlmax];
+
+                    for(int inl = 0; inl < this->inlmax; inl++)
+                    {
+                        this->v_delta_pdm_shell[iks][mu][nu][inl] = new double [mn_size];
+                        ModuleBase::GlobalFunc::ZEROS(v_delta_pdm_shell[iks][mu][nu][inl], mn_size);
+                    }                
+                }
+            }
+        }
+    }
+    else
+    {
+        this->v_delta_pdm_shell_complex = new std::complex<double>**** [nks];
+        for (int iks=0; iks<nks; iks++)
         {
-            this->v_delta_pdm_shell[iks][mu] = new double** [nlocal];
+            this->v_delta_pdm_shell_complex[iks] = new std::complex<double>*** [nlocal];
 
-            for (int nu=0; nu<nlocal; nu++)
+            for (int mu=0; mu<nlocal; mu++)
             {
-                this->v_delta_pdm_shell[iks][mu][nu] = new double* [this->inlmax];
+                this->v_delta_pdm_shell_complex[iks][mu] = new std::complex<double>** [nlocal];
 
-                for(int inl = 0; inl < this->inlmax; inl++)
+                for (int nu=0; nu<nlocal; nu++)
                 {
-                    this->v_delta_pdm_shell[iks][mu][nu][inl] = new double [mn_size];
-                    ModuleBase::GlobalFunc::ZEROS(v_delta_pdm_shell[iks][mu][nu][inl], mn_size);
-                }                
+                    this->v_delta_pdm_shell_complex[iks][mu][nu] = new std::complex<double>* [this->inlmax];
+
+                    for(int inl = 0; inl < this->inlmax; inl++)
+                    {
+                        this->v_delta_pdm_shell_complex[iks][mu][nu][inl] = new std::complex<double> [mn_size];
+                        ModuleBase::GlobalFunc::ZEROS(v_delta_pdm_shell_complex[iks][mu][nu][inl], mn_size);
+                    }                
+                }
             }
         }
     }
@@ -461,23 +495,46 @@ void LCAO_Deepks::init_v_delta_pdm_shell(const int nks,const int nlocal)
 
 void LCAO_Deepks::del_v_delta_pdm_shell(const int nks,const int nlocal)
 {
-    for (int iks=0; iks<nks; iks++)
+    if (nks==1)
     {
-        for (int mu=0; mu<nlocal; mu++)
+        for (int iks=0; iks<nks; iks++)
         {
-            for (int nu=0; nu<nlocal; nu++)
+            for (int mu=0; mu<nlocal; mu++)
             {
-                for (int inl = 0;inl < this->inlmax; inl++)
+                for (int nu=0; nu<nlocal; nu++)
                 {
-                    delete[] this->v_delta_pdm_shell[iks][mu][nu][inl];
+                    for (int inl = 0;inl < this->inlmax; inl++)
+                    {
+                        delete[] this->v_delta_pdm_shell[iks][mu][nu][inl];
+                    }
+                    delete[] this->v_delta_pdm_shell[iks][mu][nu];                
+                }
+                delete[] this->v_delta_pdm_shell[iks][mu]; 
+            }
+            delete[] this->v_delta_pdm_shell[iks]; 
+        }
+        delete[] this->v_delta_pdm_shell;    
+    }
+    else
+    {
+        for (int iks=0; iks<nks; iks++)
+        {
+            for (int mu=0; mu<nlocal; mu++)
+            {
+                for (int nu=0; nu<nlocal; nu++)
+                {
+                    for (int inl = 0;inl < this->inlmax; inl++)
+                    {
+                        delete[] this->v_delta_pdm_shell_complex[iks][mu][nu][inl];
+                    }
+                    delete[] this->v_delta_pdm_shell_complex[iks][mu][nu];                
                 }
-                delete[] this->v_delta_pdm_shell[iks][mu][nu];                
+                delete[] this->v_delta_pdm_shell_complex[iks][mu]; 
             }
-            delete[] this->v_delta_pdm_shell[iks][mu]; 
+            delete[] this->v_delta_pdm_shell_complex[iks]; 
         }
-        delete[] this->v_delta_pdm_shell[iks]; 
+        delete[] this->v_delta_pdm_shell_complex;    
     }
-    delete[] this->v_delta_pdm_shell;    
 
     return;
 }

source/module_hamilt_lcao/module_deepks/LCAO_deepks.h

@@ -54,8 +54,11 @@ class LCAO_Deepks
     ///\rho_{HL} = c_{L, \mu}c_{L,\nu} - c_{H, \mu}c_{H,\nu} \f$ (for gamma_only)
     ModuleBase::matrix o_delta;
 
-    ///(Unit: Ry) Hamiltonian matrix
+    ///(Unit: Ry) Hamiltonian matrix in k space
+    /// for gamma only
     std::vector<double> h_mat;    
+    /// for multi-k
+    std::vector<std::vector<std::complex<double>>> h_mat_k;
 
     /// Correction term to the Hamiltonian matrix: \f$\langle\psi|V_\delta|\psi\rangle\f$ (for gamma only)
     std::vector<double> H_V_delta;
@@ -159,13 +162,14 @@ class LCAO_Deepks
     // dD/dX, tensor form of gdmx
     std::vector<torch::Tensor> gdmr_vector;
 
-    // orbital_pdm_shell:[1,Inl,nm*nm]; \langle \phi_\mu|\alpha\rangle\langle\alpha|\phi_\nu\rnalge
+    // orbital_pdm_shell:[1,Inl,nm*nm]; \langle \phi_\mu|\alpha\rangle\langle\alpha|\phi_\nu\ranlge
     double**** orbital_pdm_shell;
     // orbital_precalc:[1,NAt,NDscrpt]; gvdm*orbital_pdm_shell
     torch::Tensor orbital_precalc_tensor;
 
     // v_delta_pdm_shell[nks,nlocal,nlocal,Inl,nm*nm] = overlap * overlap
     double***** v_delta_pdm_shell;
+    std::complex<double>***** v_delta_pdm_shell_complex; // for multi-k
     // v_delta_precalc[nks,nlocal,nlocal,NAt,NDscrpt] = gvdm * v_delta_pdm_shell;
     torch::Tensor v_delta_precalc_tensor;
     //for v_delta==2 , new v_delta_precalc storage method
@@ -220,6 +224,7 @@ class LCAO_Deepks
     void init(const LCAO_Orbitals& orb,
               const int nat,
               const int ntype,
+              const int nks,
               const Parallel_Orbitals& pv_in,
               std::vector<int> na);
 
@@ -437,12 +442,15 @@ class LCAO_Deepks
     // 11. cal_orbital_precalc_k : orbital_precalc is usted for training with orbital label,
     //                          for multi-k case, which equals gvdm * orbital_pdm_shell,
     //                          orbital_pdm_shell[1,Inl,nm*nm] = dm_hl_k * overlap * overlap
-    //12. cal_v_delta_precalc : v_delta_precalc is used for training with v_delta label,
+    // 12. cal_v_delta_precalc : v_delta_precalc is used for training with v_delta label,
     //                         which equals gvdm * v_delta_pdm_shell,
     //                         v_delta_pdm_shell = overlap * overlap
-    //13. check_v_delta_precalc : check v_delta_precalc
-    //14. prepare_psialpha : prepare psialpha for outputting npy file
-    //15. prepare_gevdm : prepare gevdm for outputting npy file
+    // 13. cal_v_delta_precalc_k : v_delta_precalc is used for training with v_delta label,
+    //                         for multi-k case, which equals ???
+    //                         ???
+    // 14. check_v_delta_precalc : check v_delta_precalc
+    // 15. prepare_psialpha : prepare psialpha for outputting npy file
+    // 16. prepare_gevdm : prepare gevdm for outputting npy file
 
   public:
     /// Calculates descriptors
@@ -500,6 +508,14 @@ class LCAO_Deepks
         const LCAO_Orbitals &orb,
         Grid_Driver &GridD);
 
+    void cal_v_delta_precalc_k(const int nlocal,
+        const int nat,
+        const int nks,
+        const std::vector<ModuleBase::Vector3<double>> &kvec_d,
+        const UnitCell &ucell,
+        const LCAO_Orbitals &orb,
+        Grid_Driver &GridD);
+
     void check_v_delta_precalc(const int nat, const int nks,const int nlocal);
 
     // prepare psialpha for outputting npy file
@@ -508,6 +524,13 @@ class LCAO_Deepks
         const UnitCell &ucell,
         const LCAO_Orbitals &orb,
         Grid_Driver &GridD);
+    void prepare_psialpha_k(const int nlocal,
+        const int nat,
+        const int nks,
+        const std::vector<ModuleBase::Vector3<double>> &kvec_d,
+        const UnitCell &ucell,
+        const LCAO_Orbitals &orb,
+        Grid_Driver &GridD);
     void check_vdp_psialpha(const int nat, const int nks, const int nlocal);
 
     // prepare gevdm for outputting npy file

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -291,7 +291,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot,
 
             if (PARAM.inp.deepks_scf)
             {
-                int nocc = PARAM.inp.nelec / 2;
+                int nocc = PARAM.inp.nelec / 2; // redundant!
                 ModuleBase::matrix wg_hl;
                 wg_hl.create(nks, PARAM.inp.nbands);
                 std::vector<std::vector<ModuleBase::ComplexMatrix>> dm_bandgap_k;
@@ -333,7 +333,94 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot,
         } // end bandgap label
         if(deepks_v_delta)
         {
-            ModuleBase::WARNING_QUIT("ESolver_KS_LCAO", "V_delta label has not been developed for multi-k now!");
+            std::vector<ModuleBase::ComplexMatrix> h_tot(nks);
+            for (int ik = 0; ik < nks; ik++)
+            {
+                h_tot[ik].create(nlocal, nlocal);
+            }
+
+            DeePKS_domain::collect_h_mat(*ParaV, ld->h_mat_k,h_tot,nlocal,nks);
+
+            const std::string file_htot = PARAM.globalv.global_out_dir + "deepks_htot.npy";
+            LCAO_deepks_io::save_npy_h(h_tot, file_htot, nlocal, nks, my_rank);
+
+            if(PARAM.inp.deepks_scf)
+            {
+                std::vector<ModuleBase::ComplexMatrix> v_delta(nks);
+                std::vector<ModuleBase::ComplexMatrix> hbase(nks);
+                for (int ik = 0; ik < nks; ik++)
+                {
+                    v_delta[ik].create(nlocal, nlocal);
+                    hbase[ik].create(nlocal, nlocal);
+                }
+                DeePKS_domain::collect_h_mat(*ParaV, ld->H_V_delta_k,v_delta,nlocal,nks); 
+
+                const std::string file_hbase = PARAM.globalv.global_out_dir + "deepks_hbase.npy";
+                for (int ik = 0; ik < nks; ik++)
+                {
+                    hbase[ik] = h_tot[ik] - v_delta[ik];
+                }
+                LCAO_deepks_io::save_npy_h(hbase, file_hbase, nlocal, nks, my_rank);
+
+                const std::string file_vdelta = PARAM.globalv.global_out_dir + "deepks_vdelta.npy";
+                LCAO_deepks_io::save_npy_h(v_delta, file_vdelta, nlocal, nks, my_rank);
+
+                if(deepks_v_delta==1)//v_delta_precalc storage method 1
+                {
+                    ld->cal_v_delta_precalc_k(nlocal,
+                            nat,
+                            nks,
+                            kvec_d,
+                            ucell,
+                            orb,
+                            GridD);
+                
+                    LCAO_deepks_io::save_npy_v_delta_precalc(
+                      nat, 
+                      nks, 
+                      nlocal,
+                      ld->des_per_atom,
+                      ld->v_delta_precalc_tensor,
+                      PARAM.globalv.global_out_dir,
+                      my_rank);
+
+                }
+                else if(deepks_v_delta==2)//v_delta_precalc storage method 2
+                {
+                    ld->prepare_psialpha_k(nlocal,
+                                nat,
+                                nks,
+                                kvec_d,
+                                ucell,
+                                orb,
+                                GridD);
+                    
+                    LCAO_deepks_io::save_npy_psialpha(nat, 
+                                nks, 
+                                nlocal,
+                                ld->inlmax,
+                                ld->lmaxd,
+                                ld->psialpha_tensor,
+                                PARAM.globalv.global_out_dir,
+                                my_rank);
+
+                    ld->prepare_gevdm(
+                                nat,
+                                orb);
+
+                    LCAO_deepks_io::save_npy_gevdm(nat,
+                      ld->inlmax,
+                      ld->lmaxd,
+                      ld->gevdm_tensor,
+                      PARAM.globalv.global_out_dir,
+                      my_rank);
+                }
+            }
+            else //deepks_scf == 0
+            {
+                const std::string file_hbase = PARAM.globalv.global_out_dir + "deepks_hbase.npy";
+                LCAO_deepks_io::save_npy_h(h_tot, file_hbase, nlocal, nks, my_rank);
+            }
         }
     }  // end deepks_out_labels