Merge branch 'develop' into init

Author: A-006

docs/advanced/input_files/input-main.md

@@ -2060,7 +2060,7 @@ Warning: this function is not robust enough for the current version. Please try
 - **Type**: int
 - **Availability**: numerical atomic orbital basis
 - **Description**: Include V_delta label for DeePKS training. When `deepks_out_labels` is true and `deepks_v_delta` > 0, ABACUS will output h_base.npy, v_delta.npy and h_tot.npy(h_tot=h_base+v_delta). 
-  Meanwhile, when `deepks_v_delta` equals 1, ABACUS will also output v_delta_precalc.npy, which is used to calculate V_delta during DeePKS training. However, when the number of atoms grows, the size of v_delta_precalc.npy will be very large. In this case, it's recommended to set `deepks_v_delta` as 2, and ABACUS will output psialpha.npy and grad_evdm.npy but not v_delta_precalc.npy. These two files are small and can be used to calculate v_delta_precalc in the procedure of training DeePKS.
+  Meanwhile, when `deepks_v_delta` equals 1, ABACUS will also output v_delta_precalc.npy, which is used to calculate V_delta during DeePKS training. However, when the number of atoms grows, the size of v_delta_precalc.npy will be very large. In this case, it's recommended to set `deepks_v_delta` as 2, and ABACUS will output phialpha.npy and grad_evdm.npy but not v_delta_precalc.npy. These two files are small and can be used to calculate v_delta_precalc in the procedure of training DeePKS.
 - **Default**: 0
 
 ### deepks_out_unittest

source/Makefile.Objects

@@ -189,13 +189,12 @@ OBJS_CELL=atom_pseudo.o\
     check_atomic_stru.o\
 
 OBJS_DEEPKS=LCAO_deepks.o\
-        deepks_fgamma.o\
-        deepks_fk.o\
+        deepks_force.o\
         LCAO_deepks_odelta.o\
         LCAO_deepks_io.o\
         LCAO_deepks_mpi.o\
         LCAO_deepks_pdm.o\
-        LCAO_deepks_psialpha.o\
+        LCAO_deepks_phialpha.o\
         LCAO_deepks_torch.o\
         LCAO_deepks_vdelta.o\
         deepks_hmat.o\

source/module_cell/read_atoms.cpp

