add output EDM(R) and some test funcs

Author: maki49

source/module_lr/Grad/esolver_lr_grad.cpp

@@ -24,6 +24,41 @@ inline void print_force(const std::vector<ModuleBase::matrix>& force, Tstream& o
         }
     }
 }
+
+// check C_uaC_va-C_uiC_vi of lumo-homo, nocc=1,  nk=1
+template<typename T>
+inline void test_dm_diff_H2(const T* dm, const psi::Psi<T>& c, const int nbasis)
+{
+    std::cout << "difference dm cal: " << std::endl;
+    LR_Util::print_value(dm, nbasis, nbasis);
+    std::cout << "difference dm ref: " << std::endl;
+    for (int i = 0;i < nbasis;++i)
+    {
+        for (int j = 0;j < nbasis;++j)
+        {
+            std::cout << c(0, 1, i) * c(0, 1, j) - c(0, 0, i) * c(0, 0, j) << " ";
+        }
+        std::cout << std::endl;
+    }
+}
+
+// check e_aC_uaC_va-e_iC_uiC_vi of lumo-homo, nocc=1,  nk=1
+template<typename T>
+inline void test_edm_H2(const T* const edm, const double* const eig_ks, const psi::Psi<T>& c, const int nbasis)
+{
+    std::cout << "edm cal: " << std::endl;
+    LR_Util::print_value(edm, nbasis, nbasis);
+    std::cout << "edm ref: " << std::endl;
+    for (int i = 0;i < nbasis;++i)
+    {
+        for (int j = 0;j < nbasis;++j)
+        {
+            std::cout << eig_ks[1] * c(0, 1, i) * c(0, 1, j) - eig_ks[0] * c(0, 0, i) * c(0, 0, j) << " ";
+        }
+        std::cout << std::endl;
+    }
+}
+
 template<typename T, typename TR>
 void LR::ESolver_LR<T, TR>::init_pot_groundstate(const Charge& chg_gs)
 {
@@ -145,8 +180,21 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
                 this->gint_, pot_weak, pot_hxc_gs_weak,
                 this->kv, this->gd, this->paraX_, this->paraC_, this->paraMat_,
                 has_local_xc(this->xc_kernel));
+        if (PARAM.inp.test_force && nocc[0] == 1)
+        {
+            const std::vector<ct::Tensor>& dm_diff = cal_dm_diff_pblas(this->X[0].template data<T>() + offset, this->paraX_[0], c, this->paraC_, this->nbasis, this->nocc[0], this->nvirt[0], this->paraMat_);
+            // test_dm_diff_H2<T>(relaxed_diff_dm.get_DMK_pointer(0), c, this->nbasis);
+            test_dm_diff_H2<T>(dm_diff[0].data<T>(), c, this->nbasis);
+            test_edm_H2<T>(edm_k[0].data<T>(), this->eig_ks.c, c, this->nbasis);
+        }
         elecstate::DensityMatrix<T, double> edm_real = LR_Util::build_dm_from_dmk<T, double>(edm_k,
             this->paraMat_, this->nk, this->kv.kvec_d, this->ucell, this->gd, this->orb_cutoff_);
+        // print edm_real (R)
+        if (PARAM.inp.test_force)
+        {
+            LR_Util::save_DMR(edm_real, "data-EDMR-sparse", this->paraMat_);
+            // LR_Util::print_DMR(edm_real, this->ucell.nat, "edm_real (R) of istate " + std::to_string(istate));
+        }
 
         ModuleBase::matrix force_hxc_dmtrans = lr_force.cal_force_hxc_dmtrans(dm_trans_real, *this->pot[ispin]);
         std::cout << "Force (Hxc-DMTrans term) of state " << istate << ": " << std::endl;

source/module_lr/utils/lr_util_hcontainer.h

@@ -4,6 +4,8 @@
 #include  "module_base/parallel_reduce.h"
 #include "module_base/macros.h"
 #include <ATen/core/tensor.h>
+#include "module_parameter/parameter.h"
+#include "module_io/single_R_io.h"
 namespace LR_Util
 {
     template <typename T>
@@ -189,4 +191,109 @@ namespace LR_Util
         if (cal_dmr) { dm.cal_DMR(); }
         return dm;
     }
