Restore deepks unit test in gamma-only case and make H_V_delta have full size of Hamiltonian in multi-k case

source/module_basis/module_ao/ORB_read.cpp

@@ -82,7 +82,7 @@ void LCAO_Orbitals::init(
             this->orbital_file.push_back(orbital_dir + orbital_file[it]);
         }
     }
-
+    this->descriptor_file = descriptor_file;
 #ifdef __MPI
     bcast_files(ntype, my_rank);
 #endif
@@ -263,7 +263,6 @@ void LCAO_Orbitals::Read_Orbitals(std::ofstream& ofs_in,
 
         delete[] this->Alpha;
         this->Alpha = new Numerical_Orbital[1]; // not related to atom type -- remain to be discussed
-
         this->Read_Descriptor(ofs_in, force_flag, my_rank);
     }
 

source/module_basis/module_nao/two_center_bundle.cpp

@@ -107,7 +107,7 @@ void TwoCenterBundle::tabulate()
     {
         overlap_orb_alpha = std::unique_ptr<TwoCenterIntegrator>(new TwoCenterIntegrator);
         overlap_orb_alpha->tabulate(*orb_, *alpha_, 'S', nr, cutoff);
-        ModuleBase::Memory::record("TwoCenterTable: Descriptor", overlap_orb_beta->table_memory());
+        ModuleBase::Memory::record("TwoCenterTable: Descriptor", overlap_orb_alpha->table_memory());
     }
 
     if (orb_onsite_)

source/module_cell/unitcell.cpp

@@ -719,6 +719,7 @@ void UnitCell::cal_nwfc(std::ofstream& log) {
     //========================
     this->lmax = 0;
     this->nmax = 0;
+    this->nmax_total = 0;
     for (int it = 0; it < ntype; it++) {
         lmax = std::max(lmax, atoms[it].nwl);
         for (int l = 0; l < atoms[it].nwl + 1; l++) {

source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp

@@ -55,10 +55,10 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::initialize_HR(const Grid_Driv
     auto* paraV = this->hR->get_paraV(); // get parallel orbitals from HR
     // TODO: if paraV is nullptr, AtomPair can not use paraV for constructor, I will repair it in the future.
 
-    // this->H_V_delta = new HContainer<TR>(paraV);
+    this->H_V_delta = new HContainer<TR>(paraV);
     if (std::is_same<TK, double>::value)
     {
-        this->H_V_delta = new HContainer<TR>(paraV);
+        //this->H_V_delta = new HContainer<TR>(paraV);
         this->H_V_delta->fix_gamma();
     }
 
@@ -123,10 +123,10 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::initialize_HR(const Grid_Driv
                                          R_index2.y - R_index1.y,
                                          R_index2.z - R_index1.z,
                                          paraV);
-                if (std::is_same<TK, double>::value)
-                {
-                    this->H_V_delta->insert_pair(tmp);
-                }
+                // if (std::is_same<TK, double>::value)
+                // {
+                this->H_V_delta->insert_pair(tmp);
+                // }
             }
         }
         if (pre_cal_nlm)
@@ -135,15 +135,14 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::initialize_HR(const Grid_Driv
         }
     }
     // allocate the memory of BaseMatrix in HR, and set the new values to zero
-    if (std::is_same<TK, double>::value)
-    {
-        // only gamma-only has full size of Hamiltonian of DeePKS now,
-        // multi-k keep same size of nonlocal operator, H_V_delta will be allocated by hR
-        this->H_V_delta->allocate(nullptr, true);
-        // expand hR with H_V_delta, only gamma-only case now
-        this->hR->add(*this->H_V_delta);
-        this->hR->allocate(nullptr, false);
-    }
+    // if (std::is_same<TK, double>::value)
+    // {
+    this->H_V_delta->allocate(nullptr, true);
+        // expand hR with H_V_delta
+        // update : for computational rigor, gamma-only and multi-k cases both have full size of Hamiltonian of DeePKS now
+    this->hR->add(*this->H_V_delta);
+    this->hR->allocate(nullptr, false);
+    // }
 
     ModuleBase::timer::tick("DeePKS", "initialize_HR");
 }
@@ -165,11 +164,11 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::contributeHR()
         GlobalC::ld.cal_descriptor(this->ucell->nat);
         GlobalC::ld.cal_gedm(this->ucell->nat);
 
-        // recalculate the H_V_delta
-        if (this->H_V_delta == nullptr)
-        {
-            this->H_V_delta = new hamilt::HContainer<TR>(*this->hR);
-        }
+        // // recalculate the H_V_delta
+        // if (this->H_V_delta == nullptr)
+        // {
+        //     this->H_V_delta = new hamilt::HContainer<std::complex<double>>(*this->hR);
+        // }
         this->H_V_delta->set_zero();
         this->calculate_HR();
 

source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.h

@@ -61,6 +61,7 @@ class DeePKS<OperatorLCAO<TK, TR>> : public OperatorLCAO<TK, TR>
     // LCAO_Deepks* ld = nullptr;
 
     const UnitCell* ucell = nullptr;
+    Grid_Driver* gridD = nullptr;
 
     const Grid_Driver* gd = nullptr;
 

source/module_hamilt_lcao/module_deepks/CMakeLists.txt

@@ -30,8 +30,8 @@ if(ENABLE_DEEPKS)
   endif()
 
 # I will rewrite the test later, the current test rely on too many modules
-#  if(BUILD_TESTING)
-#      add_subdirectory(test)
-#  endif()
+  if(BUILD_TESTING)
+      add_subdirectory(test)
+  endif()
 endif()
 

source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp

