Paw : added accumulate_rhoij

source/module_cell/module_paw/CMakeLists.txt

@@ -2,7 +2,6 @@ add_library(
     paw
     OBJECT
     paw_element.cpp
-    paw_interface.cpp
     paw_sphbes.cpp
     paw_cell.cpp
 )

source/module_cell/module_paw/paw_atom.cpp

@@ -9,6 +9,8 @@ void Paw_Atom::init_paw_atom(const int nproj_in)
     rhoij.resize(nproj*(nproj + 1) / 2);
     rhoijp.resize(nproj*(nproj + 1) / 2);
     rhoijselect.resize(nproj*(nproj + 1) / 2);
+
+    this -> reset_rhoij();
 }
 
 void Paw_Atom::set_ca(std::vector<std::complex<double>> & ca_in, const double weight_in)

source/module_cell/module_paw/paw_cell.cpp

@@ -4,6 +4,7 @@
 
 void Paw_Cell::init_paw_cell(
     const double ecutwfc_in, const double cell_factor_in,
+    const double omega_in,
     const int nat_in, const int ntyp_in,
     const int * atom_type_in, const double ** atom_coord_in,
     const std::vector<std::string> & filename_list_in,
@@ -17,6 +18,7 @@ void Paw_Cell::init_paw_cell(
     this -> nx = nx_in;
     this -> ny = ny_in;
     this -> nz = nz_in;
+    this -> omega = omega_in;
 
     atom_coord.resize(nat);
     atom_type.resize(nat);
@@ -31,6 +33,27 @@ void Paw_Cell::init_paw_cell(
         atom_type[iat] = atom_type_in[iat];
     }
 
+    nat_type.resize(ntyp);
+    atom_map.resize(nat);
+
+    // Not sure if these will be used, but it doesn't
+    // hurt to prepare them I suppose
+    // Doesn't take much time, also not much space
+    int ind = 0;
+    for(int it = 0; it < ntyp; it ++)
+    {
+        nat_type[it] = 0;
+        for(int iat = 0; iat < nat; iat ++)
+        {
+            if(atom_type[iat] == it)
+            {
+                nat_type[it] ++;
+                atom_map[ind] = iat;
+                ind ++;
+            }
+        }
+    }
+
     paw_element_list.resize(ntyp);
     assert(filename_list_in.size() == ntyp);
     for(int ityp = 0; ityp < ntyp; ityp ++)
@@ -41,6 +64,14 @@ void Paw_Cell::init_paw_cell(
 
     this->map_paw_proj();
 
+    paw_atom_list.resize(nat);
+    for(int iat = 0; iat < nat; iat ++)
+    {
+        int it = atom_type[iat];
+        int nproj = paw_element_list[it].get_mstates();
+        paw_atom_list[iat].init_paw_atom(nproj);
+    }
+
     eigts1.resize(nat);
     eigts2.resize(nat);
     eigts3.resize(nat);
@@ -70,15 +101,17 @@ void Paw_Cell::init_paw_cell(
 
 // exp(-i(k+G)R_I) = exp(-ikR_I) exp(-iG_xR_Ix) exp(-iG_yR_Iy) exp(-iG_zR_Iz)
 void Paw_Cell::set_paw_k(
-    const int npw, const double * kpt,
+    const int npw_in, const double * kpt,
     const int * ig_to_ix, const int * ig_to_iy, const int * ig_to_iz,
-    const double ** kpg)
+    const double ** kpg, const double tpiba)
 {
     ModuleBase::TITLE("Paw_Element","set_paw_k");
 
     const double pi = 3.141592653589793238462643383279502884197;
     const double twopi = 2.0 * pi;
 
+    this -> npw = npw_in;
+
     struc_fact.resize(nat);
     for(int iat = 0; iat < nat; iat ++)
     {
@@ -103,6 +136,12 @@ void Paw_Cell::set_paw_k(
     }
 
     this -> set_ylm(npw, kpg);
+
+    gnorm.resize(npw);
+    for(int ipw = 0; ipw < npw; ipw ++)
+    {
+        gnorm[ipw] = std::sqrt(kpg[ipw][0]*kpg[ipw][0] + kpg[ipw][1]*kpg[ipw][1] + kpg[ipw][2]*kpg[ipw][2]) * tpiba;
+    }
 }
 
 void Paw_Cell::map_paw_proj()
@@ -126,10 +165,12 @@ void Paw_Cell::map_paw_proj()
     iprj_to_il.resize(nproj_tot);
     iprj_to_l.resize(nproj_tot);
     iprj_to_m.resize(nproj_tot);
+    start_iprj_ia.resize(nat);
 
     int iproj = 0;
     for(int ia = 0; ia < nat; ia ++)
     {
+        start_iprj_ia[ia] = iproj;
         int it = atom_type[ia];
         int mstates = paw_element_list[it].get_mstates();
         std::vector<int> im_to_istate = paw_element_list[it].get_im_to_istate();
@@ -156,10 +197,10 @@ void Paw_Cell::map_paw_proj()
     assert(iproj == nproj_tot);
 }
 
-void Paw_Cell::set_ylm(const int npw, const double ** kpg)
+void Paw_Cell::set_ylm(const int npw_in, const double ** kpg)
 {
-    ylm_k.resize(npw);
-    for(int ipw = 0; ipw < npw; ipw ++)
+    ylm_k.resize(npw_in);
+    for(int ipw = 0; ipw < npw_in; ipw ++)
     {
         ylm_k[ipw] = calc_ylm(lmax, kpg[ipw]);
     }
@@ -289,4 +330,161 @@ double Paw_Cell::ass_leg_pol(const int l, const int m, const double arg)
             return pll;
         }
     }
+}
+
+void Paw_Cell::get_vkb()
+{
+    ModuleBase::TITLE("Paw_Cell","get_vkb");
+
+    // Not sure if the loop structure is the best
+    // may need to improve the efficiency later
+
+    // vkb = ptilde(G)
+    // = [\int f(r)r^2j_l(r)dr] * [(-i)^l] * [ylm(\hat{G})] * [exp(-GR_I)]
+    // = ptilde * fact * ylm * sk (in the code below)
+    vkb.resize(nproj_tot);
+    for(int iproj = 0; iproj < nproj_tot; iproj ++)
+    {
+        vkb[iproj].resize(npw);
+    }
+
+    const std::complex<double> neg_i = std::complex<double>(0.0,-1.0);
+
+    for(int ipw = 0; ipw < npw; ipw ++)
+    {
+        int ind_at = 0;
+        for(int it = 0; it < ntyp; it ++)
+        {
+            const int nstates = paw_element_list[it].get_nstates(); // number of l channels
+            const int nproj = paw_element_list[it].get_mstates(); // number of (l,m) channels
+
+            std::vector<double> ptilde;
+            ptilde.resize(nproj);
+
+            int iproj = 0;
+            for(int istate = 0; istate < nstates; istate ++)
+            {
+                const double tmp = paw_element_list[it].get_ptilde(istate,gnorm[ipw],omega);
+                const int l = paw_element_list[it].get_lstate()[istate];
+
+                for(int im = 0; im < 2*l+1; im ++)
+                {
+                    ptilde[iproj] = tmp;
+                    iproj ++;
+                }
+            }
+            assert(iproj == nproj);
+
+            const int na = nat_type[it];
+            for(int ia = 0; ia < na; ia ++)
+            {
+                const int iat = atom_map[ind_at+ia];
+                const int proj_start = start_iprj_ia[iat];
+
+                for(int ip = 0; ip < nproj; ip ++)
+                {
+                    const int il = paw_element_list[it].get_im_to_istate()[ip];
+                    const int l  = paw_element_list[it].get_lstate()[il];
+                    const int m  = paw_element_list[it].get_mstate()[ip];
+
+                    // I put these codes in comment here as a reminder of
+                    // how this part is treated differently from nonlocal PP
+                    // First of all, the order of m quantum number
+                    // in nonlocal PP (calculated using ylm_real) is 0,1,-1,2,-2,...
+                    // but -l,-l+1,-l+2,...,l here
+                    // Secondly, there is a factor (-1)^m which is present here
+                    // but not in ylm_real
+
+                    /*
+                    int m1 = 0;
+                    if(l==1 && m == -1) m1 = 0;
+                    if(l==1 && m == 0 ) m1 = 1;
+                    if(l==1 && m == 1 ) m1 = -1;
+                    */
+
+                    //const int lm_ind = l*l + l + m1;
+                    const int lm_ind = l*l + l + m;
+
+                    const std::complex<double> fact = std::pow(neg_i,l);
+
+                    //const double ylm = ylm_k[ipw][lm_ind] * std::pow(-1.0,m1);
+                    const double ylm = ylm_k[ipw][lm_ind];
+
+                    const std::complex<double> sk = struc_fact[iat][ipw];
+
+                    vkb[ip+proj_start][ipw] = fact * ptilde[ip] * ylm * sk;
+                }
+            }
+
+            ind_at += na;
+        }
+        assert(ind_at == nat);
+    }
+}
+
+void Paw_Cell::accumulate_rhoij(const std::complex<double> * psi, const double weight)
+{
+    ModuleBase::TITLE("Paw_Cell","accumulate_rhoij");
+
+    for(int iat = 0; iat < nat; iat ++)
+    {
+        // ca : <ptilde(G)|psi(G)>
+        // = \sum_G [\int f(r)r^2j_l(r)dr] * [(-i)^l] * [ylm(\hat{G})] * [exp(-GR_I)] *psi(G)
+        // = \sum_ipw ptilde * fact * ylm * sk * psi (in the code below)
+        // This is what is called 'becp' in nonlocal pp
+        std::vector<std::complex<double>> ca;
+
+        const int it = atom_type[iat];
+        const int nproj = paw_element_list[it].get_mstates();
+        const int proj_start = start_iprj_ia[iat];
+
+        ca.resize(nproj);
+
+        for(int iproj = 0; iproj < nproj; iproj ++)
+        {
+            ca[iproj] = 0.0;
+
+            // consider use blas subroutine for this part later
+            for(int ipw = 0; ipw < npw; ipw ++)
+            {
+                ca[iproj] += psi[ipw] * std::conj(vkb[iproj+proj_start][ipw]);
+            }
+        }
+
+        // ca should be summed over MPI ranks since planewave basis is distributed
+        // but not for now (I'll make sure serial version works first)
+        // Parallel_Reduce::reduce_complex_double_pool(ca.data(), nproj);
+
+        paw_atom_list[iat].set_ca(ca, weight);
+        paw_atom_list[iat].accumulate_rhoij();
+    }
+}
+
+std::vector<std::vector<double>> Paw_Cell::get_rhoij()
+{
+    std::vector<std::vector<double>> rhoij_all;
+    rhoij_all.resize(nat);
+    for(int iat = 0; iat < nat; iat ++)
+    {
+        rhoij_all[iat] = paw_atom_list[iat].get_rhoij();
+    }
+
+    return rhoij_all;
+}
+
+void Paw_Cell::get_rhoijp(std::vector<std::vector<double>> & rhoijp,
+        std::vector<std::vector<int>> & rhoijselect, std::vector<int> & nrhoijsel)
+{
+
+    rhoijp.resize(nat);
+    rhoijselect.resize(nat);
+    nrhoijsel.resize(nat);
+
+    for(int iat = 0; iat < nat; iat ++)
+    {
+        paw_atom_list[iat].convert_rhoij();
+        rhoijp[iat] = paw_atom_list[iat].get_rhoijp();
+        rhoijselect[iat] = paw_atom_list[iat].get_rhoijselect();
+        nrhoijsel[iat] = paw_atom_list[iat].get_nrhoijsel();
+    }
 }