fix: reproducing gs-dm and gs-force

Author: maki49

source/module_lr/Grad/esolver_lr_grad.cpp

@@ -120,7 +120,7 @@ ct::Tensor LR::ESolver_LR<T, TR>::solve_zvector_eqation(const int ispin)
 template<typename T, typename TR>
 std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin)
 {
-    if (PARAM.inp.test_force) { this->test_force(ispin); }
+    if (PARAM.inp.test_force && ispin == 0) { this->test_force(); }
 
     const ct::Tensor& Z = this->solve_zvector_eqation(ispin);
 
@@ -130,7 +130,7 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
 
     // calculate the force (the partial gradient of Lagrangian)
     LR_Force<T> lr_force(this->ucell, this->kv.kvec_d, this->paraMat_, *this->pw_rho, this->locpp, this->sf, this->gd, this->gint_, this->two_center_bundle_);
-    GlobalV::ofs_running << "Start to calculate excited-state force of  " << this->spin_types[ispin] << std::endl;
+    GlobalV::ofs_running << "Start to calculate excited-state force of " << this->spin_types[ispin] << std::endl;
     // ground state dm for currrent spin (only for test the correctness of the force)
     // elecstate::DensityMatrix<T, T> dm_gs(this->paraMat_, 1, this->kv.kvec_d, this->nk);
 
@@ -199,7 +199,7 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
         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;
         LR_Util::print_value(force_hxc_dmtrans.c, ucell.nat, 3);
-        ModuleBase::matrix force_hamiltgs_relaxed_diff = lr_force.cal_force_hamilt_gs_dm_relaxed_diff(relaxed_diff_dm_real, *pot_gs);
+        ModuleBase::matrix force_hamiltgs_relaxed_diff = lr_force.cal_force_hamilt_gs_dm_relaxed_diff(relaxed_diff_dm_real, *pot_gs, /*with_ewald=*/false);
         std::cout << "Force (GS-(T+Z) term) of state " << istate << ": " << std::endl;
         LR_Util::print_value(force_hamiltgs_relaxed_diff.c, ucell.nat, 3);
         ModuleBase::matrix force_overlap_edm = lr_force.cal_force_overlap_edm(edm_real);    // "-" sign has been included in the force factor
@@ -215,21 +215,22 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
 }
 
 template<typename T, typename TR>
-void LR::ESolver_LR<T, TR>::test_force(const int ispin)
+void LR::ESolver_LR<T, TR>::test_force()
 {
     LR_Force<T> lr_force(this->ucell, this->kv.kvec_d, this->paraMat_, *this->pw_rho, this->locpp, this->sf, this->gd, this->gint_, this->two_center_bundle_);
 
     elecstate::DensityMatrix<T, double> dm_gs(&this->paraMat_, this->nspin, this->kv.kvec_d, this->nk);   //DX
-    elecstate::cal_dm_psi(&this->paraMat_, this->wg_ks, *this->psi_ks, dm_gs);
-    LR_Util::initialize_DMR(dm_gs, this->paraMat_, this->ucell, this->gd, this->orb_cutoff_);
+    elecstate::cal_dm_psi(&this->paraMat_all_, this->wg_ks_all, *this->psi_ks_all, dm_gs);   // nbands is important here
+    LR_Util::initialize_DMR(dm_gs, this->paraMat_, this->ucell, this->gd, this->orb_cutoff_);   // nbands is not important here
     dm_gs.cal_DMR();
-
+    LR_Util::print_DMR(dm_gs, this->ucell.nat, "DM(R) of ground state");
     ///========================== test 1: reproduce the force of ground state =========================
     // energy density matrix of the ground state
     elecstate::DensityMatrix<T, double> edm_gs(&this->paraMat_, this->nspin, this->kv.kvec_d, this->nk);   //DX
-    ModuleBase::matrix wg_ekb_ks(nspin, nbands);
-    std::transform(this->wg_ks.c, this->wg_ks.c + nspin * nbands, this->eig_ks.c, wg_ekb_ks.c, std::multiplies<double>());
-    elecstate::cal_dm_psi(&this->paraMat_, wg_ekb_ks, *this->psi_ks, edm_gs);
+    ModuleBase::matrix wg_ekb_ks_all(nspin, PARAM.inp.nbands);
+    std::transform(this->wg_ks_all.c, this->wg_ks_all.c + nspin * PARAM.inp.nbands,
+        this->eig_ks_all.c, wg_ekb_ks_all.c, std::multiplies<double>());
+    elecstate::cal_dm_psi(&this->paraMat_all_, wg_ekb_ks_all, *this->psi_ks_all, edm_gs);
     LR_Util::initialize_DMR(edm_gs, this->paraMat_, this->ucell, this->gd, this->orb_cutoff_);
     edm_gs.cal_DMR();
     // ground-state force
@@ -239,7 +240,7 @@ void LR::ESolver_LR<T, TR>::test_force(const int ispin)
     ///========================== test 2: reproduce the DX Hartree term =========================
     ModuleBase::matrix f_hxc_potgs = lr_force.reproduce_force_gs_loc(dm_gs, *this->pot_gs_hartree);
     ModuleIO::print_force(GlobalV::ofs_running, this->ucell, "GS Hartree force calculated by 'cal_pulay_fs' from potential (eV/Angstrom)", f_hxc_potgs, false);
-    ModuleBase::matrix f_hxc_potlr = lr_force.cal_force_hxc_dmtrans(dm_gs, *this->pot[ispin]);
+    ModuleBase::matrix f_hxc_potlr = lr_force.cal_force_hxc_dmtrans(dm_gs, *this->pot[0]);
     ModuleIO::print_force(GlobalV::ofs_running, this->ucell, "2* GS Hxc force calculated by 'LR_Force' from kernel (eV/Angstrom)", f_hxc_potlr * 2, false);
     // 2 for spin in f->v. Spin in v->f is already multiplied in the singlet Hartree factor 2.
     /// ======================================= END test 2 =========================================

source/module_lr/Grad/force/lr_force.cpp

@@ -26,12 +26,13 @@ namespace LR
     }
 
     template<typename TK>
