Test: Add unit tests and integrate tests for new features of DeePKS. (#6140)

* Add integrate test for v_delta(k) label in DeePKS multi-k calculations.

* Add orbital related unit test for DeePKS.

* Add vdpre related unit test for DeePKS.

* Add check for complex numbers.

* Change DeePKS multik UT to minimize file size.

* Update gedm_ref.dat.

* Rearrange and modify tests in DeePKS.

* Fix nscf bug caused by different filename.

* Update DeePKS multi-k UT.

* Rename jd as README.

* Update result.ref.

* Update result.ref.

* Adjust integrate tests in DeePKS for acceleration.

* Remove useless files in DeePKS integrate tests.

* Add nscf and relax integrate tests for DeePKS.

* Add input check.

* Add missing files for DeePKS test.

* Temporarily close the check for gedm in DeePKS UT.

Author: ErjieWu

source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp

@@ -690,42 +690,39 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
         } // end symmetry
 
 #ifdef __DEEPKS
-        if (PARAM.inp.deepks_out_labels) // not parallelized yet
+        if (PARAM.inp.deepks_out_labels == 1)
         {
-            if (PARAM.inp.deepks_out_labels == 1)
+            const std::string file_stot = PARAM.globalv.global_out_dir + "deepks_stot.npy";
+            LCAO_deepks_io::save_matrix2npy(file_stot,
+                                            scs,
+                                            GlobalV::MY_RANK,
+                                            ucell.omega,
+                                            'U'); // change to energy unit Ry when printing, S_tot;
+
+            const std::string file_sbase = PARAM.globalv.global_out_dir + "deepks_sbase.npy";
+            if (PARAM.inp.deepks_scf)
             {
-                const std::string file_stot = PARAM.globalv.global_out_dir + "deepks_stot.npy";
-                LCAO_deepks_io::save_matrix2npy(file_stot,
-                                                scs,
+                LCAO_deepks_io::save_matrix2npy(file_sbase,
+                                                scs - svnl_dalpha,
                                                 GlobalV::MY_RANK,
                                                 ucell.omega,
-                                                'U'); // change to energy unit Ry when printing, S_tot;
-
-                const std::string file_sbase = PARAM.globalv.global_out_dir + "deepks_sbase.npy";
-                if (PARAM.inp.deepks_scf)
-                {
-                    LCAO_deepks_io::save_matrix2npy(file_sbase,
-                                                    scs - svnl_dalpha,
-                                                    GlobalV::MY_RANK,
-                                                    ucell.omega,
-                                                    'U'); // change to energy unit Ry when printing, S_base;
-                }
-                else
-                {
-                    LCAO_deepks_io::save_matrix2npy(file_sbase,
-                                                    scs,
-                                                    GlobalV::MY_RANK,
-                                                    ucell.omega,
-                                                    'U'); // sbase = stot
-                }
+                                                'U'); // change to energy unit Ry when printing, S_base;
             }
-            else if (PARAM.inp.deepks_out_labels == 2)
+            else
             {
-                const std::string file_stot = PARAM.globalv.global_out_dir + "deepks_stress.npy";
-                LCAO_deepks_io::save_matrix2npy(file_stot, scs, GlobalV::MY_RANK, ucell.omega,
-                                                'F'); // flat mode
+                LCAO_deepks_io::save_matrix2npy(file_sbase,
+                                                scs,
+                                                GlobalV::MY_RANK,
+                                                ucell.omega,
+                                                'U'); // sbase = stot
             }
         }
+        else if (PARAM.inp.deepks_out_labels == 2)
+        {
+            const std::string file_stot = PARAM.globalv.global_out_dir + "deepks_stress.npy";
+            LCAO_deepks_io::save_matrix2npy(file_stot, scs, GlobalV::MY_RANK, ucell.omega,
+                                            'F'); // flat mode
+        }
 #endif
 
         // print Rydberg stress or not

source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp

@@ -233,7 +233,6 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::calculate_HR()
     const Parallel_Orbitals* paraV = this->V_delta_R->get_paraV();
     const int npol = this->ucell->get_npol();
 
-    // 1. calculate <phi|alpha> for each pair of atoms
     #pragma omp parallel for schedule(dynamic)
     for (int iat0 = 0; iat0 < this->ucell->nat; iat0++)
     {
@@ -290,9 +289,7 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::calculate_HR()
             }
         }
         const int trace_alpha_size = trace_alpha_row.size();
