Fix: various matrix elements calculations in orbital generation (#3940)

jinzx10 · kirk0830 · WHUweiqingzhou · web-flow · commit 840f0d5015fa · 2024-05-21T10:41:06.000+08:00
* jy-jy two-center integration

* change index set to vector

* more accurate algorithm for jy-jy &amp; unit fix

* enable non-LCAO build

* fix

* initialize local variables

* change the usage of ucell in Numerical_Basis from GlobalC to function arguments

* clean up function interface

* various bug fixed

* fix unit test

* add const qualifiers

---------

Co-authored-by: kirk0830 &lt;67682086+kirk0830@users.noreply.github.com&gt;
Co-authored-by: wqzhou &lt;33364058+WHUweiqingzhou@users.noreply.github.com&gt;
diff --git a/source/Makefile.Objects b/source/Makefile.Objects
@@ -429,6 +429,7 @@ OBJS_IO=input.o\
     dos_nao.o\
     numerical_descriptor.o\
     numerical_basis.o\
+    numerical_basis_jyjy.o\
     output.o\
     print_info.o\
     read_cube.o\
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -1314,7 +1314,7 @@ void ESolver_KS_PW<T, Device>::after_all_runners(void)
             if(INPUT.bessel_nao_rcuts.size() == 1)
             {
                 Numerical_Basis numerical_basis;
-                numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc);
+                numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc, GlobalC::ucell);
             }
             else
             {
@@ -1339,7 +1339,7 @@ void ESolver_KS_PW<T, Device>::after_all_runners(void)
                         Will be refactored in the future.
                     */
                     Numerical_Basis numerical_basis;
-                    numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc);
+                    numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc, GlobalC::ucell);
 					std::string old_fname_header = winput::spillage_outdir 
 						+ "/" 
 						+ "orb_matrix.";
diff --git a/source/module_io/CMakeLists.txt b/source/module_io/CMakeLists.txt
@@ -8,6 +8,7 @@ list(APPEND objects
     nscf_band.cpp
     write_istate_info.cpp
     numerical_basis.cpp
+    numerical_basis_jyjy.cpp
     numerical_descriptor.cpp
     output.cpp
     print_info.cpp
diff --git a/source/module_io/bessel_basis.cpp b/source/module_io/bessel_basis.cpp
@@ -99,9 +99,9 @@ double Bessel_Basis::Polynomial_Interpolation2
 	const double x3 = 3.0 - x0;
 	const double y=
 		this->TableOne(l, ie, iq) * x1 * x2 * x3 / 6.0 +
-        this->TableOne(l, ie, iq) * x0 * x2 * x3 / 2.0 -
-        this->TableOne(l, ie, iq) * x1 * x0 * x3 / 2.0 +
-        this->TableOne(l, ie, iq) * x1 * x2 * x0 / 6.0 ;
+        this->TableOne(l, ie, iq+1) * x0 * x2 * x3 / 2.0 -
+        this->TableOne(l, ie, iq+2) * x1 * x0 * x3 / 2.0 +
+        this->TableOne(l, ie, iq+3) * x1 * x2 * x0 / 6.0 ;
 	return y;
 }
 
@@ -117,9 +117,9 @@ double Bessel_Basis::Polynomial_Interpolation(
 	const double x3 = 3.0 - x0;
 	const double y=
 		this->Faln(it, l, ic, iq) * x1 * x2 * x3 / 6.0 +
-        this->Faln(it, l, ic, iq) * x0 * x2 * x3 / 2.0 -
-        this->Faln(it, l, ic, iq) * x1 * x0 * x3 / 2.0 +
-        this->Faln(it, l, ic, iq) * x1 * x2 * x0 / 6.0 ;
+        this->Faln(it, l, ic, iq+1) * x0 * x2 * x3 / 2.0 -
+        this->Faln(it, l, ic, iq+2) * x1 * x0 * x3 / 2.0 +
+        this->Faln(it, l, ic, iq+3) * x1 * x2 * x0 / 6.0 ;
 	return y;
 }
 
diff --git a/source/module_io/numerical_basis.cpp b/source/module_io/numerical_basis.cpp
diff --git a/source/module_io/numerical_basis.h b/source/module_io/numerical_basis.h
@@ -29,49 +29,58 @@ class Numerical_Basis
     Numerical_Basis();
     ~Numerical_Basis();
 