-    ModuleBase::matrix LR_Force<TK>::cal_force_hamilt_gs_dm_relaxed_diff(const elecstate::DensityMatrix<TK, double>& relax_diff_dm, const elecstate::Potential& pot_gs)
+    ModuleBase::matrix LR_Force<TK>::cal_force_hamilt_gs_dm_relaxed_diff(const elecstate::DensityMatrix<TK, double>& relax_diff_dm,
+        const elecstate::Potential& pot_gs, const bool with_ewald)
     {
         const Charge chr_diff_relaxed = dm_to_charge(relax_diff_dm);
 
         // 1. local pp (Hellmann-Feynman)(fvl_dvl) + ewald + core correction (+ self-consistent charge)
-        ModuleBase::matrix f_pw = ForcePWTerms<double>()(this->ucell_, chr_diff_relaxed, this->rhopw_, this->locpp_, this->sf_, /*with_ewald=*/false);
+        ModuleBase::matrix f_pw = ForcePWTerms<double>()(this->ucell_, chr_diff_relaxed, this->rhopw_, this->locpp_, this->sf_, with_ewald);
 
         // 2. nonlocal pp (Hellmann-Feynman + Pulay)
         ModuleBase::matrix fvnl = cal_force_nonlocal(this->ucell_, this->kvec_d_, this->gd_, this->two_center_bundle_, relax_diff_dm);

source/module_lr/Grad/force/lr_force.h

@@ -25,7 +25,8 @@ namespace LR
         }
 
         /// 1. $Tr[H_{GS}^x * (T+D^Z)]$, where GS=groud state and $(T+D^Z)$ is the relaxed difference density matrix
-        ModuleBase::matrix cal_force_hamilt_gs_dm_relaxed_diff(const elecstate::DensityMatrix<TK, double>& relaxed_diff_dm, const elecstate::Potential& pot_gs);
+        ModuleBase::matrix cal_force_hamilt_gs_dm_relaxed_diff(const elecstate::DensityMatrix<TK, double>& relaxed_diff_dm,
+            const elecstate::Potential& pot_gs, const bool with_ewald = true);
 
         /// 2. $Tr[S^x * (EDM)]
         ModuleBase::matrix cal_force_overlap_edm(const elecstate::DensityMatrix<TK, double>& edm);

source/module_lr/Grad/multipliers/hamilt_zeq_right.h