+
+    namespace sparse_format
+    {
+        // ref: sparse_format::cal_HContainer_d/cd and sparse_format::cal_HSR
+        // but more general(not depend on LCAO_HS_Arrays)
+        template<typename TR>
+        std::map<Abfs::Vector3_Order<int>, std::map<size_t, std::map<size_t, TR>>>
+            get_sparse_format(
+                const hamilt::HContainer<TR>& hR,
+                const Parallel_Orbitals& pv,
+                const double& sparse_thr = 1e-10)
+        {
+            std::map<Abfs::Vector3_Order<int>, std::map<size_t, std::map<size_t, TR>>> target;
+            auto row_indexes = pv.get_indexes_row();
+            auto col_indexes = pv.get_indexes_col();
+            for (int iap = 0; iap < hR.size_atom_pairs(); ++iap) {
+                int atom_i = hR.get_atom_pair(iap).get_atom_i();
+                int atom_j = hR.get_atom_pair(iap).get_atom_j();
+                int start_i = pv.atom_begin_row[atom_i];
+                int start_j = pv.atom_begin_col[atom_j];
+                int row_size = pv.get_row_size(atom_i);
+                int col_size = pv.get_col_size(atom_j);
+                for (int iR = 0; iR < hR.get_atom_pair(iap).get_R_size(); ++iR) {
+                    auto& matrix = hR.get_atom_pair(iap).get_HR_values(iR);
+                    const ModuleBase::Vector3<int> r_index
+                        = hR.get_atom_pair(iap).get_R_index(iR);
+                    Abfs::Vector3_Order<int> dR(r_index.x, r_index.y, r_index.z);
+                    for (int i = 0; i < row_size; ++i) {
+                        int mu = row_indexes[start_i + i];
+                        for (int j = 0; j < col_size; ++j) {
+                            int nu = col_indexes[start_j + j];
+                            const auto& value_tmp = matrix.get_value(i, j);
+                            if (std::abs(value_tmp) > sparse_thr) {
+                                target[dR][mu][nu] = value_tmp;
+                            }
+                        }
+                    }
+                }
+            }
+            return target;
+        }
+
+        //a more general version of save_HSR_sparse (not depend on LCAO_HS_Arrays)
+        template<typename TR>
+        void save_sparse(
+            const std::map<Abfs::Vector3_Order<int>, std::map<size_t, std::map<size_t, TR>>>& smat,
+            // const std::set<Abfs::Vector3_Order<int>>& all_R_coor,
+            // const bool& binary,
+            const std::string& filename,
+            const Parallel_Orbitals& pv,
+            const double& sparse_thr = 1e-10)
+        {
+            // calculate the total number of non-zero elements of the (nbasis, nbasis) matrix for each R
+            std::vector<int> non_zero_counts(smat.size(), 0);   // number of Rs
+            int i = 0;
+            for (const auto& Rij : smat)
+                non_zero_counts[i++] = std::accumulate(Rij.second.begin(), Rij.second.end(), 0,
+                    [](int sum, const auto& line) { return sum + line.second.size(); });
+            Parallel_Reduce::reduce_all(non_zero_counts.data(), non_zero_counts.size());
+
+
+            std::string out_dir = PARAM.globalv.global_out_dir + filename;
+            std::ofstream ofs;
+            if (GlobalV::DRANK == 0)
+            {
+                ofs.open(out_dir);
+                // if (binary) ofs.open(out_dir, std::ios::binary);
+                ofs << "STEP: 0" << std::endl;
+                ofs << "Matrix Dimension: " << PARAM.globalv.nlocal << std::endl;
+                ofs << "Matrix number: " << non_zero_counts.size() << std::endl;
+            }
+            i = 0;
+            for (const auto& Rij : smat)
+            {
+                const auto& R = Rij.first;
+                ofs << R.x << " " << R.y << " " << R.z << " " << non_zero_counts[i++] << std::endl;
+                ModuleIO::output_single_R(ofs, Rij.second, sparse_thr, false, pv);
+            }
+            if (GlobalV::DRANK == 0) { ofs.close(); }
+        }
+    }
+
+    template<typename TR>
+    void save_HR(
+        const hamilt::HContainer<TR>& hR,
+        // const std::set<Abfs::Vector3_Order<int>>& all_R_coor,
+        // const bool& binary,
+        const std::string& filename,
+        const Parallel_Orbitals& pv,
+        const double& sparse_thr = 1e-10)
+    {
+        sparse_format::save_sparse(sparse_format::get_sparse_format(hR, pv, sparse_thr),
+            filename, pv, sparse_thr);
+    }
+
+    template <typename TK, typename TR>
+    void save_DMR(const elecstate::DensityMatrix<TK, TR>& DMR,
+        const std::string& filename,
+        const Parallel_Orbitals& pv,
+        const double& sparse_thr = 1e-10)
+    {
+        int is = 0;
+        for (auto& dr : DMR.get_DMR_vector())
+            save_HR(*dr, filename + "_s" + std::to_string(is), pv, sparse_thr);
+    }
 }