-    void start_from_file_k(const int& ik, ModuleBase::ComplexMatrix& psi, const Structure_Factor& sf, const ModulePW::PW_Basis_K* wfcpw);
-    void output_overlap(const psi::Psi<std::complex<double>>& psi, const Structure_Factor& sf, const K_Vectors& kv, const ModulePW::PW_Basis_K* wfcpw);
-
+    void start_from_file_k(const int& ik, ModuleBase::ComplexMatrix& psi, const Structure_Factor& sf, const ModulePW::PW_Basis_K* wfcpw, const UnitCell& ucell);
+    void output_overlap(const psi::Psi<std::complex<double>>& psi, const Structure_Factor& sf, const K_Vectors& kv, const ModulePW::PW_Basis_K* wfcpw, const UnitCell& ucell);
 
   private:
     bool init_label = false;
 
     Bessel_Basis bessel_basis;
 
     std::vector<ModuleBase::IntArray> mu_index;
-    static std::vector<ModuleBase::IntArray> init_mu_index(void);
+    static std::vector<ModuleBase::IntArray> init_mu_index(const UnitCell& ucell);
 
     void numerical_atomic_wfc(const int& ik,
                               const ModulePW::PW_Basis_K* wfcpw,
                               ModuleBase::ComplexMatrix& psi,
-                              const Structure_Factor& sf);
+                              const Structure_Factor& sf,
+                              const UnitCell& ucell);
 
     ModuleBase::ComplexArray cal_overlap_Q(const int& ik,
                                            const int& np,
                                            const ModulePW::PW_Basis_K* wfcpw,
                                            const psi::Psi<std::complex<double>>& psi,
                                            const double derivative_order,
-                                           const Structure_Factor& sf) const;
+                                           const Structure_Factor& sf,
+                                           const UnitCell& ucell) const;
 
+    // computed in the plane-wave basis
     ModuleBase::ComplexArray cal_overlap_Sq(const int& ik,
                                             const int& np,
                                             const double derivative_order,
                                             const Structure_Factor& sf,
-                                            const ModulePW::PW_Basis_K* wfcpw) const;
+                                            const ModulePW::PW_Basis_K* wfcpw,
+                                            const UnitCell& ucell) const;
 
     static ModuleBase::matrix cal_overlap_V(const ModulePW::PW_Basis_K* wfcpw,
                                             const psi::Psi<std::complex<double>>& psi,
                                             const double derivative_order,
-                                            const K_Vectors& kv);
+                                            const K_Vectors& kv,
+                                            const double tpiba);
 
-    ModuleBase::realArray cal_flq(const int ik, const std::vector<ModuleBase::Vector3<double>> &gk) const;
+    // gk should be in the atomic unit (Bohr)
+    ModuleBase::realArray cal_flq(const std::vector<ModuleBase::Vector3<double>> &gk,
+                                  const int ucell_lmax) const;
 
-    static ModuleBase::matrix cal_ylm(const std::vector<ModuleBase::Vector3<double>> &gk);
+    // Ylm does not depend on the magnitude so unit is not important
+    static ModuleBase::matrix cal_ylm(const std::vector<ModuleBase::Vector3<double>> &gk,
+                                      const int ucell_lmax);
 
+    // gk and the returned gpow are both in the atomic unit (Bohr)
     static std::vector<double> cal_gpow(const std::vector<ModuleBase::Vector3<double>> &gk,
                                         const double derivative_order);
 
-    static void output_info(std::ofstream& ofs, const Bessel_Basis& bessel_basis, const K_Vectors& kv);
+    static void output_info(std::ofstream& ofs, const Bessel_Basis& bessel_basis, const K_Vectors& kv, const UnitCell& ucell);
 
     static void output_k(std::ofstream& ofs, const K_Vectors& kv);
 
