Feature&Doc&Tests: Support for deepks_v_delta=-1&-2 (#6245)

* Move DeePKS printing message from screen to log file.

* Add support for DeePKS Real space Hamiltonian.

* Fix a merge bug.

* Add DeePKS V_delta_R support in doc&value_check.

* Add integrate test for deepks_v_delta<0.

* Remove redundant code.

* Fix the bug that DeePKS integrate test output 0/1 file.

* Add UT for DeePKS H(R).

* clang-format change.

---------

Co-authored-by: Mohan Chen <[email protected]>

Author: ErjieWu

docs/advanced/input_files/input-main.md

@@ -2145,7 +2145,7 @@ Warning: this function is not robust enough for the current version. Please try
 - **Description**: Print labels and descriptors for DeePKS in OUT.${suffix}. The names of these files start with "deepks".
   - 0 : No output.
   - 1 : Output intermediate files needed during DeePKS training.
-  - 2 : Output target labels for label preperation. The label files are named as `deepks_<property>.npy`, where the units and formats are the same as label files `<property>.npy` required for training, except that the first dimension (`nframes`) is excluded. System structrue files are also given in `deepks_atom.npy` and `deepks_box.npy` in the unit of *Bohr*, which means `lattice_constant` should be set to 1 when training. 
+  - 2 : Output target labels for label preperation. The label files are named as `deepks_<property>.npy` or `deepks_<property>.csr`, where the units and formats are the same as label files `<property>.npy` or `<property>.csr` required for training, except that the first dimension (`nframes`) is excluded. System structrue files are also given in `deepks_atom.npy` and `deepks_box.npy` in the unit of *Bohr*, which means `lattice_constant` should be set to 1 when training. 
 - **Note**: When `deepks_out_labels` equals **1**, the path of a numerical descriptor (an `orb` file) is needed to be specified under the `NUMERICAL_DESCRIPTOR` tag in the `STRU` file. For example:
 
   ```text
@@ -2252,8 +2252,11 @@ Warning: this function is not robust enough for the current version. Please try
 
 - **Type**: int
 - **Availability**: numerical atomic orbital basis
-- **Description**: Include V_delta label for DeePKS training. When `deepks_out_labels` is true and `deepks_v_delta` > 0, ABACUS will output h_base.npy, v_delta.npy and h_tot.npy(h_tot=h_base+v_delta). 
-  Meanwhile, when `deepks_v_delta` equals 1, ABACUS will also output v_delta_precalc.npy, which is used to calculate V_delta during DeePKS training. However, when the number of atoms grows, the size of v_delta_precalc.npy will be very large. In this case, it's recommended to set `deepks_v_delta` as 2, and ABACUS will output phialpha.npy and grad_evdm.npy but not v_delta_precalc.npy. These two files are small and can be used to calculate v_delta_precalc in the procedure of training DeePKS.
+- **Description**: Include V_delta/V_delta_R (Hamiltonian in k/real space) label for DeePKS training. When `deepks_out_labels` is true and `deepks_v_delta` > 0 (k space), ABACUS will output `deepks_hbase.npy`, `deepks_vdelta.npy` and `deepks_htot.npy`(htot=hbase+vdelta). When `deepks_out_labels` is true and `deepks_v_delta` < 0 (real space), ABACUS will output `deepks_hrtot.csr`, `deepks_hrdelta.csr`. Some more files output for different settings.
+  - `deepks_v_delta` = 1: `deepks_vdpre.npy`, which is used to calculate V_delta during DeePKS training.
+  - `deepks_v_delta` = 2: `deepks_phialpha.npy` and `deepks_gevdm.npy`, which can be used to calculate `deepks_vdpre.npy`. A recommanded method for memory saving.
+  - `deepks_v_delta` = -1: `deepks_vdrpre.npy`, which is used to calculate V_delta_R during DeePKS training.
+  - `deepks_v_delta` = -2: `deepks_phialpha_r.npy` and `deepks_gevdm.npy`, which can be used to calculate `deepks_vdrpre.npy`. A recommanded method for memory saving.
 - **Default**: 0
 
 ### deepks_out_unittest

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -2,12 +2,12 @@
 #include "LCAO_deepks_interface.h"
 
 #include "LCAO_deepks_io.h" // mohan add 2024-07-22
-#include "source_base/global_variable.h"
-#include "source_base/tool_title.h"
 #include "module_elecstate/cal_dm.h"
 #include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
 #include "module_hamilt_lcao/module_hcontainer/output_hcontainer.h"
 #include "module_parameter/parameter.h"
+#include "source_base/global_variable.h"
+#include "source_base/tool_title.h"
 
 template <typename TK, typename TR>
 LCAO_Deepks_Interface<TK, TR>::LCAO_Deepks_Interface(std::shared_ptr<LCAO_Deepks<TK>> ld_in) : ld(ld_in)
@@ -353,14 +353,15 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
             }                                                                 // end deepks_out_labels == 1
         }
 
-        // not add deepks_out_labels = 2 for HR yet
         // H(R) matrix part, for HR, base will not be calculated since they are HContainer objects
         if (PARAM.inp.deepks_v_delta < 0)
         {
             // set the output
             const double sparse_threshold = 1e-10;
             const int precision = 8;
-            const std::string file_hrtot = PARAM.globalv.global_out_dir + "deepks_hrtot.csr";
+            const std::string file_hrtot
+                = PARAM.globalv.global_out_dir
+                  + (PARAM.inp.deepks_out_labels == 1 ? "deepks_hrtot.csr" : "deepks_hamiltonian_r.csr");
             hamilt::HContainer<TR>* hR_tot = (p_ham->getHR());
 
             if (rank == 0)
@@ -369,47 +370,75 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
                 ofs_hr << "Matrix Dimension of H(R): " << hR_tot->get_nbasis() << std::endl;
                 ofs_hr << "Matrix number of H(R): " << hR_tot->size_R_loop() << std::endl;
                 hamilt::Output_HContainer<TR> out_hr(hR_tot, ofs_hr, sparse_threshold, precision);
-                out_hr.write();
+                out_hr.write(true); // write all the matrices, including empty ones
                 ofs_hr.close();
             }
 
             if (PARAM.inp.deepks_scf)
             {
-                const std::string file_vdeltar = PARAM.globalv.global_out_dir + "deepks_hrdelta.csr";
-                hamilt::HContainer<TR>* h_deltaR = p_ham->get_V_delta_R();
-
-                if (rank == 0)
+                if (PARAM.inp.deepks_out_labels == 1)
                 {
-                    std::ofstream ofs_hr(file_vdeltar, std::ios::out);
-                    ofs_hr << "Matrix Dimension of H_delta(R): " << h_deltaR->get_nbasis() << std::endl;
-                    ofs_hr << "Matrix number of H_delta(R): " << h_deltaR->size_R_loop() << std::endl;
-                    hamilt::Output_HContainer<TR> out_hr(h_deltaR, ofs_hr, sparse_threshold, precision);
-                    out_hr.write();
-                    ofs_hr.close();
-                }
+                    const std::string file_vdeltar = PARAM.globalv.global_out_dir + "deepks_hrdelta.csr";
+                    hamilt::HContainer<TR>* h_deltaR = p_ham->get_V_delta_R();
 
-                torch::Tensor phialpha_r_out;
-                torch::Tensor R_query;
-                DeePKS_domain::prepare_phialpha_r(nlocal,
-                                                  lmaxd,
-                                                  inlmax,
-                                                  nat,
-                                                  phialpha,
-                                                  ucell,
-                                                  orb,
-                                                  *ParaV,
-                                                  GridD,
-                                                  phialpha_r_out,
-                                                  R_query);
-                const std::string file_phialpha_r = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy";
-                const std::string file_R_query = PARAM.globalv.global_out_dir + "deepks_R_query.npy";
-                LCAO_deepks_io::save_tensor2npy<double>(file_phialpha_r, phialpha_r_out, rank);
-                LCAO_deepks_io::save_tensor2npy<int>(file_R_query, R_query, rank);
-
-                torch::Tensor gevdm_out;
-                DeePKS_domain::prepare_gevdm(nat, lmaxd, inlmax, orb, gevdm, gevdm_out);
-                const std::string file_gevdm = PARAM.globalv.global_out_dir + "deepks_gevdm.npy";
-                LCAO_deepks_io::save_tensor2npy<double>(file_gevdm, gevdm_out, rank);
+                    if (rank == 0)
+                    {
+                        std::ofstream ofs_hr(file_vdeltar, std::ios::out);
+                        ofs_hr << "Matrix Dimension of H_delta(R): " << h_deltaR->get_nbasis() << std::endl;
+                        ofs_hr << "Matrix number of H_delta(R): " << h_deltaR->size_R_loop() << std::endl;
+                        hamilt::Output_HContainer<TR> out_hr(h_deltaR, ofs_hr, sparse_threshold, precision);
+                        out_hr.write(true); // write all the matrices, including empty ones
+                        ofs_hr.close();
+                    }
+
+                    if (PARAM.inp.deepks_v_delta == -1)
+                    {
+                        int R_size = DeePKS_domain::get_R_size(*h_deltaR);
+                        torch::Tensor vdr_precalc;
+                        DeePKS_domain::cal_vdr_precalc(nlocal,
+                                                       lmaxd,
+                                                       inlmax,
+                                                       nat,
+                                                       nks,
+                                                       R_size,
+                                                       inl2l,
+                                                       kvec_d,
+                                                       phialpha,
+                                                       gevdm,
+                                                       inl_index,
+                                                       ucell,
+                                                       orb,
+                                                       *ParaV,
+                                                       GridD,
+                                                       vdr_precalc);
+
+                        const std::string file_vdrpre = PARAM.globalv.global_out_dir + "deepks_vdrpre.npy";
+                        LCAO_deepks_io::save_tensor2npy<double>(file_vdrpre, vdr_precalc, rank);
+                    }
+                    else if (PARAM.inp.deepks_v_delta == -2)
+                    {
+                        int R_size = DeePKS_domain::get_R_size(*h_deltaR);
+                        torch::Tensor phialpha_r_out;
+                        DeePKS_domain::prepare_phialpha_r(nlocal,
+                                                          lmaxd,
+                                                          inlmax,
+                                                          nat,
+                                                          R_size,
+                                                          phialpha,
+                                                          ucell,
+                                                          orb,
+                                                          *ParaV,
+                                                          GridD,
+                                                          phialpha_r_out);
+                        const std::string file_phialpha_r = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy";
+                        LCAO_deepks_io::save_tensor2npy<double>(file_phialpha_r, phialpha_r_out, rank);
+
+                        torch::Tensor gevdm_out;
+                        DeePKS_domain::prepare_gevdm(nat, lmaxd, inlmax, orb, gevdm, gevdm_out);
+                        const std::string file_gevdm = PARAM.globalv.global_out_dir + "deepks_gevdm.npy";
+                        LCAO_deepks_io::save_tensor2npy<double>(file_gevdm, gevdm_out, rank);
+                    }
+                }
             }
         }
 
@@ -544,7 +573,7 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
                         << " = " << std::setprecision(8) << e_delta_band * ModuleBase::Ry_to_eV << " eV" << std::endl;
             ofs_running << "  E_delta_NN = " << std::setprecision(8) << E_delta << " Ry"
                         << " = " << std::setprecision(8) << E_delta * ModuleBase::Ry_to_eV << " eV" << std::endl;
-                        ofs_running << " -----------------------------------------------" << std::endl;
+            ofs_running << " -----------------------------------------------" << std::endl;
         }
         if (PARAM.inp.deepks_out_unittest)
         {

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.h

@@ -3,9 +3,9 @@
 
 #ifdef __MLALGO
 #include "LCAO_deepks.h"
+#include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"
 #include "source_base/complexmatrix.h"
 #include "source_base/matrix.h"
-#include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"
 
 #include <memory>
 

source/module_hamilt_lcao/module_deepks/deepks_check.cpp

@@ -9,7 +9,8 @@ void DeePKS_domain::check_tensor(const torch::Tensor& tensor, const std::string&
     {
         return;
     }
-    using T_tensor = typename std::conditional<std::is_same<T, std::complex<double>>::value, c10::complex<double>, T>::type;
+    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);
@@ -21,15 +22,18 @@ void DeePKS_domain::check_tensor(const torch::Tensor& tensor, const std::string&
 
     // stride for each dimension
     std::vector<int64_t> strides(ndim, 1);
-    for (int i = ndim - 2; i >= 0; --i) {
+    for (int i = ndim - 2; i >= 0; --i)
+    {
         strides[i] = strides[i + 1] * sizes[i + 1];
     }
 
-    for (int64_t idx = 0; idx < numel; ++idx) {
+    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) {
+        for (int d = 0; d < ndim; ++d)
+        {
             indices[d] = tmp / strides[d];
             tmp = tmp % strides[d];
         }
@@ -39,16 +43,21 @@ void DeePKS_domain::check_tensor(const torch::Tensor& tensor, const std::string&
         ofs << *tmp_ptr;
 
         // print space or newline
-        if ( ( (idx+1) % sizes[ndim-1] ) == 0 ) {
+        if (((idx + 1) % sizes[ndim - 1]) == 0)
+        {
             ofs << std::endl;
-        } else {
+        }
+        else
+        {
             ofs << " ";
         }
     }
 
     ofs.close();
 }
 
+
+
 template void DeePKS_domain::check_tensor<int>(const torch::Tensor& tensor, const std::string& filename, const int rank);
 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);