Combine some checking functions in DeePKS.

Author: ErjieWu

source/Makefile.Objects

@@ -201,6 +201,7 @@ OBJS_CELL=atom_pseudo.o\
 
 OBJS_DEEPKS=LCAO_deepks.o\
         deepks_basic.o\
+        deepks_check.o\
         deepks_descriptor.o\
         deepks_force.o\
         deepks_fpre.o\

source/module_hamilt_lcao/module_deepks/CMakeLists.txt

@@ -2,6 +2,7 @@ if(ENABLE_DEEPKS)
   list(APPEND objects
       LCAO_deepks.cpp
       deepks_basic.cpp
+      deepks_check.cpp
       deepks_descriptor.cpp
       deepks_force.cpp
       deepks_fpre.cpp

source/module_hamilt_lcao/module_deepks/LCAO_deepks.h

@@ -4,6 +4,7 @@
 #ifdef __DEEPKS
 
 #include "deepks_basic.h"
+#include "deepks_check.h"
 #include "deepks_descriptor.h"
 #include "deepks_force.h"
 #include "deepks_fpre.h"

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -175,8 +175,8 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
 
                 if (PARAM.inp.deepks_out_unittest)
                 {
-                    DeePKS_domain::check_gdmx(gdmx);
-                    DeePKS_domain::check_gvx(gvx, rank);
+                    DeePKS_domain::check_tensor<double>(gdmx, "gdmx.dat", rank);
+                    DeePKS_domain::check_tensor<double>(gvx, "gvx.dat", rank);
                 }
             }
         }
@@ -198,8 +198,8 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
 
                 if (PARAM.inp.deepks_out_unittest)
                 {
-                    DeePKS_domain::check_gdmepsl(gdmepsl);
-                    DeePKS_domain::check_gvepsl(gvepsl, rank);
+                    DeePKS_domain::check_tensor<double>(gdmepsl, "gdmepsl.dat", rank);
+                    DeePKS_domain::check_tensor<double>(gvepsl, "gvepsl.dat", rank);
                 }
             }
         }

source/module_hamilt_lcao/module_deepks/deepks_check.cpp

@@ -0,0 +1,55 @@
+#ifdef __DEEPKS
+
+#include "deepks_check.h"
+
+template <typename T>
+void DeePKS_domain::check_tensor(const torch::Tensor& tensor, const std::string& filename, const int rank)
+{
+    if (rank != 0)
+    {
+        return;
+    }
+    using T_tensor = typename std::conditional<std::is_same<T, std::complex<double>>::value, c10::complex<double>, T>::type;
+
+    std::ofstream ofs(filename.c_str());
+    ofs << std::setprecision(10);
+
+    auto sizes = tensor.sizes();
+    int ndim = sizes.size();
+    auto data_ptr = tensor.data_ptr<T_tensor>();
+    int64_t numel = tensor.numel();
+
+    // stride for each dimension
+    std::vector<int64_t> strides(ndim, 1);
+    for (int i = ndim - 2; i >= 0; --i) {
+        strides[i] = strides[i + 1] * sizes[i + 1];
+    }
+
+    for (int64_t idx = 0; idx < numel; ++idx) {
+        // index to multi-dimensional indices
+        std::vector<int64_t> indices(ndim);
+        int64_t tmp = idx;
+        for (int d = 0; d < ndim; ++d) {
+            indices[d] = tmp / strides[d];
+            tmp = tmp % strides[d];
+        }
+
+        T_tensor tmp_val = data_ptr[idx];
+        T* tmp_ptr = reinterpret_cast<T*>(&tmp_val);
+        ofs << *tmp_ptr;
+
+        // print space or newline
+        if ( ( (idx+1) % sizes[ndim-1] ) == 0 ) {
+            ofs << std::endl;
+        } else {
+            ofs << " ";
+        }
+    }
+
+    ofs.close();
+}
+
+template void DeePKS_domain::check_tensor<double>(const torch::Tensor& tensor, const std::string& filename, const int rank);
+template void DeePKS_domain::check_tensor<std::complex<double>>(const torch::Tensor& tensor, const std::string& filename, const int rank);
+
+#endif

source/module_hamilt_lcao/module_deepks/deepks_check.h

@@ -0,0 +1,27 @@
+#ifndef DEEPKS_CHECK_H
+#define DEEPKS_CHECK_H
+
+#ifdef __DEEPKS
+
+#include <string>
+#include <torch/script.h>
+#include <torch/torch.h>
+
+namespace DeePKS_domain
+{
+//------------------------
+// deepks_check.cpp
+//------------------------
+
+// This file contains subroutines for checking files
+
+// There are 1 subroutines in this file:
+// 1. check_tensor, which is used for tensor data checking
+
+template <typename T>
+void check_tensor(const torch::Tensor& tensor, const std::string& filename, const int rank);
+
+} // namespace DeePKS_domain
+
+#endif
+#endif

source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp

@@ -142,53 +142,6 @@ void DeePKS_domain::cal_gdmx(const int lmaxd,
     return;
 }
 