@@ -103,7 +103,7 @@ namespace LR
             for (int ib = 0;ib < nband;++ib)
             {
                 const int offset = ib * ld_psi;
-                this->cal_dm_trans(0, psi + offset);  // transition density matrix
+                // this->cal_dm_trans(0, psi + offset);  // transition density matrix, only for test
                 this->cal_dm_diff(0, psi + offset);  // difference density matrix
                 hamilt::Operator<T>* node(this->ops);
                 while (node != nullptr)

source/module_lr/esolver_lrtd_lcao.cpp

@@ -181,12 +181,14 @@ LR::ESolver_LR<T, TR>::ESolver_LR(ModuleESolver::ESolver_KS_LCAO<T, TR>&& ks_sol
 
     this->set_dimension();
 
-    // setup_wd_division is not need to be covered in #ifdef __MPI, see its implementation
+    // setup_2d_division is not need to be covered in #ifdef __MPI, see its implementation
     LR_Util::setup_2d_division(this->paraMat_, 1, this->nbasis, this->nbasis);
-
-    this->paraMat_.atom_begin_row = std::move(ks_sol.pv.atom_begin_row);
-    this->paraMat_.atom_begin_col = std::move(ks_sol.pv.atom_begin_col);
-    this->paraMat_.iat2iwt_ = ucell.get_iat2iwt();
+    this->set_parallel_orbitals_band(this->paraMat_, this->nbands);
+    if (PARAM.inp.cal_force)
+    {
+        LR_Util::setup_2d_division(this->paraMat_all_, 1, this->nbasis, this->nbasis);
+        this->set_parallel_orbitals_band(this->paraMat_all_, PARAM.inp.nbands);
+    }
 
     LR_Util::setup_2d_division(this->paraC_, 1, this->nbasis, this->nbands
 #ifdef __MPI
@@ -195,32 +197,42 @@ LR::ESolver_LR<T, TR>::ESolver_LR(ModuleESolver::ESolver_KS_LCAO<T, TR>&& ks_sol
     );
     auto move_gs = [&, this]() -> void  // move the ground state info
         {
-            this->psi_ks = ks_sol.psi;
+            this->psi_ks_all = ks_sol.psi;
             ks_sol.psi = nullptr;
             //only need the eigenvalues. the 'elecstates' of excited states is different from ground state.
-            this->eig_ks = std::move(ks_sol.pelec->ekb);
+            this->eig_ks_all = std::move(ks_sol.pelec->ekb);
         };
+    move_gs();
+    // allocate psi_ks and eig_ks in the [nocc, nvirt] window
 #ifdef __MPI
-    if (this->nbands == PARAM.inp.nbands) { move_gs(); }
-    else    // copy the part of ground state info according to paraC_
+    this->psi_ks = new psi::Psi<T>(this->kv.get_nks(),
+        this->paraC_.get_col_size(),
+        this->paraC_.get_row_size(),
+        this->kv.ngk,
+        true);
+#else
+    this->psi_ks = new psi::Psi<T>(this->kv.get_nks(), this->nbands, this->nbasis, this->kv.ngk, true);
+#endif
+    this->eig_ks.create(this->kv.get_nks(), this->nbands);
+    const int start_band = this->nocc_max - *std::max_element(nocc.begin(), nocc.end());
+
+    for (int ik = 0;ik < this->kv.get_nks();++ik)
     {
-        this->psi_ks = new psi::Psi<T>(this->kv.get_nks(), 
-                                       this->paraC_.get_col_size(), 
-                                       this->paraC_.get_row_size(),
-                                       this->kv.ngk,
-                                       true);
-        this->eig_ks.create(this->kv.get_nks(), this->nbands);
-        const int start_band = this->nocc_max - *std::max_element(nocc.begin(), nocc.end());
-        for (int ik = 0;ik < this->kv.get_nks();++ik)
+        // copy the KS orbitals in the [nocc, nvirt] window
+#ifdef __MPI
+        Cpxgemr2d(this->nbasis, this->nbands, &(*this->psi_ks_all)(ik, 0, 0), 1, start_band + 1, ks_sol.pv.desc_wfc,
+            &(*this->psi_ks)(ik, 0, 0), 1, 1, this->paraC_.desc, this->paraC_.blacs_ctxt);
+#else 
+        for (int ib = 0;ib < this->nbands;++ib)
         {
-            Cpxgemr2d(this->nbasis, this->nbands, &(*ks_sol.psi)(ik, 0, 0), 1, start_band + 1, ks_sol.pv.desc_wfc,
-                &(*this->psi_ks)(ik, 0, 0), 1, 1, this->paraC_.desc, this->paraC_.blacs_ctxt);
-            for (int ib = 0;ib < this->nbands;++ib) { this->eig_ks(ik, ib) = ks_sol.pelec->ekb(ik, start_band + ib); }
+            auto* start = &(*this->psi_ks_all)(ik, start_band + ib, 0);
+            auto* to = &(*this->psi_ks)(ik, ib, 0);
         }
-    }
-#else
-    move_gs();
 #endif
+        // copy the KS bands in the [nocc, nvirt] window
+        for (int ib = 0;ib < this->nbands;++ib) { this->eig_ks(ik, ib) = this->eig_ks_all(ik, start_band + ib); }
+    }
+
     if (nspin == 2)
     {
         this->nupdown = cal_nupdown_form_occ(ks_sol.pelec->wg);
@@ -309,14 +321,12 @@ LR::ESolver_LR<T, TR>::ESolver_LR(const Input_para& inp, UnitCell& ucell) : inpu
     this->set_dimension();
     //  setup 2d-block distribution for AO-matrix and KS wfc
     LR_Util::setup_2d_division(this->paraMat_, 1, this->nbasis, this->nbasis);
-#ifdef __MPI
-    this->paraMat_.set_desc_wfc_Eij(this->nbasis, this->nbands, paraMat_.get_row_size());
-    int err = this->paraMat_.set_nloc_wfc_Eij(this->nbands, GlobalV::ofs_running, GlobalV::ofs_warning);
-    if (input.ri_hartree_benchmark != "aims") { this->paraMat_.set_atomic_trace(ucell.get_iat2iwt(), ucell.nat, this->nbasis); }
-#else
-    this->paraMat_.nrow_bands = this->nbasis;
-    this->paraMat_.ncol_bands = this->nbands;
-#endif
+    this->set_parallel_orbitals_band(this->paraMat_, this->nbands);
+    if (PARAM.inp.cal_force)
+    {
+        LR_Util::setup_2d_division(this->paraMat_all_, 1, this->nbasis, this->nbasis);
+        this->set_parallel_orbitals_band(this->paraMat_all_, PARAM.inp.nbands);
+    }
 
     // read the ground state info
     // now ModuleIO::read_wfc_nao needs `Parallel_Orbitals` and can only read all the bands
@@ -594,6 +604,18 @@ void LR::ESolver_LR<T, TR>::after_all_runners(UnitCell& ucell)
         if (PARAM.inp.cal_force) { this->cal_force(is); }
     }
 }
+template<typename T, typename TR>
+void LR::ESolver_LR<T, TR>::set_parallel_orbitals_band(Parallel_Orbitals& pmat, const int nbands_in)
+{
+#ifdef __MPI
+    pmat.set_desc_wfc_Eij(this->nbasis, nbands_in, pmat.get_row_size());
+    int err = pmat.set_nloc_wfc_Eij(nbands_in, GlobalV::ofs_running, GlobalV::ofs_warning);
+    if (input.ri_hartree_benchmark != "aims") { pmat.set_atomic_trace(ucell.get_iat2iwt(), ucell.nat, this->nbasis); }
+#else
+    pmat.nrow_bands = this->nbasis;
+    pmat.ncol_bands = nbands_in;
+#endif
+}
 
 template<typename T, typename TR>
 void LR::ESolver_LR<T, TR>::setup_eigenvectors_X()
@@ -717,8 +739,17 @@ void LR::ESolver_LR<T, TR>::read_ks_wfc()
         /*skip_bands=*/this->nocc_max - this->nocc_in)) {
         ModuleBase::WARNING_QUIT("ESolver_LR", "read ground-state wavefunction failed.");
     }
