Refactor: Combine some checking functions in DeePKS. (#6285)

* Combine some checking functions in DeePKS.

* Remove check_f_delta().

* Remove check_o_delta().

* Update vdpre_ref.dat

* Fix a bug.

* Fix a merge bug.

* Reduce the memory cost for DeePKS GO-UT.

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/hamilt_lcaodft/FORCE_gamma.cpp

@@ -272,7 +272,14 @@ void Force_LCAO<double>::ftable(const bool isforce,
 #ifdef __MLALGO
     if (PARAM.inp.deepks_scf && PARAM.inp.deepks_out_unittest)
     {
-        DeePKS_domain::check_f_delta(ucell.nat, fvnl_dalpha, svnl_dalpha);
+        std::ofstream ofs_f("F_delta.dat");
+        std::ofstream ofs_s("stress_delta.dat");
+        ofs_f << std::setprecision(10);
+        ofs_s << std::setprecision(10);
+        fvnl_dalpha.print(ofs_f);
+        ofs_f.close();
+        svnl_dalpha.print(ofs_s);
+        ofs_s.close();
     }
 #endif
 

source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp

@@ -312,7 +312,14 @@ void Force_LCAO<std::complex<double>>::ftable(const bool isforce,
 #ifdef __MLALGO
     if (PARAM.inp.deepks_scf && PARAM.inp.deepks_out_unittest)
     {
-        DeePKS_domain::check_f_delta(ucell.nat, fvnl_dalpha, svnl_dalpha);
+        std::ofstream ofs_f("F_delta.dat");
+        std::ofstream ofs_s("stress_delta.dat");
+        ofs_f << std::setprecision(10);
+        ofs_s << std::setprecision(10);
+        fvnl_dalpha.print(ofs_f);
+        ofs_f.close();
+        svnl_dalpha.print(ofs_s);
+        ofs_s.close();
     }
 #endif
 

source/module_hamilt_lcao/module_deepks/CMakeLists.txt

@@ -2,6 +2,7 @@ if(ENABLE_MLALGO)
   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 __MLALGO
 
 #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 __MLALGO
+
+#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 __MLALGO
+
+#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_force.cpp

@@ -263,32 +263,6 @@ void DeePKS_domain::cal_f_delta(const hamilt::HContainer<double>* dmr,
     return;
 }
 
-// prints forces and stress from DeePKS (LCAO)
-void DeePKS_domain::check_f_delta(const int nat, ModuleBase::matrix& f_delta, ModuleBase::matrix& svnl_dalpha)
-{
-    ModuleBase::TITLE("DeePKS_domain", "check_F_delta");
-
-    std::ofstream ofs("F_delta.dat");
-    ofs << std::setprecision(10);
-
-    for (int iat = 0; iat < nat; iat++)
-    {
-        ofs << f_delta(iat, 0) << " " << f_delta(iat, 1) << " " << f_delta(iat, 2) << std::endl;
-    }
-
-    std::ofstream ofs1("stress_delta.dat");
-    ofs1 << std::setprecision(10);
-    for (int ipol = 0; ipol < 3; ipol++)
-    {
-        for (int jpol = 0; jpol < 3; jpol++)
-        {
-            ofs1 << svnl_dalpha(ipol, jpol) << " ";
-        }
-        ofs1 << std::endl;
-    }
-    return;
-}
-
 template void DeePKS_domain::cal_f_delta<double>(const hamilt::HContainer<double>* dmr,
                                                  const UnitCell& ucell,
                                                  const LCAO_Orbitals& orb,

source/module_hamilt_lcao/module_deepks/deepks_force.h

@@ -21,9 +21,8 @@ namespace DeePKS_domain
 // This file contains subroutines for calculating F_delta,
 // which is defind as sum_mu,nu rho_mu,nu d/dX (<chi_mu|alpha>V(D)<alpha|chi_nu>)
 
-// There are 2 subroutines in this file:
+// There are 1 subroutine in this file:
 // 1. cal_f_delta, which is used for F_delta calculation
-// 2. check_f_delta, which prints F_delta into F_delta.dat for checking
 
 template <typename TK>
 void cal_f_delta(const hamilt::HContainer<double>* dmr,
@@ -39,8 +38,6 @@ void cal_f_delta(const hamilt::HContainer<double>* dmr,
                  ModuleBase::matrix& f_delta,
                  const bool isstress,
                  ModuleBase::matrix& svnl_dalpha);
-
-void check_f_delta(const int nat, ModuleBase::matrix& f_delta, ModuleBase::matrix& svnl_dalpha);
 } // namespace DeePKS_domain
 
 #endif