-void DeePKS_domain::check_gdmx(const torch::Tensor& gdmx)
-{
-    std::stringstream ss;
-    std::ofstream ofs_x;
-    std::ofstream ofs_y;
-    std::ofstream ofs_z;
-
-    ofs_x << std::setprecision(10);
-    ofs_y << std::setprecision(10);
-    ofs_z << std::setprecision(10);
-
-    // size: [3][natom][inlmax][nm][nm]
-    auto size = gdmx.sizes();
-    auto accessor = gdmx.accessor<double, 5>();
-    for (int ia = 0; ia < size[1]; ia++)
-    {
-        ss.str("");
-        ss << "gdmx_" << ia << ".dat";
-        ofs_x.open(ss.str().c_str());
-        ss.str("");
-        ss << "gdmy_" << ia << ".dat";
-        ofs_y.open(ss.str().c_str());
-        ss.str("");
-        ss << "gdmz_" << ia << ".dat";
-        ofs_z.open(ss.str().c_str());
-
-        for (int inl = 0; inl < size[2]; inl++)
-        {
-            for (int m1 = 0; m1 < size[3]; m1++)
-            {
-                for (int m2 = 0; m2 < size[4]; m2++)
-                {
-                    ofs_x << accessor[0][ia][inl][m1][m2] << " ";
-                    ofs_y << accessor[1][ia][inl][m1][m2] << " ";
-                    ofs_z << accessor[2][ia][inl][m1][m2] << " ";
-                }
-            }
-            ofs_x << std::endl;
-            ofs_y << std::endl;
-            ofs_z << std::endl;
-        }
-        ofs_x.close();
-        ofs_y.close();
-        ofs_z.close();
-    }
-}
-
 // calculates gradient of descriptors from gradient of projected density matrices
 void DeePKS_domain::cal_gvx(const int nat,
                             const int inlmax,
@@ -243,55 +196,6 @@ void DeePKS_domain::cal_gvx(const int nat,
     return;
 }
 
-void DeePKS_domain::check_gvx(const torch::Tensor& gvx, const int rank)
-{
-    std::stringstream ss;
-    std::ofstream ofs_x;
-    std::ofstream ofs_y;
-    std::ofstream ofs_z;
-
-    if (rank != 0)
-    {
-        return;
-    }
-
-    auto size = gvx.sizes();
-    auto accessor = gvx.accessor<double, 4>();
-
-    for (int ia = 0; ia < size[0]; ia++)
-    {
-        ss.str("");
-        ss << "gvx_" << ia << ".dat";
-        ofs_x.open(ss.str().c_str());
-        ss.str("");
-        ss << "gvy_" << ia << ".dat";
-        ofs_y.open(ss.str().c_str());
-        ss.str("");
-        ss << "gvz_" << ia << ".dat";
-        ofs_z.open(ss.str().c_str());
-
-        ofs_x << std::setprecision(10);
-        ofs_y << std::setprecision(10);
-        ofs_z << std::setprecision(10);
-
-        for (int ib = 0; ib < size[2]; ib++)
-        {
-            for (int nlm = 0; nlm < size[3]; nlm++)
-            {
-                ofs_x << accessor[ia][0][ib][nlm] << " ";
-                ofs_y << accessor[ia][1][ib][nlm] << " ";
-                ofs_z << accessor[ia][2][ib][nlm] << " ";
-            }
-            ofs_x << std::endl;
-            ofs_y << std::endl;
-            ofs_z << std::endl;
-        }
-        ofs_x.close();
-        ofs_y.close();
-        ofs_z.close();
-    }
-}
-
 template void DeePKS_domain::cal_gdmx<double>(const int lmaxd,
                                               const int inlmax,
                                               const int nks,

source/module_hamilt_lcao/module_deepks/deepks_fpre.h

@@ -21,13 +21,11 @@ namespace DeePKS_domain
 // deepks_fpre.cpp
 //------------------------
 
-// This file contains 4 subroutines for calculating,
+// This file contains 2 subroutines for calculating,
 // 1. cal_gdmx, calculating gdmx
-// 2. check_gdmx, which prints gdmx to a series of .dat files
-// 3. cal_gvx : gvx is used for training with force label, which is gradient of descriptors,
+// 2. 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
-// 4. check_gvx : prints gvx into gvx.dat for checking
 
 // calculate the gradient of pdm with regard to atomic positions
 // d/dX D_{Inl,mm'}
@@ -45,8 +43,6 @@ void cal_gdmx(const int lmaxd,
               const Grid_Driver& GridD,
               torch::Tensor& gdmx);
 
-void check_gdmx(const torch::Tensor& gdmx);
-
 /// calculates gradient of descriptors w.r.t atomic positions
 ///----------------------------------------------------
 /// m, n: 2*l+1
@@ -64,7 +60,6 @@ void cal_gvx(const int nat,
              const torch::Tensor& gdmx,
              torch::Tensor& gvx,
              const int rank);
-void check_gvx(const torch::Tensor& gvx, const int rank);
 
 } // namespace DeePKS_domain
 #endif

source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp

@@ -295,25 +295,6 @@ 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,8 +41,6 @@ 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