-        //--------------------------------------------------
 
-        // 2. calculate <phi_I|beta>D<beta|phi_{J,R}> for each pair of <IJR> atoms
         for (int ad1 = 0; ad1 < adjs.adj_num + 1; ++ad1)
         {
             const int T1 = adjs.ntype[ad1];
@@ -363,17 +360,6 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::calculate_HR()
                         ps2t[i] = col_ptr[trace_alpha_col[i]];
                     }
                 }
-                /*for(int irow = 0;irow<row_size;irow++)
-                {
-                    for(int icol=0;icol<col_size;icol++)
-                    {
-                        for(int ialpha=0;ialpha<trace_alpha_size;ialpha++)
-                        {
-                            tmp->get_pointer()[irow*col_size+icol] +=
-                                s_1t[irow*trace_alpha_size+ialpha] * s_2t[icol*trace_alpha_size+ialpha];
-                        }
-                    }
-                }*/
                 // dgemm for s_2t and s_1t to get HR_12
                 constexpr char transa = 'T', transb = 'N';
                 const double gemm_alpha = 1.0, gemm_beta = 1.0;

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -58,14 +58,15 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
     const int nspin = PARAM.inp.nspin;
     const int nk = nks / nspin;
 
+    // Update DMR in any case of deepks_out_labels/deepks_scf
+    DeePKS_domain::update_dmr(kvec_d, dm->get_DMK_vector(), ucell, orb, *ParaV, GridD, dmr);
+
     // Note : update PDM and all other quantities with the current dm
     // DeePKS PDM and descriptor
     if (PARAM.inp.deepks_out_labels == 1 || PARAM.inp.deepks_scf)
     {
         // this part is for integrated test of deepks
         // so it is printed no matter even if deepks_out_labels is not used
-        DeePKS_domain::update_dmr(kvec_d, dm->get_DMK_vector(), ucell, orb, *ParaV, GridD, dmr);
-        
         DeePKS_domain::cal_pdm<
             TK>(init_pdm, inlmax, lmaxd, inl2l, inl_index, kvec_d, dmr, phialpha, ucell, orb, GridD, *ParaV, pdm);
 

source/module_hamilt_lcao/module_deepks/deepks_orbital.cpp

@@ -56,6 +56,16 @@ void DeePKS_domain::cal_o_delta(const std::vector<TH>& dm_hl,
     return;
 }
 
+void DeePKS_domain::check_o_delta(ModuleBase::matrix& o_delta)
+{
+    std::ofstream ofs("o_delta.dat");
+    ofs << std::setprecision(10);
+    for (int ik = 0; ik < o_delta.nr; ik++)
+    {
+        ofs << o_delta(ik, 0) << std::endl;
+    }
+}
+
 template void DeePKS_domain::cal_o_delta<double, ModuleBase::matrix>(const std::vector<ModuleBase::matrix>& dm_hl,
                                                                      const std::vector<std::vector<double>>& h_delta,
                                                                      //  std::vector<double>& o_delta,

source/module_hamilt_lcao/module_deepks/deepks_orbital.h

@@ -31,6 +31,8 @@ void cal_o_delta(const std::vector<TH>& dm_hl,
                  const Parallel_Orbitals& pv,
                  const int nks,
                  const int nspin);
+
+void check_o_delta(ModuleBase::matrix& o_delta);
 } // namespace DeePKS_domain
 
 #endif

source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp

@@ -295,6 +295,25 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector<TH>& dm_hl,
     return;
 }
 
+void DeePKS_domain::check_orbpre(const torch::Tensor& orbpre)
+{
+    auto sizes = orbpre.sizes();
+    auto accessor = orbpre.accessor<double, 3>();
+    std::ofstream ofs("orbital_precalc.dat");
+    for (int iknb = 0; iknb < sizes[0]; iknb++)
+    {
+        for (int iat = 0; iat < sizes[1]; iat++)
+        {
+            for (int m = 0; m < sizes[2]; m++)
+            {
+                ofs << accessor[iknb][iat][m] << " ";
+            }
+            ofs << std::endl;
+        }
+        ofs << std::endl;
+    }
+}
+
 template void DeePKS_domain::cal_orbital_precalc<double, ModuleBase::matrix>(
     const std::vector<ModuleBase::matrix>& dm_hl,
     const int lmaxd,

source/module_hamilt_lcao/module_deepks/deepks_orbpre.h

@@ -41,6 +41,8 @@ void cal_orbital_precalc(const std::vector<TH>& dm_hl,
                          const Parallel_Orbitals& pv,
                          const Grid_Driver& GridD,
                          torch::Tensor& orbital_precalc);
+
+void check_orbpre(const torch::Tensor& orbpre);
 } // namespace DeePKS_domain
 #endif
 #endif