-    this->eig_ks = std::move(this->pelec->ekb);
-    this->wg_ks = std::move(this->pelec->wg);
+
+    if (PARAM.inp.cal_force)
+    {    // allocate psi_ks_all and eig_ks_all to read all the bands
+        this->psi_ks_all = new psi::Psi<T>(this->kv.get_nks(), paraMat_all_.ncol_bands, paraMat_all_.get_row_size(), this->kv.ngk, true);
+        this->eig_ks_all.create(this->kv.get_nks(), PARAM.inp.nbands);
+        this->wg_ks_all.create(this->kv.get_nks(), PARAM.inp.nbands);
+        if (!ModuleIO::read_wfc_nao(PARAM.globalv.global_readin_dir, paraMat_all_, *this->psi_ks_all, this->wg_ks_all, this->eig_ks_all,/*skip_bands=*/0))
+        {
+            GlobalV::ofs_running << " Read in all the KS wavefunctions for force calculation. " << std::endl;
+        }
+    }
 }
 
 template<typename T, typename TR>

source/module_lr/esolver_lrtd_lcao.h

@@ -35,6 +35,7 @@ namespace LR
         ESolver_LR(const Input_para& inp, UnitCell& ucell);
         ~ESolver_LR() {
             delete this->psi_ks;
+            delete this->psi_ks_all;
         }
 
         ///input: input, call, basis(LCAO), psi(ground state), elecstate