diff --git a/source/module_io/numerical_basis_jyjy.cpp b/source/module_io/numerical_basis_jyjy.cpp
@@ -0,0 +1,96 @@
+#include "module_io/numerical_basis_jyjy.h"
+#include "module_basis/module_nao/two_center_integrator.h"
+
+namespace NumericalBasis
+{
+#ifdef __LCAO
+using index_t = std::tuple<int, int, int, int>;
+
+std::vector<index_t> indexgen(const std::vector<int>& natom,
+                              const std::vector<int>& lmax)
+{
+#ifdef __DEBUG
+    assert(natom.size() == lmax.size());
+#endif
+    std::vector<index_t> index;
+    for (size_t itype = 0; itype < natom.size(); ++itype)
+    {
+        for (int iatom = 0; iatom < natom[itype]; ++iatom)
+        {
+            for (int l = 0; l <= lmax[itype]; ++l)
+            {
+                for (int M = 0; M < 2*l+1; ++M)
+                {
+                    // convert the "abacus M" to the conventional m
+                    const int m = (M % 2 == 0) ? -M/2 : (M+1)/2;
+                    index.emplace_back(itype, iatom, l, m);
+                }
+            }
+        }
+    }
+    return index;
+}
+
+
+ModuleBase::ComplexArray cal_overlap_Sq(
+    const char type,
+    const int lmax,
+    const int nbes,
+    const double rcut,
+    const std::vector<std::vector<ModuleBase::Vector3<double>>>& R,
+    const std::vector<index_t>& mu_index
+)
+{
+    // allocate output array
+    const int nlocal = mu_index.size();
+    ModuleBase::ComplexArray overlap_Sq(nlocal, nlocal, nbes, nbes);
+    overlap_Sq.zero_out();
+
+    // build a RadialCollection of spherical Bessel functions
+    const double dr = 0.005; // grid spacing for SphbesRadials; smaller for higher precision
+    RadialCollection orb;
+    orb.build(lmax, nbes, rcut, 0.0, dr);
+
+    ModuleBase::SphericalBesselTransformer sbt;
+    orb.set_transformer(sbt);
+
+    const double rmax = orb.rcut_max() * 2.0;
+    const int nr = static_cast<int>(rmax / dr) + 1;
+    
+    orb.set_uniform_grid(true, nr, rmax, 'i', true);
+
+    // build the two-center integrator
+    TwoCenterIntegrator intor;
+    intor.tabulate(orb, orb, type, nr, rmax);
+
+    // traverse the vector of composite index (itype, iatom, l, m)
+    int t1 = 0, a1 = 0, l1 = 0, m1 = 0;
+    int t2 = 0, a2 = 0, l2 = 0, m2 = 0;
+    for (auto it1 = mu_index.cbegin(); it1 != mu_index.cend(); ++it1)
+    {
+        std::tie(t1, a1, l1, m1) = *it1;
+        for (auto it2 = mu_index.cbegin(); it2 != mu_index.cend(); ++it2)
+        {
+            std::tie(t2, a2, l2, m2) = *it2;
+            ModuleBase::Vector3<double> dR = R[t2][a2] - R[t1][a1];
+            for (int zeta1 = 0; zeta1 < nbes; ++zeta1)
+            {
+                for (int zeta2 = 0; zeta2 < nbes; ++zeta2)
+                {
+                    double elem = 0.0;
+                    intor.calculate(t1, l1, zeta1, m1,
+                                    t2, l2, zeta2, m2,
+                                    dR, &elem);
+                    overlap_Sq(it1 - mu_index.begin(),
+                               it2 - mu_index.begin(),
+                               zeta1, zeta2) = elem;
+                }
+            }
+        }
+    }
+
+    return overlap_Sq;
+}
+#endif
+
+} // end of namespace NumericalBasis
diff --git a/source/module_io/numerical_basis_jyjy.h b/source/module_io/numerical_basis_jyjy.h
@@ -0,0 +1,38 @@
+#ifndef NUMERICAL_BASIS_JYJY_H
+#define NUMERICAL_BASIS_JYJY_H
+
+#include <tuple>
+#include <vector>
+#include "module_base/vector3.h"
+#include "module_base/complexarray.h"
+
+namespace NumericalBasis
+{
+std::vector<std::tuple<int, int, int, int>> indexgen(const std::vector<int>& natom,
+                                                     const std::vector<int>& lmax);
+
+/**
+ * @brief <jy|op|jy> overlap matrix (two-center integration)
+ *
+ *
+ * @param[in]   type                'S' (op = 1) or 'T' (kinetic, op = -\nabla^2)
+ * @param[in]   lmax                maximum angular momentum
+ * @param[in]   nbes                number of Bessel functions
+ * @param[in]   rcut                cutoff radius
+ * @param[in]   sigma               smoothing parameter
+ * @param[in]   tau_cart            atomic positions (in Bohr)
+ * @param[in]   mu_index            composite index
+ *
+ */
+ModuleBase::ComplexArray cal_overlap_Sq(
+    const char type, // 'S' or 'T'
+    const int lmax,
+    const int nbes,
+    const double rcut,
+    const std::vector<std::vector<ModuleBase::Vector3<double>>>& tau_cart,
+    const std::vector<std::tuple<int, int, int, int>>& mu_index
+);
+
+}
+
+#endif
diff --git a/source/module_io/test/CMakeLists.txt b/source/module_io/test/CMakeLists.txt
@@ -173,3 +173,9 @@ add_test(NAME read_wfc_pw_test_parallel
       COMMAND mpirun -np 4 ./read_wfc_pw_test
       WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
 )