source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp

@@ -458,7 +458,7 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
 
 void DeePKS_domain::check_pdm(const int inlmax, const std::vector<int>& inl2l, const std::vector<torch::Tensor>& pdm)
 {
-    const std::string file_projdm = PARAM.globalv.global_out_dir + "pdm.dat";
+    const std::string file_projdm = PARAM.globalv.global_out_dir + "deepks_projdm.dat";
     std::ofstream ofs(file_projdm.c_str());
 
     ofs << std::setprecision(10);

source/module_hamilt_lcao/module_deepks/deepks_vdpre.cpp

@@ -172,28 +172,25 @@ void DeePKS_domain::cal_v_delta_precalc(const int nlocal,
 }
 
 template <typename TK>
-void DeePKS_domain::check_v_delta_precalc(const int nat,
-                                          const int nks,
-                                          const int nlocal,
-                                          const int des_per_atom,
-                                          const torch::Tensor& v_delta_precalc)
+void DeePKS_domain::check_v_delta_precalc(const torch::Tensor& v_delta_precalc)
 {
     using TK_tensor =
         typename std::conditional<std::is_same<TK, std::complex<double>>::value, c10::complex<double>, TK>::type;
+    auto sizes = v_delta_precalc.sizes();
     std::ofstream ofs("v_delta_precalc.dat");
     ofs << std::setprecision(10);
     auto accessor
         = v_delta_precalc
               .accessor<std::conditional_t<std::is_same<TK, double>::value, double, c10::complex<double>>, 5>();
-    for (int iks = 0; iks < nks; ++iks)
+    for (int iks = 0; iks < sizes[0]; ++iks)
     {
-        for (int mu = 0; mu < nlocal; ++mu)
+        for (int mu = 0; mu < sizes[1]; ++mu)
         {
-            for (int nu = 0; nu < nlocal; ++nu)
+            for (int nu = 0; nu < sizes[2]; ++nu)
             {
-                for (int iat = 0; iat < nat; ++iat)
+                for (int iat = 0; iat < sizes[3]; ++iat)
                 {
-                    for (int p = 0; p < des_per_atom; ++p)
+                    for (int p = 0; p < sizes[4]; ++p)
                     {
                         TK_tensor tmp = accessor[iks][mu][nu][iat][p];
                         TK* tmp_ptr = reinterpret_cast<TK*>(&tmp);
@@ -432,16 +429,8 @@ template void DeePKS_domain::cal_v_delta_precalc<std::complex<double>>(
     const Grid_Driver& GridD,
     torch::Tensor& v_delta_precalc);
 
-template void DeePKS_domain::check_v_delta_precalc<double>(const int nat,
-                                                           const int nks,
-                                                           const int nlocal,
-                                                           const int des_per_atom,
-                                                           const torch::Tensor& v_delta_precalc);
-template void DeePKS_domain::check_v_delta_precalc<std::complex<double>>(const int nat,
-                                                                         const int nks,
-                                                                         const int nlocal,
-                                                                         const int des_per_atom,
-                                                                         const torch::Tensor& v_delta_precalc);
+template void DeePKS_domain::check_v_delta_precalc<double>(const torch::Tensor& v_delta_precalc);
+template void DeePKS_domain::check_v_delta_precalc<std::complex<double>>(const torch::Tensor& v_delta_precalc);
 
 template void DeePKS_domain::prepare_phialpha<double>(const int nlocal,
                                                       const int lmaxd,

source/module_hamilt_lcao/module_deepks/deepks_vdpre.h

@@ -51,11 +51,7 @@ void cal_v_delta_precalc(const int nlocal,
                          torch::Tensor& v_delta_precalc);
 
 template <typename TK>
-void check_v_delta_precalc(const int nat,
-                           const int nks,
-                           const int nlocal,
-                           const int des_per_atom,
-                           const torch::Tensor& v_delta_precalc);
+void check_v_delta_precalc(const torch::Tensor& v_delta_precalc);
 
 // for deepks_v_delta = 2
 // prepare phialpha for outputting npy file