@@ -361,7 +361,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix<TK, double>* d
 
 void LCAO_Deepks::check_projected_dm()
 {
-    const std::string file_projdm = PARAM.globalv.global_out_dir + "deepks_projdm.dat";
+    const std::string file_projdm = PARAM.globalv.global_out_dir + "pdm.dat";
     std::ofstream ofs(file_projdm.c_str());
 
     const int pdm_size = (this->lmaxd * 2 + 1) * (this->lmaxd * 2 + 1);

source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp

@@ -45,9 +45,6 @@ void LCAO_Deepks::cal_descriptor(const int nat) {
         return;
     }
 
-    // init pdm_tensor and d_tensor
-    torch::Tensor tmp;
-
     // if pdm_tensor and d_tensor is not empty, clear it !!
     if (!this->d_tensor.empty()) 
     {
@@ -63,23 +60,15 @@ void LCAO_Deepks::cal_descriptor(const int nat) {
     for (int inl = 0; inl < this->inlmax; ++inl) 
     {
         const int nm = 2 * inl_l[inl] + 1;
-        tmp = torch::ones({nm, nm},
-                          torch::TensorOptions().dtype(torch::kFloat64));
 
-        for (int m1 = 0; m1 < nm; ++m1) 
-        {
-            for (int m2 = 0; m2 < nm; ++m2) 
-            {
-                tmp.index_put_({m1, m2}, this->pdm[inl][m1 * nm + m2]);
-            }
-        }
+        torch::TensorOptions options = torch::TensorOptions()
+            .dtype(torch::kFloat64)
+            .requires_grad(true); 
 
-        // torch::Tensor tmp = torch::from_blob(this->pdm[inl], { nm, nm },
-        // torch::requires_grad());
+        torch::Tensor tmp = torch::from_blob(this->pdm[inl], { nm, nm }, options).clone();
 
-        tmp.requires_grad_(true);
         this->pdm_tensor.push_back(tmp);
-        this->d_tensor.push_back(torch::ones({nm}, torch::requires_grad(true)));
+        this->d_tensor.push_back(torch::ones({nm}, options));
     }
 
     // cal d_tensor

source/module_hamilt_lcao/module_deepks/test/CMakeLists.txt

@@ -1,44 +1,46 @@
 add_executable(
     test_deepks
-    main_deepks.cpp klist_1.cpp LCAO_deepks_test_prep.cpp LCAO_deepks_test.cpp nnr.cpp
+    main_deepks.cpp klist_1.cpp LCAO_deepks_test_prep.cpp LCAO_deepks_test.cpp
+    ../../../module_cell/unitcell.cpp
+    ../../../module_cell/atom_spec.cpp
+    ../../../module_cell/atom_pseudo.cpp
+    ../../../module_cell/read_atoms.cpp
+    ../../../module_cell/setup_nonlocal.cpp
+    ../../../module_cell/pseudo.cpp
+    ../../../module_cell/read_pp.cpp
+    ../../../module_cell/read_pp_complete.cpp
+    ../../../module_cell/read_pp_upf100.cpp
+    ../../../module_cell/read_pp_upf201.cpp
+    ../../../module_cell/read_pp_vwr.cpp
+    ../../../module_cell/read_pp_blps.cpp
+    ../../../module_hamilt_pw/hamilt_pwdft/soc.cpp
+    ../../../module_io/output.cpp
+    ../../../module_elecstate/read_pseudo.cpp
+    ../../../module_elecstate/cal_nelec_nband.cpp
+    ../../../module_elecstate/module_dm/density_matrix.cpp
+    ../../../module_elecstate/module_dm/density_matrix_io.cpp
+    ../../../module_hamilt_lcao/module_hcontainer/base_matrix.cpp 
+    ../../../module_hamilt_lcao/module_hcontainer/hcontainer.cpp 
+    ../../../module_hamilt_lcao/module_hcontainer/atom_pair.cpp
+    ../../../module_hamilt_lcao/module_hcontainer/func_transfer.cpp
+    ../../../module_hamilt_lcao/module_hcontainer/func_folding.cpp
+    ../../../module_hamilt_lcao/module_hcontainer/transfer.cpp
+    ../../../module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp
+    ../../../module_hamilt_lcao/hamilt_lcaodft/operator_lcao/operator_lcao.cpp
+    ../../../module_hamilt_general/operator.cpp
 )
 
-if(ENABLE_COVERAGE)
-  add_coverage(test_deepks)
-endif()
-
-get_target_property(ABACUS_LINK_LIBRARIES ${ABACUS_BIN_NAME} LINK_LIBRARIES)
 target_link_libraries(
     test_deepks
-    base cell symmetry md surchem xc_ 
-    neighbor orb io_basic io_advanced relax gint driver esolver hsolver psi elecstate
-    hamilt_general hamilt_pwdft hamilt_lcao tddft hamilt_ofdft hamilt_stodft planewave
-    pthread vdw dftu deltaspin hcontainer
-    deepks device numerical_atomic_orbitals container psi_initializer psi_overall_init
-    ${ABACUS_LINK_LIBRARIES}
+    base device  parameter deepks psi planewave neighbor container
+    orb gint numerical_atomic_orbitals paw
+    ${math_libs}
 )
 
-if(ENABLE_PAW)
-  target_link_libraries(
-    test_deepks
-    paw
-  )
-endif()
-
-if(USE_ELPA)
-  target_link_libraries(
-      test_deepks
-      genelpa
-  )
-endif()
-if(USE_CUDA)
-  target_link_libraries(diag_cusolver)
+if(ENABLE_COVERAGE)
+  add_coverage(test_deepks)
 endif()
 
-if (ENABLE_LIBRI)
-  target_link_libraries(test_deepks
-          ri)
-endif()
 
 install(
     TARGETS test_deepks