Move functions for calculating descriptor from LCAO_deepks to DeePKS_domain.

ErjieWu · ErjieWu · commit 865f3dadb4c5 · 2025-01-04T22:44:44.000+08:00
diff --git a/source/Makefile.Objects b/source/Makefile.Objects
@@ -192,6 +192,7 @@ OBJS_CELL=atom_pseudo.o\
 
 OBJS_DEEPKS=LCAO_deepks.o\
         deepks_force.o\
+        deepks_descriptor.o\
         deepks_orbital.o\
         deepks_orbpre.o\
         deepks_vdpre.o\
@@ -206,7 +207,6 @@ OBJS_DEEPKS=LCAO_deepks.o\
         cal_gdmepsl.o\
         cal_gedm.o\
         cal_gvx.o\
-        cal_descriptor.o\
         
 
 OBJS_ELECSTAT=elecstate.o\
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp
@@ -254,7 +254,12 @@ void Force_LCAO<double>::ftable(const bool isforce,
         // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm
         // GlobalC::ld.cal_projected_DM(dm, ucell, orb, gd);
 
-        GlobalC::ld.cal_descriptor(ucell.nat, descriptor);
+        DeePKS_domain::cal_descriptor(ucell.nat,
+                                      GlobalC::ld.inlmax,
+                                      GlobalC::ld.inl_l,
+                                      GlobalC::ld.pdm,
+                                      descriptor,
+                                      GlobalC::ld.des_per_atom);
         GlobalC::ld.cal_gedm(ucell.nat, descriptor);
 
         const int nks = 1;
@@ -306,7 +311,12 @@ void Force_LCAO<double>::ftable(const bool isforce,
 
         GlobalC::ld.check_projected_dm();
 
-        GlobalC::ld.check_descriptor(ucell, PARAM.globalv.global_out_dir, descriptor);
+        DeePKS_domain::check_descriptor(GlobalC::ld.inlmax,
+                                        GlobalC::ld.des_per_atom,
+                                        GlobalC::ld.inl_l,
+                                        ucell,
+                                        PARAM.globalv.global_out_dir,
+                                        descriptor);
 
         GlobalC::ld.check_gedm();
 
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp
@@ -350,7 +350,12 @@ void Force_LCAO<std::complex<double>>::ftable(const bool isforce,
         // GlobalC::ld.cal_projected_DM(dm, ucell, orb, gd);
 
         std::vector<torch::Tensor> descriptor;
-        GlobalC::ld.cal_descriptor(ucell.nat, descriptor);
+        DeePKS_domain::cal_descriptor(ucell.nat,
+                                      GlobalC::ld.inlmax,
+                                      GlobalC::ld.inl_l,
+                                      GlobalC::ld.pdm,
+                                      descriptor,
+                                      GlobalC::ld.des_per_atom);
         GlobalC::ld.cal_gedm(ucell.nat, descriptor);
 
         DeePKS_domain::cal_f_delta<std::complex<double>>(dm_k,
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp
@@ -163,7 +163,12 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::contributeHR()
         GlobalC::ld.cal_projected_DM<TK>(this->DM, *this->ucell, *ptr_orb_, *(this->gd));
 
         std::vector<torch::Tensor> descriptor;
-        GlobalC::ld.cal_descriptor(this->ucell->nat, descriptor);
+        DeePKS_domain::cal_descriptor(this->ucell->nat,
+                                      GlobalC::ld.inlmax,
+                                      GlobalC::ld.inl_l,
+                                      GlobalC::ld.pdm,
+                                      descriptor,
+                                      GlobalC::ld.des_per_atom);
         GlobalC::ld.cal_gedm(this->ucell->nat, descriptor);
 
         // // recalculate the H_V_delta
diff --git a/source/module_hamilt_lcao/module_deepks/CMakeLists.txt b/source/module_hamilt_lcao/module_deepks/CMakeLists.txt
@@ -1,6 +1,7 @@
 if(ENABLE_DEEPKS)
   list(APPEND objects
       LCAO_deepks.cpp
+      deepks_descriptor.cpp
       deepks_force.cpp
       deepks_orbital.cpp
       deepks_orbpre.cpp
@@ -16,8 +17,6 @@ if(ENABLE_DEEPKS)
       cal_gdmepsl.cpp
       cal_gedm.cpp
       cal_gvx.cpp      
-      cal_descriptor.cpp
-      
   )
 
   add_library(
diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h
@@ -3,6 +3,7 @@
 
 #ifdef __DEEPKS
 
+#include "deepks_descriptor.h"
 #include "deepks_force.h"
 #include "deepks_hmat.h"
 #include "deepks_orbital.h"
@@ -309,9 +310,6 @@ class LCAO_Deepks
     // as well as subroutines that prints the results for checking
 
     // The file contains 8 subroutines:
-    // 1. cal_descriptor : obtains descriptors which are eigenvalues of pdm
-    //       by calling torch::linalg::eigh
-    // 2. check_descriptor : prints descriptor for checking
     // 3. cal_gvx : gvx is used for training with force label, which is gradient of descriptors,
     //       calculated by d(des)/dX = d(pdm)/dX * d(des)/d(pdm) = gdmx * gvdm
     //       using einsum
@@ -329,16 +327,6 @@ class LCAO_Deepks
     // 9. check_gedm : prints gedm for checking
 
   public:
-    /// Calculates descriptors
-    /// which are eigenvalues of pdm in blocks of I_n_l
-    void cal_descriptor(const int nat, std::vector<torch::Tensor>& descriptor);
-    /// print descriptors based on LCAO basis
-    void check_descriptor(const UnitCell& ucell,
-                          const std::string& out_dir,
-                          const std::vector<torch::Tensor>& descriptor);
-
-    void cal_descriptor_equiv(const int nat, std::vector<torch::Tensor>& descriptor);
-
     /// calculates gradient of descriptors w.r.t atomic positions
     ///----------------------------------------------------
     /// m, n: 2*l+1
diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp
@@ -298,8 +298,18 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
         ld->check_projected_dm(); // print out the projected dm for NSCF calculaiton
 
         std::vector<torch::Tensor> descriptor;
-        ld->cal_descriptor(nat, descriptor); // final descriptor
-        ld->check_descriptor(ucell, PARAM.globalv.global_out_dir, descriptor);
+        DeePKS_domain::cal_descriptor(nat,
+                                      ld->inlmax,
+                                      ld->inl_l,
+                                      ld->pdm,
+                                      descriptor,
+                                      ld->des_per_atom); // final descriptor
+        DeePKS_domain::check_descriptor(ld->inlmax,
+                                        ld->des_per_atom,
+                                        ld->inl_l,
+                                        ucell,
+                                        PARAM.globalv.global_out_dir,
+                                        descriptor);
 
         if (PARAM.inp.deepks_out_labels)
         {
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_descriptor.cpp b/source/module_hamilt_lcao/module_deepks/deepks_descriptor.cpp
@@ -1,10 +1,12 @@
 /// 1. cal_descriptor : obtains descriptors which are eigenvalues of pdm
 ///       by calling torch::linalg::eigh
 /// 2. check_descriptor : prints descriptor for checking
+/// 3. cal_descriptor_equiv : calculates descriptor in equivalent version
 
 #ifdef __DEEPKS
 
-#include "LCAO_deepks.h"
+#include "deepks_descriptor.h"
+
 #include "LCAO_deepks_io.h" // mohan add 2024-07-22
 #include "module_base/blas_connector.h"
 #include "module_base/constants.h"
@@ -13,44 +15,46 @@
 #include "module_hamilt_lcao/module_hcontainer/atom_pair.h"
 #include "module_parameter/parameter.h"
 
-void LCAO_Deepks::cal_descriptor_equiv(const int nat, std::vector<torch::Tensor>& descriptor)
+void DeePKS_domain::cal_descriptor_equiv(const int nat,
+                                         const int des_per_atom,
+                                         const std::vector<torch::Tensor>& pdm,
+                                         std::vector<torch::Tensor>& descriptor)
 {
-    ModuleBase::TITLE("LCAO_Deepks", "cal_descriptor_equiv");
-    ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor_equiv");
+    ModuleBase::TITLE("DeePKS_domain", "cal_descriptor_equiv");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_descriptor_equiv");
 
+    assert(des_per_atom > 0);
     for (int iat = 0; iat < nat; iat++)
     {
         auto tmp = torch::zeros(des_per_atom, torch::kFloat64);
-        std::memcpy(tmp.data_ptr(), this->pdm[iat].data_ptr<double>(), sizeof(double) * tmp.numel());
+        std::memcpy(tmp.data_ptr(), pdm[iat].data_ptr<double>(), sizeof(double) * tmp.numel());
         descriptor.push_back(tmp);
     }
 
-    ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor_equiv");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_descriptor_equiv");
 }
 
 // calculates descriptors from projected density matrices
-void LCAO_Deepks::cal_descriptor(const int nat, std::vector<torch::Tensor>& descriptor)
+void DeePKS_domain::cal_descriptor(const int nat,
+                                   const int inlmax,
+                                   const int* inl_l,
+                                   const std::vector<torch::Tensor>& pdm,
+                                   std::vector<torch::Tensor>& descriptor,
+                                   const int des_per_atom = -1)
 {
-    ModuleBase::TITLE("LCAO_Deepks", "cal_descriptor");
-    ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor");
-
-    // init descriptor
-    // if descriptor is not empty, clear it !!
-    if (!descriptor.empty())
-    {
-        descriptor.erase(descriptor.begin(), descriptor.end());
-    }
+    ModuleBase::TITLE("DeePKS_domain", "cal_descriptor");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_descriptor");
 
     if (PARAM.inp.deepks_equiv)
     {
-        this->cal_descriptor_equiv(nat, descriptor);
+        DeePKS_domain::cal_descriptor_equiv(nat, des_per_atom, pdm, descriptor);
         return;
     }
 
-    for (int inl = 0; inl < this->inlmax; ++inl)
+    for (int inl = 0; inl < inlmax; ++inl)
     {
         const int nm = 2 * inl_l[inl] + 1;
-        this->pdm[inl].requires_grad_(true);
+        pdm[inl].requires_grad_(true);
         descriptor.push_back(torch::ones({nm}, torch::requires_grad(true)));
     }
 
@@ -64,15 +68,18 @@ void LCAO_Deepks::cal_descriptor(const int nat, std::vector<torch::Tensor>& desc
         d_v = torch::linalg::eigh(pdm[inl], /*uplo*/ "U");
         descriptor[inl] = std::get<0>(d_v);
     }
-    ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_descriptor");
     return;
 }
 
-void LCAO_Deepks::check_descriptor(const UnitCell& ucell,
-                                   const std::string& out_dir,
-                                   const std::vector<torch::Tensor>& descriptor)
+void DeePKS_domain::check_descriptor(const int inlmax,
+                                     const int des_per_atom,
+                                     const int* inl_l,
+                                     const UnitCell& ucell,
+                                     const std::string& out_dir,
+                                     const std::vector<torch::Tensor>& descriptor)
 {
-    ModuleBase::TITLE("LCAO_Deepks", "check_descriptor");
+    ModuleBase::TITLE("DeePKS_domain", "check_descriptor");
 
     if (GlobalV::MY_RANK != 0)
     {
@@ -91,7 +98,7 @@ void LCAO_Deepks::check_descriptor(const UnitCell& ucell,
             for (int ia = 0; ia < ucell.atoms[it].na; ia++)
             {
                 int iat = ucell.itia2iat(it, ia);
-                ofs << ucell.atoms[it].label << " atom_index " << ia + 1 << " n_descriptor " << this->des_per_atom
+                ofs << ucell.atoms[it].label << " atom_index " << ia + 1 << " n_descriptor " << des_per_atom
                     << std::endl;
                 int id = 0;
                 for (int inl = 0; inl < inlmax / ucell.nat; inl++)
@@ -118,10 +125,9 @@ void LCAO_Deepks::check_descriptor(const UnitCell& ucell,
         for (int iat = 0; iat < ucell.nat; iat++)
         {
             const int it = ucell.iat2it[iat];
-            ofs << ucell.atoms[it].label << " atom_index " << iat + 1 << " n_descriptor " << this->des_per_atom
-                << std::endl;
+            ofs << ucell.atoms[it].label << " atom_index " << iat + 1 << " n_descriptor " << des_per_atom << std::endl;
             auto accessor = descriptor[iat].accessor<double, 1>();
-            for (int i = 0; i < this->des_per_atom; i++)
+            for (int i = 0; i < des_per_atom; i++)
             {
                 ofs << accessor[i] << " ";
                 if (i % 8 == 7)
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_descriptor.h b/source/module_hamilt_lcao/module_deepks/deepks_descriptor.h
@@ -0,0 +1,51 @@
+#ifndef DEEPKS_DESCRIPTOR_H
+#define DEEPKS_DESCRIPTOR_H
+
+#ifdef __DEEPKS
+
+#include "module_base/intarray.h"
+#include "module_base/timer.h"
+#include "module_cell/unitcell.h"
+
+#include <torch/script.h>
+#include <torch/torch.h>
+
+namespace DeePKS_domain
+{
+//------------------------
+// deepks_descriptor.cpp
+//------------------------
+
+// This file contains interfaces with libtorch,
+// including loading of model and calculating gradients
+// as well as subroutines that prints the results for checking
+
+// The file contains 8 subroutines:
+// 1. cal_descriptor : obtains descriptors which are eigenvalues of pdm
+//       by calling torch::linalg::eigh
+// 2. check_descriptor : prints descriptor for checking
+// 3. cal_descriptor_equiv : calculates descriptor in equivalent version
+
+/// Calculates descriptors
+/// which are eigenvalues of pdm in blocks of I_n_l
+void cal_descriptor(const int nat,
+                    const int inlmax,
+                    const int* inl_l,
+                    const std::vector<torch::Tensor>& pdm,
+                    std::vector<torch::Tensor>& descriptor,
+                    const int des_per_atom);
+/// print descriptors based on LCAO basis
+void check_descriptor(const int inlmax,
+                      const int des_per_atom,
+                      const int* inl_l,
+                      const UnitCell& ucell,
+                      const std::string& out_dir,
+                      const std::vector<torch::Tensor>& descriptor);
+
+void cal_descriptor_equiv(const int nat,
+                          const int des_per_atom,
+                          const std::vector<torch::Tensor>& pdm,
+                          std::vector<torch::Tensor>& descriptor);
+} // namespace DeePKS_domain
+#endif
+#endif
diff --git a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp
@@ -198,8 +198,18 @@ void test_deepks::check_gdmx(torch::Tensor& gdmx)
 
 void test_deepks::check_descriptor(std::vector<torch::Tensor>& descriptor)
 {
-    GlobalC::ld.cal_descriptor(ucell.nat, descriptor);
-    GlobalC::ld.check_descriptor(ucell, "./", descriptor);
+    DeePKS_domain::cal_descriptor(ucell.nat,
+                                  GlobalC::ld.inlmax,
+                                  GlobalC::ld.inl_l,
+                                  GlobalC::ld.pdm,
+                                  descriptor,
+                                  GlobalC::ld.des_per_atom);
+    DeePKS_domain::check_descriptor(GlobalC::ld.inlmax,
+                                    GlobalC::ld.des_per_atom,
+                                    GlobalC::ld.inl_l,
+                                    ucell,
+                                    "./",
+                                    descriptor);
     this->compare_with_ref("deepks_desc.dat", "descriptor_ref.dat");
 }
 