+
+AddTest(
+  TARGET numerical_basis_test
+  LIBS base ${math_libs} device numerical_atomic_orbitals container orb
+  SOURCES numerical_basis_test.cpp ../numerical_basis_jyjy.cpp
+)
diff --git a/source/module_io/test/bessel_basis_test.cpp b/source/module_io/test/bessel_basis_test.cpp
@@ -484,9 +484,9 @@ TEST_F(TestBesselBasis, PolynomialInterpolation2Test) {
         i_Lmax, d_dr, d_dk
     );
     double d_yExpected = vvv_d_TableOne[0][0][i_position]*d_x1*d_x2*d_x3/6.0+
-                         vvv_d_TableOne[0][0][i_position]*d_x0*d_x2*d_x3/2.0-
-                         vvv_d_TableOne[0][0][i_position]*d_x1*d_x0*d_x3/2.0+
-                         vvv_d_TableOne[0][0][i_position]*d_x1*d_x2*d_x0/6.0;
+                         vvv_d_TableOne[0][0][i_position+1]*d_x0*d_x2*d_x3/2.0-
+                         vvv_d_TableOne[0][0][i_position+2]*d_x1*d_x0*d_x3/2.0+
+                         vvv_d_TableOne[0][0][i_position+3]*d_x1*d_x2*d_x0/6.0;
     double d_yTested = besselBasis.Polynomial_Interpolation2(0, 0, d_Gnorm);
     EXPECT_NEAR(d_yExpected, d_yTested, 0.01);
 }
@@ -552,9 +552,9 @@ TEST_F(TestBesselBasis, PolynomialInterpolationTest) {
         "./support/BesselBasis_UnitTest_C4_AtomType0.html", i_Ntype, i_Lmax, 1, i_Nq, vvv_d_TableOne
     );
     double d_yExpected = vvvv_d_Faln[0][0][0][i_position]*d_x1*d_x2*d_x3/6.0+
-                         vvvv_d_Faln[0][0][0][i_position]*d_x0*d_x2*d_x3/2.0-
-                         vvvv_d_Faln[0][0][0][i_position]*d_x1*d_x0*d_x3/2.0+
-                         vvvv_d_Faln[0][0][0][i_position]*d_x1*d_x2*d_x0/6.0;
+                         vvvv_d_Faln[0][0][0][i_position+1]*d_x0*d_x2*d_x3/2.0-
+                         vvvv_d_Faln[0][0][0][i_position+2]*d_x1*d_x0*d_x3/2.0+
+                         vvvv_d_Faln[0][0][0][i_position+3]*d_x1*d_x2*d_x0/6.0;
     double d_yTested = besselBasis.Polynomial_Interpolation(0, 0, 0, d_Gnorm);
     EXPECT_NEAR(d_yExpected, d_yTested, 0.01);
 }
diff --git a/source/module_io/test/numerical_basis_test.cpp b/source/module_io/test/numerical_basis_test.cpp

Original file line number	Diff line number	Diff line change
`@@ -1314,7 +1314,7 @@ void ESolver_KS_PW<T, Device>::after_all_runners(void)`
`1314`	`1314`	`if(INPUT.bessel_nao_rcuts.size() == 1)`
`1315`	`1315`	`{`
`1316`	`1316`	`Numerical_Basis numerical_basis;`
`1317`		`- numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc);`
	`1317`	`+ numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc, GlobalC::ucell);`
`1318`	`1318`	`}`
`1319`	`1319`	`else`
`1320`	`1320`	`{`
`@@ -1339,7 +1339,7 @@ void ESolver_KS_PW<T, Device>::after_all_runners(void)`
`1339`	`1339`	`Will be refactored in the future.`
`1340`	`1340`	`*/`
`1341`	`1341`	`Numerical_Basis numerical_basis;`
`1342`		`- numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc);`
	`1342`	`+ numerical_basis.output_overlap(this->psi[0], this->sf, this->kv, this->pw_wfc, GlobalC::ucell);`
`1343`	`1343`	`std::string old_fname_header = winput::spillage_outdir`
`1344`	`1344`	`+ "/"`
`1345`	`1345`	`+ "orb_matrix.";`