@@ -676,45 +676,44 @@ bool UnitCell::read_atom_positions(std::ifstream &ifpos, std::ofstream &ofs_runn
                     std::string mags;
                     //cout<<"mag"<<atoms[it].mag[ia]<<"angle1"<<atoms[it].angle1[ia]<<"angle2"<<atoms[it].angle2[ia]<<'\n';
 
-                    if(PARAM.inp.nspin==4)
-                    {
-                        if(PARAM.inp.noncolin)
+                    // ----------------------------------------------------------------------------
+                    // recalcualte mag and m_loc_ from read in angle1, angle2 and mag or mx, my, mz
+                    if(input_angle_mag)
+                    {// angle1 or angle2 are given, calculate mx, my, mz from angle1 and angle2 and mag
+                        atoms[it].m_loc_[ia].z = atoms[it].mag[ia] *
+                            cos(atoms[it].angle1[ia]);
+                        if(std::abs(sin(atoms[it].angle1[ia])) > 1e-10 )
                         {
-                            if(input_angle_mag)
-                            {
-                                atoms[it].m_loc_[ia].z = atoms[it].mag[ia] *
-                                    cos(atoms[it].angle1[ia]);
-                                if(std::abs(sin(atoms[it].angle1[ia])) > 1e-10 )
-                                {
-                                    atoms[it].m_loc_[ia].x = atoms[it].mag[ia] *
-                                        sin(atoms[it].angle1[ia]) * cos(atoms[it].angle2[ia]);
-                                    atoms[it].m_loc_[ia].y = atoms[it].mag[ia] *
-                                        sin(atoms[it].angle1[ia]) * sin(atoms[it].angle2[ia]);
-                                }
-                            }
-                            else if (input_vec_mag)
-                            {
-                                double mxy=sqrt(pow(atoms[it].m_loc_[ia].x,2)+pow(atoms[it].m_loc_[ia].y,2));
-                                atoms[it].angle1[ia]=atan2(mxy,atoms[it].m_loc_[ia].z);
-                                if(mxy>1e-8)
-                                {
-                                    atoms[it].angle2[ia]=atan2(atoms[it].m_loc_[ia].y,atoms[it].m_loc_[ia].x);
-                                }
-                            }
-                            else
-                            {
-                                atoms[it].m_loc_[ia].x = 0;
-                                atoms[it].m_loc_[ia].y = 0;
-                                atoms[it].m_loc_[ia].z = atoms[it].mag[ia];
-                            }
+                            atoms[it].m_loc_[ia].x = atoms[it].mag[ia] *
+                                sin(atoms[it].angle1[ia]) * cos(atoms[it].angle2[ia]);
+                            atoms[it].m_loc_[ia].y = atoms[it].mag[ia] *
+                                sin(atoms[it].angle1[ia]) * sin(atoms[it].angle2[ia]);
                         }
-                        else
+                    }
+                    else if (input_vec_mag)
+                    {// mx, my, mz are given, calculate angle1 and angle2 from mx, my, mz
+                        double mxy=sqrt(pow(atoms[it].m_loc_[ia].x,2)+pow(atoms[it].m_loc_[ia].y,2));
+                        atoms[it].angle1[ia]=atan2(mxy,atoms[it].m_loc_[ia].z);
+                        if(mxy>1e-8)
                         {
+                            atoms[it].angle2[ia]=atan2(atoms[it].m_loc_[ia].y,atoms[it].m_loc_[ia].x);
+                        }
+                    }
+                    else// only one mag is given, assume it is z
+                    {
+                        atoms[it].m_loc_[ia].x = 0;
+                        atoms[it].m_loc_[ia].y = 0;
+                        atoms[it].m_loc_[ia].z = atoms[it].mag[ia];
+                    }
+
+                    if(PARAM.inp.nspin==4)
+                    {
+                        if(!PARAM.inp.noncolin)
+                        {
+                            //collinear case with nspin = 4, only z component is used
                             atoms[it].m_loc_[ia].x = 0;
                             atoms[it].m_loc_[ia].y = 0;
-                            atoms[it].m_loc_[ia].z = atoms[it].mag[ia];
                         }
-
                         //print only ia==0 && mag>0 to avoid too much output
                         //print when ia!=0 && mag[ia] != mag[0] to avoid too much output
                         //  'A || (!A && B)' is equivalent to 'A || B',so the following 
@@ -735,8 +734,8 @@ bool UnitCell::read_atom_positions(std::ifstream &ifpos, std::ofstream &ofs_runn
                         ModuleBase::GlobalFunc::ZEROS(magnet.ux_ ,3);
                     }
                     else if(PARAM.inp.nspin==2)
-                    {
-                        atoms[it].m_loc_[ia].x = atoms[it].mag[ia];
+                    {// collinear case with nspin = 2, only z component is used
+                        atoms[it].mag[ia] = atoms[it].m_loc_[ia].z;
                         //print only ia==0 && mag>0 to avoid too much output
                         //print when ia!=0 && mag[ia] != mag[0] to avoid too much output
                         if(ia==0 || (atoms[it].mag[ia] != atoms[it].mag[0]))
@@ -751,6 +750,8 @@ bool UnitCell::read_atom_positions(std::ifstream &ifpos, std::ofstream &ofs_runn
                             ModuleBase::GlobalFunc::OUT(ofs_running, ss.str(),atoms[it].mag[ia]);
                         }
                     }
+                    // end of calculating initial magnetization of each atom
+                    // ----------------------------------------------------------------------------
 
                     if(Coordinate=="Direct")
                     {

source/module_elecstate/elecstate_pw.cpp

@@ -31,36 +31,55 @@ ElecStatePW<T, Device>::ElecStatePW(ModulePW::PW_Basis_K* wfc_basis_in,
 template<typename T, typename Device>
 ElecStatePW<T, Device>::~ElecStatePW() 
 {
-    if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+    if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single")
     {
         delmem_var_op()(this->ctx, this->rho_data);
+        delete[] this->rho;
+
+        if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp)
+        {
+            delmem_complex_op()(this->ctx, this->rhog_data);
+            delete[] this->rhog;
+        }
         if (get_xc_func_type() == 3 || PARAM.inp.out_elf[0] > 0)
         {
             delmem_var_op()(this->ctx, this->kin_r_data);
+            delete[] this->kin_r;
         }
     }
-    if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single") {
-        delete[] this->rho;
-        delete[] this->kin_r;
+    if (PARAM.globalv.use_uspp)
+    {
+        delmem_var_op()(this->ctx, this->becsum);
     }
-    delmem_var_op()(this->ctx, becsum);
     delmem_complex_op()(this->ctx, this->wfcr);
     delmem_complex_op()(this->ctx, this->wfcr_another_spin);
 }
 
 template<typename T, typename Device>
 void ElecStatePW<T, Device>::init_rho_data() 
 {
-    if(this->init_rho) {
+    if (this->init_rho)
+    {
         return;
     }
-    
-    if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single") {
+
+    if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single")
+    {
         this->rho = new Real*[this->charge->nspin];
         resmem_var_op()(this->ctx, this->rho_data, this->charge->nspin * this->charge->nrxx);
-        for (int ii = 0; ii < this->charge->nspin; ii++) {
+        for (int ii = 0; ii < this->charge->nspin; ii++)
+        {
             this->rho[ii] = this->rho_data + ii * this->charge->nrxx;
         }
+        if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp)
+        {
+            this->rhog = new T*[this->charge->nspin];
+            resmem_complex_op()(this->ctx, this->rhog_data, this->charge->nspin * this->charge->rhopw->npw);
+            for (int ii = 0; ii < this->charge->nspin; ii++)
+            {
+                this->rhog[ii] = this->rhog_data + ii * this->charge->rhopw->npw;
+            }
+        }
         if (get_xc_func_type() == 3 || PARAM.inp.out_elf[0] > 0)
         {
             this->kin_r = new Real*[this->charge->nspin];
@@ -70,8 +89,13 @@ void ElecStatePW<T, Device>::init_rho_data()
             }
         }
     }
-    else {
+    else
+    {
         this->rho = reinterpret_cast<Real **>(this->charge->rho);
+        if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp)
+        {
+            this->rhog = reinterpret_cast<T**>(this->charge->rhog);
+        }
         if (get_xc_func_type() == 3 || PARAM.inp.out_elf[0] > 0)
         {
             this->kin_r = reinterpret_cast<Real **>(this->charge->kin_r);
@@ -100,19 +124,24 @@ void ElecStatePW<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)
             // ModuleBase::GlobalFunc::ZEROS(this->charge->kin_r[is], this->charge->nrxx);
             setmem_var_op()(this->ctx, this->kin_r[is], 0,  this->charge->nrxx);
         }
-	}
+        if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp)
+        {
+            setmem_complex_op()(this->ctx, this->rhog[is], 0, this->charge->rhopw->npw);
+        }
+    }
 
     for (int ik = 0; ik < psi.get_nk(); ++ik)
     {
         psi.fix_k(ik);
         this->updateRhoK(psi);
     }
-    if (PARAM.globalv.use_uspp)
+
+    this->add_usrho(psi);
+
+    if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single")
     {
-        this->add_usrho(psi);
-    }
-    if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single") {
-        for (int ii = 0; ii < PARAM.inp.nspin; ii++) {
+        for (int ii = 0; ii < PARAM.inp.nspin; ii++)
+        {
             castmem_var_d2h_op()(cpu_ctx, this->ctx, this->charge->rho[ii], this->rho[ii], this->charge->nrxx);
             if (get_xc_func_type() == 3)
             {
@@ -397,32 +426,39 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
 template <typename T, typename Device>
 void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)
 {
-    this->cal_becsum(psi);
+    if (PARAM.globalv.use_uspp)
+    {
+        this->cal_becsum(psi);
+    }
 
     // transform soft charge to recip space using smooth grids
-    T* rhog = nullptr;
-    resmem_complex_op()(this->ctx, rhog, this->charge->rhopw->npw * PARAM.inp.nspin, "ElecState<PW>::rhog");
-    setmem_complex_op()(this->ctx, rhog, 0, this->charge->rhopw->npw * PARAM.inp.nspin);
-    for (int is = 0; is < PARAM.inp.nspin; is++)
+    if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp)
     {
-        this->rhopw_smooth->real2recip(this->rho[is], &rhog[is * this->charge->rhopw->npw]);
+        for (int is = 0; is < PARAM.inp.nspin; is++)
+        {
+            this->rhopw_smooth->real2recip(this->rho[is], this->rhog[is]);
+        }
     }
 
     // \sum_lm Q_lm(r) \sum_i <psi_i|beta_l><beta_m|psi_i> w_i
     // add to the charge density in reciprocal space the part which is due to the US augmentation.
-    this->addusdens_g(becsum, rhog);
+    if (PARAM.globalv.use_uspp)
+    {
+        this->addusdens_g(becsum, rhog);
+    }
 
     // transform back to real space using dense grids
-    for (int is = 0; is < PARAM.inp.nspin; is++)
+    if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp)
     {
-        this->charge->rhopw->recip2real(&rhog[is * this->charge->rhopw->npw], this->rho[is]);
+        for (int is = 0; is < PARAM.inp.nspin; is++)
+        {
+            this->charge->rhopw->recip2real(this->rhog[is], this->rho[is]);
+        }
     }
-
-    delmem_complex_op()(this->ctx, rhog);
 }
 
 template <typename T, typename Device>
-void ElecStatePW<T, Device>::addusdens_g(const Real* becsum, T* rhog)
+void ElecStatePW<T, Device>::addusdens_g(const Real* becsum, T** rhog)
 {
     const T one{1, 0};
     const T zero{0, 0};
@@ -506,7 +542,7 @@ void ElecStatePW<T, Device>::addusdens_g(const Real* becsum, T* rhog)
                         this->ppcell->radial_fft_q(this->ctx, npw, ih, jh, it, qmod, ylmk0, qgm);
                         for (int ig = 0; ig < npw; ig++)
                         {
-                            rhog[is * npw + ig] += qgm[ig] * aux2[ijh * npw + ig];
+                            rhog[is][ig] += qgm[ig] * aux2[ijh * npw + ig];
                         }
                         ijh++;
                     }

source/module_elecstate/elecstate_pw.h

@@ -42,8 +42,9 @@ class ElecStatePW : public ElecState
 
     //! init rho_data and kin_r_data
     void init_rho_data();
-    Real** rho = nullptr; 
-    Real** kin_r = nullptr; //[Device] [spin][nrxx] rho and kin_r
+    Real** rho = nullptr;   // [Device] [spin][nrxx] rho
+    T** rhog = nullptr;     // [Device] [spin][nrxx] rhog
+    Real** kin_r = nullptr; // [Device] [spin][nrxx] kin_r
 
   protected:
 
@@ -70,15 +71,16 @@ class ElecStatePW : public ElecState
 
     //! Non-local pseudopotentials
     //! \sum_lm Q_lm(r) \sum_i <psi_i|beta_l><beta_m|psi_i> w_i
-    void addusdens_g(const Real* becsum, T* rhog);
+    void addusdens_g(const Real* becsum, T** rhog);
 
     Device * ctx = {};
 
     bool init_rho = false;
 
     mutable T* vkb = nullptr;
 
-    Real* rho_data = nullptr; 
+    Real* rho_data = nullptr;
+    T* rhog_data = nullptr;
     Real* kin_r_data = nullptr;
     T* wfcr = nullptr; 
     T* wfcr_another_spin = nullptr;

source/module_esolver/lcao_before_scf.cpp

@@ -205,17 +205,19 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
     }
 
 #ifdef __DEEPKS
-    // for each ionic step, the overlap <psi|alpha> must be rebuilt
+    // for each ionic step, the overlap <phi|alpha> must be rebuilt
     // since it depends on ionic positions
     if (PARAM.globalv.deepks_setorb)
     {
         const Parallel_Orbitals* pv = &this->pv;
-        // build and save <psi(0)|alpha(R)> at beginning
-        GlobalC::ld.build_psialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));
+        // allocate <phi(0)|alpha(R)>, phialpha is different every ion step, so it is allocated here
+        GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+        // build and save <phi(0)|alpha(R)> at beginning
+        GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));
 
         if (PARAM.inp.deepks_out_unittest)
         {
-            GlobalC::ld.check_psialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+            GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
         }
     }
 #endif

source/module_esolver/lcao_others.cpp

@@ -211,17 +211,19 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)
     }
 
 #ifdef __DEEPKS
-    // for each ionic step, the overlap <psi|alpha> must be rebuilt
+    // for each ionic step, the overlap <phi|alpha> must be rebuilt
     // since it depends on ionic positions
     if (PARAM.globalv.deepks_setorb)
     {
         const Parallel_Orbitals* pv = &this->pv;
-        // build and save <psi(0)|alpha(R)> at beginning
-        GlobalC::ld.build_psialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));
+        // allocate <phi(0)|alpha(R)>, phialpha is different every ion step, so it is allocated here
+        GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+        // build and save <phi(0)|alpha(R)> at beginning
+        GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));
 
         if (PARAM.inp.deepks_out_unittest)
         {
-            GlobalC::ld.check_psialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+            GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
         }
     }
 #endif