@@ -61,17 +62,22 @@ namespace LR
         // ground state info 
 
         /// @brief ground state wave function
-        psi::Psi<T>* psi_ks = nullptr;
+        psi::Psi<T>* psi_ks = nullptr;  ///< KS orbitals used in the [nocc+nvirt] window
+        psi::Psi<T>* psi_ks_all = nullptr;  ///< all KS orbitals, read from the file, or moved from ESolver_FP::pelec.psi
 
         /// @brief ground state bands, read from the file, or moved from ESolver_FP::pelec.ekb
-        ModuleBase::matrix eig_ks;///< energy of ground state
+        ModuleBase::matrix eig_ks;///< ground state eigenvalues in the [nocc+nvirt] window
+        ModuleBase::matrix eig_ks_all; ///< all eigenvalues of ground state, read from the file, or moved from ESolver_FP::pelec.ekb
+        ModuleBase::matrix wg_ks;   /// occupation numbers of ground state in the [nocc+nvirt] window
+        ModuleBase::matrix wg_ks_all;   /// occupation number of all bands of ground state
+
 
         // @brief only needed for force calculation 
         std::unique_ptr<elecstate::Potential> pot_gs;
         std::unique_ptr<elecstate::Potential> pot_gs_hartree;   /// ground-state Hartree potential, only used for test_force
         double etxc_gs = 0.;
         double vtxc_gs = 0.;
-        ModuleBase::matrix wg_ks;   /// occupation number of ground state
+
         std::shared_ptr<PotHxcLR> pot_hxc_gs; /// used in lr-grad, in the ground-state Hxc gradient term coming from dF/dC
 
         /// @brief Excited state wavefunction (locc, lvirt are local size of nocc and nvirt in each process)
@@ -112,6 +118,7 @@ namespace LR
         std::vector<Parallel_2D> paraX_;
         /// @brief variables for parallel distribution of matrix in AO representation
         Parallel_Orbitals paraMat_;
+        Parallel_Orbitals paraMat_all_; // for the parallelized size of the KS orbitals
 
         TwoCenterBundle two_center_bundle_;
 
@@ -133,11 +140,14 @@ namespace LR
         /// reset nocc, nvirt, npairs after read ground-state wavefunction when nspin=2
         void reset_dim_spin2();
 
+        /// setup Parallel_Orbitals info. beyond Parallel_2D
+        void set_parallel_orbitals_band(Parallel_Orbitals& p, const int nbands_in);
+
         ///========================== for gradient calculation =========================
         void init_pot_groundstate(const Charge& chg_gs);
         ct::Tensor solve_zvector_eqation(const int ispin);
         std::vector<ModuleBase::matrix> cal_force(const int ispin);
-        void test_force(const int ispin);   // test: reproduce the force of ground state
+        void test_force();   // test: reproduce the force of ground state
 
 #ifdef __EXX
         /// Tdata of Exx_LRI is same as T, for the reason, see operator_lr_exx.h

source/module_lr/utils/lr_util_hcontainer.h

@@ -42,7 +42,7 @@ namespace LR_Util
         std::cout << label << "\n";
         int is = 0;
         for (auto& dr : DMR.get_DMR_vector())
-            print_HR(*dr, nat, "DMR[" + std::to_string(is++) + "]", threshold);
+            print_HR(*dr, nat, "DMR[ispin=s" + std::to_string(is++) + "]", threshold);
     }
 
     template<typename T>