traverse states outside of act() (still bug)

source/module_lr/hamilt_casida.h

@@ -1,9 +1,11 @@
 #pragma once
+#include <typeinfo>
 #include "module_hamilt_general/hamilt.h"
 #include "module_elecstate/module_dm/density_matrix.h"
 #include "module_lr/operator_casida/operator_lr_diag.h"
 #include "module_lr/operator_casida/operator_lr_hxc.h"
 #include "module_basis/module_ao/parallel_orbitals.h"
+#include "module_lr/dm_trans/dm_trans.h"
 #ifdef __EXX
 #include "module_lr/operator_casida/operator_lr_exx.h"
 #include "module_lr/ri_benchmark/operator_ri_hartree.h"
@@ -42,8 +44,10 @@ namespace LR
         {
             ModuleBase::TITLE("HamiltLR", "HamiltLR");
             if (ri_hartree_benchmark != "aims") { assert(aims_nbasis.empty()); }
-            this->DM_trans.resize(1);
-            this->DM_trans[0] = LR_Util::make_unique<elecstate::DensityMatrix<T, T>>(&pmat_in, nspin, kv_in.kvec_d, nk);
+            // always use nspin=1 for transition density matrix
+            this->DM_trans = LR_Util::make_unique<elecstate::DensityMatrix<T, T>>(&pmat_in, 1, kv_in.kvec_d, nk);
+            this->DM_trans->init_DMR(&gd_in, &ucell_in);
+
             // add the diag operator  (the first one)
             this->ops = new OperatorLRDiag<T>(eig_ks.c, pX[0], nk, nocc[0], nvirt[0]);
             //add Hxc operator
@@ -104,11 +108,24 @@ namespace LR
                 hamilt::Operator<T>* lr_exx = new OperatorLREXX<T>(nspin, naos, nocc[0], nvirt[0], ucell_in, psi_ks_in,
                     this->DM_trans, exx_lri_in, kv_in, pX_in[0], pc_in, pmat_in,
                     xc_kernel == "hf" ? 1.0 : exx_alpha, //alpha
-                    ri_hartree_benchmark != "none"/*whether to cal_dm_trans first here*/,
                     aims_nbasis);
                 this->ops->add(lr_exx);
             }
 #endif
+
+            this->cal_dm_trans = [&, this](const int& is, const T* X)->void
+                {
+                    const auto psi_ks_is = LR_Util::get_psi_spin(psi_ks_in, is, nk);
+#ifdef __MPI
+                    std::vector<ct::Tensor>  dm_trans_2d = cal_dm_trans_pblas(X, pX[is], psi_ks_is, pc_in, naos, nocc[is], nvirt[is], pmat_in);
+                    if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos, pmat_in);
+#else
+                    std::vector<ct::Tensor>  dm_trans_2d = cal_dm_trans_blas(X, psi_ks_is, nocc[is], nvirt[is]);
+                    if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos);
+#endif
+                    // tensor to vector, then set DMK
+                    for (int ik = 0;ik < nk;++ik) { this->DM_trans->set_DMK_pointer(ik, dm_trans_2d[ik].data<T>()); }
+                };
         }
         ~HamiltLR()
         {
@@ -123,11 +140,16 @@ namespace LR
         virtual void hPsi(const T* psi_in, T* hpsi, const int ld_psi, const int& nband) const
         {
             assert(ld_psi == nk * pX[0].get_local_size());
-            hamilt::Operator<T>* node(this->ops);
-            while (node != nullptr)
+            for (int ib = 0;ib < nband;++ib)
             {
-                node->act(nband, ld_psi, /*npol=*/1, psi_in, hpsi);
-                node = (hamilt::Operator<T>*)(node->next_op);
+                const int offset = ib * ld_psi;
+                this->cal_dm_trans(0, psi_in + offset);  // calculate transition density matrix here
+                hamilt::Operator<T>* node(this->ops);
+                while (node != nullptr)
+                {
+                    node->act(/*nband=*/1, ld_psi, /*npol=*/1, psi_in + offset, hpsi + offset);
+                    node = (hamilt::Operator<T>*)(node->next_op);
+                }
             }
         }
 
@@ -136,14 +158,16 @@ namespace LR
         const std::vector<int>& nvirt;
         const int nspin = 1;
         const int nk = 1;
+        const bool tdm_sym = false;     ///< whether to symmetrize the transition density matrix
         const std::vector<Parallel_2D>& pX;
         T one()const;
         /// transition density matrix in AO representation
-        /// Hxc only: size=1, calculate on the same address for each bands
-        /// Hxc+Exx: size=nbands, store the result of each bands for common use
-        std::vector<std::unique_ptr<elecstate::DensityMatrix<T, T>>> DM_trans;
+        /// calculate on the same address for each bands, and commonly used by all the operators
+        std::unique_ptr<elecstate::DensityMatrix<T, T>> DM_trans;
 
         /// first node operator, add operations from each operators
         hamilt::Operator<T, base_device::DEVICE_CPU>* ops = nullptr;
+
+        std::function<void(const int&, const T*)> cal_dm_trans;
     };
 }

source/module_lr/hamilt_ulr.hpp

@@ -38,8 +38,8 @@ namespace LR
             nloc_per_band(nk* pX[0].get_local_size() + nk * pX[1].get_local_size())
         {
             ModuleBase::TITLE("HamiltULR", "HamiltULR");
-            this->DM_trans.resize(1);
-            this->DM_trans[0] = LR_Util::make_unique<elecstate::DensityMatrix<T, T>>(&pmat_in, nspin, kv_in.kvec_d, nk);
+            this->DM_trans = LR_Util::make_unique<elecstate::DensityMatrix<T, T>>(&pmat_in, 1, kv_in.kvec_d, nk);
+            this->DM_trans->init_DMR(&gd_in, &ucell_in);
 
             // how to change the index of eig_ks and psi_ks_In?
             // modify the interface of opetators to support different left- and right- spin-pairs
@@ -68,6 +68,20 @@ namespace LR
                 }
             }
 #endif
+
+            this->cal_dm_trans = [&, this](const int& is, const T* X)->void
+                {
+                    const auto psi_ks_is = LR_Util::get_psi_spin(psi_ks_in, is, nk);
+#ifdef __MPI
+                    std::vector<ct::Tensor>  dm_trans_2d = cal_dm_trans_pblas(X, pX[is], psi_ks_is, pc_in, naos, nocc[is], nvirt[is], pmat_in);
+                    if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos, pmat_in);
+#else
+                    std::vector<ct::Tensor>  dm_trans_2d = cal_dm_trans_blas(X, psi_ks_is, nocc[is], nvirt[is]);
+                    if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos);
+#endif
+                    // tensor to vector, then set DMK
+                    for (int ik = 0;ik < nk;++ik) { this->DM_trans->set_DMK_pointer(ik, dm_trans_2d[ik].data<T>()); }
+                };
         }
         ~HamiltULR()
         {
@@ -85,6 +99,7 @@ namespace LR
                 for (int is_bj : {0, 1})
                 {
                     const int offset = offset_b + is_bj * xdim_is[0];
+                    cal_dm_trans(is_bj, psi_in + offset);   // calculate transition density matrix here
                     for (int is_ai : {0, 1})
                     {
                         hamilt::Operator<T>* node(this->ops[(is_ai << 1) + is_bj]);
@@ -195,6 +210,9 @@ namespace LR
         /// transition density matrix in AO representation
         /// Hxc only: size=1, calculate on the same address for each bands
         /// Hxc+Exx: size=nbands, store the result of each bands for common use
-        std::vector<std::unique_ptr<elecstate::DensityMatrix<T, T>>> DM_trans;
+        std::unique_ptr<elecstate::DensityMatrix<T, T>> DM_trans;
+
+        std::function<void(const int&, const T*)> cal_dm_trans;
+        const bool tdm_sym = false;     ///< whether to symmetrize the transition density matrix
     };
 }

source/module_lr/lr_spectrum.cpp

@@ -33,7 +33,7 @@ void LR::LR_Spectrum<double>::oscillator_strength()
     osc.resize(nstate, 0.0);
     // const int nspin0 = (this->nspin == 2) ? 2 : 1;   use this in NSPIN=4 implementation
     double osc_tot = 0.0;
-    elecstate::DensityMatrix<double, double> DM_trans(&this->pmat, this->nspin, this->kv.kvec_d, this->nk);
+    elecstate::DensityMatrix<double, double> DM_trans(&this->pmat, 1, this->kv.kvec_d, this->nk);
     DM_trans.init_DMR(&GlobalC::GridD, &this->ucell);
     this->transition_dipole_.resize(nstate, ModuleBase::Vector3<double>(0.0, 0.0, 0.0));
     for (int istate = 0;istate < nstate;++istate)
@@ -87,9 +87,9 @@ void LR::LR_Spectrum<std::complex<double>>::oscillator_strength()
     osc.resize(nstate, 0.0);
     // const int nspin0 = (this->nspin == 2) ? 2 : 1;   use this in NSPIN=4 implementation
     double osc_tot = 0.0;
-    elecstate::DensityMatrix<std::complex<double>, std::complex<double>> DM_trans(&this->pmat, this->nspin, this->kv.kvec_d, this->nk);
+    elecstate::DensityMatrix<std::complex<double>, std::complex<double>> DM_trans(&this->pmat, 1, this->kv.kvec_d, this->nk);
     DM_trans.init_DMR(&GlobalC::GridD, &this->ucell);
-    elecstate::DensityMatrix<std::complex<double>, double> DM_trans_real_imag(&this->pmat, this->nspin, this->kv.kvec_d, this->nk);
+    elecstate::DensityMatrix<std::complex<double>, double> DM_trans_real_imag(&this->pmat, 1, this->kv.kvec_d, this->nk);
     DM_trans_real_imag.init_DMR(&GlobalC::GridD, &this->ucell);
 
     this->transition_dipole_.resize(nstate, ModuleBase::Vector3<std::complex<double>>(0.0, 0.0, 0.0));

source/module_lr/operator_casida/operator_lr_exx.cpp

@@ -89,25 +89,12 @@ namespace LR
         for (int ib = 0;ib < nbands;++ib)
         {
             const int xstart_b = ib * nk * pX.get_local_size();
-            // suppose Cs，Vs, have already been calculated in the ion-step of ground state, 
-            // DM_trans(k) and DM_trans(R) has already been calculated from psi_in in OperatorLRHxc::act
-            // but int RI_benchmark, DM_trans(k) should be first calculated here
-            if (cal_dm_trans)
-            {
-#ifdef __MPI
-                std::vector<container::Tensor>  dm_trans_2d = cal_dm_trans_pblas(psi_in + xstart_b, pX, psi_ks, pc, naos, nocc, nvirt, pmat);
-                if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos, pmat);
-#else
-                std::vector<container::Tensor>  dm_trans_2d = cal_dm_trans_blas(psi_in + xstart, psi_ks, nocc, nvirt);
-                if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos);
-#endif
-                // tensor to vector, then set DMK
-                for (int ik = 0;ik < nk;++ik) { this->DM_trans[ib]->set_DMK_pointer(ik, dm_trans_2d[ik].data<T>()); }
-            }
+            // suppose Cs，Vs, have already been calculated in the ion-step of ground state
+            // and DM_trans has been calculated in hPsi() outside.
 
             // 1. set_Ds (once)
             // convert to vector<T*> for the interface of RI_2D_Comm::split_m2D_ktoR (interface will be unified to ct::Tensor)
-            std::vector<std::vector<T>> DMk_trans_vector = this->DM_trans[ib]->get_DMK_vector();
+            std::vector<std::vector<T>> DMk_trans_vector = this->DM_trans->get_DMK_vector();
             // assert(DMk_trans_vector.size() == nk);
             std::vector<const std::vector<T>*> DMk_trans_pointer(nk);
             for (int ik = 0;ik < nk;++ik) {DMk_trans_pointer[ik] = &DMk_trans_vector[ik];}

source/module_lr/operator_casida/operator_lr_exx.h

@@ -23,19 +23,18 @@ namespace LR
             const int& nvirt,
             const UnitCell& ucell_in,
             const psi::Psi<T>& psi_ks_in,
-            std::vector<std::unique_ptr<elecstate::DensityMatrix<T, T>>>& DM_trans_in,
+            std::unique_ptr<elecstate::DensityMatrix<T, T>>& DM_trans_in,
             // HContainer<double>* hR_in,
             std::weak_ptr<Exx_LRI<T>> exx_lri_in,
             const K_Vectors& kv_in,
             const Parallel_2D& pX_in,
             const Parallel_2D& pc_in,
             const Parallel_Orbitals& pmat_in,
             const double& alpha = 1.0,
-            const bool& cal_dm_trans = false,
             const std::vector<int>& aims_nbasis = {})
             : nspin(nspin), naos(naos), nocc(nocc), nvirt(nvirt),
             psi_ks(psi_ks_in), DM_trans(DM_trans_in), exx_lri(exx_lri_in), kv(kv_in),
-            pX(pX_in), pc(pc_in), pmat(pmat_in), ucell(ucell_in), alpha(alpha), cal_dm_trans(cal_dm_trans),
+            pX(pX_in), pc(pc_in), pmat(pmat_in), ucell(ucell_in), alpha(alpha),
             aims_nbasis(aims_nbasis)
         {
             ModuleBase::TITLE("OperatorLREXX", "OperatorLREXX");
@@ -75,7 +74,7 @@ namespace LR
         const std::vector<int> aims_nbasis={};    ///< number of basis functions for each type of atom in FHI-aims
 
         /// transition density matrix 
-        std::vector<std::unique_ptr<elecstate::DensityMatrix<T, T>>>& DM_trans;
+        std::unique_ptr<elecstate::DensityMatrix<T, T>>& DM_trans;
 
         /// density matrix of a certain (i, a, k), with full naos*naos size for each key
         /// D^{iak}_{\mu\nu}(k): 1/N_k * c^*_{ak,\mu} c_{ik,\nu}

source/module_lr/operator_casida/operator_lr_hxc.cpp

@@ -8,7 +8,6 @@
 #include "module_lr/utils/lr_util_print.h"
 // #include "module_hamilt_lcao/hamilt_lcaodft/DM_gamma_2d_to_grid.h"
 #include "module_hamilt_lcao/module_hcontainer/hcontainer_funcs.h"
-#include "module_lr/dm_trans/dm_trans.h"
 #include "module_lr/AX/AX.h"
 #include "module_hamilt_pw/hamilt_pwdft/global.h"
 
@@ -22,41 +21,16 @@ namespace LR
     {
         ModuleBase::TITLE("OperatorLRHxc", "act");
         const int& sl = ispin_ks[0];
-        const int& sr = ispin_ks.size() == 1 ? sl : ispin_ks[1];
-        const auto psir_ks = LR_Util::get_psi_spin(psi_ks, sr, nk);
         const auto psil_ks = LR_Util::get_psi_spin(psi_ks, sl, nk);
 
-        this->init_DM_trans(nbands, this->DM_trans);    // initialize transion density matrix
-
         const int& lgd = gint->gridt->lgd;
         for (int ib = 0;ib < nbands;++ib)
         {
-            // if Hxc-only, the memory of single-band DM_trans is enough.
-            // if followed by EXX, we need to allocate memory for all bands.
-            const int ib_dm = (this->next_op == nullptr) ? 0 : ib;
             const int xstart_b = ib * nbasis;
 
-            // 1. transition density matrix
-#ifdef __MPI
-            std::vector<container::Tensor>  dm_trans_2d = cal_dm_trans_pblas(psi_in + xstart_b, pX[sr], psir_ks, pc, naos, nocc[sr], nvirt[sr], pmat);
-            if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos, pmat);
-#else
-            std::vector<container::Tensor>  dm_trans_2d = cal_dm_trans_blas(psi_in + xstart_b, psir_ks, nocc[sr], nvirt[sr]);
-            if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data<T>(), naos);
-#endif
-            // tensor to vector, then set DMK
-            for (int ik = 0;ik < nk;++ik) { this->DM_trans[ib_dm]->set_DMK_pointer(ik, dm_trans_2d[ik].data<T>()); }
-
-            // if (this->first_print)
-            //     for (int ik = 0;ik < nk;++ik)
-            //         LR_Util::print_tensor<T>(dm_trans_2d[ik], "1.DMK[ik=" + std::to_string(ik) + "]", this->pmat);
-
-            // use cal_DMR to get DMR form DMK by FT
-            this->DM_trans[ib_dm]->cal_DMR();  //DM_trans->get_DMR_vector() is 2d-block parallized
-            // LR_Util::print_DMR(*this->DM_trans[0], ucell.nat, "DM(R) (complex)");
-
+            this->DM_trans->cal_DMR();  //DM_trans->get_DMR_vector() is 2d-block parallized
             // ========================= begin grid calculation=========================
-            this->grid_calculation(nbands, ib_dm);   //DM(R) to H(R)
+            this->grid_calculation(nbands);   //DM(R) to H(R)
             // ========================= end grid calculation =========================
 
             // V(R)->V(k) 
@@ -80,11 +54,11 @@ namespace LR
 
 
     template<>
-    void OperatorLRHxc<double, base_device::DEVICE_CPU>::grid_calculation(const int& nbands, const int& iband_dm) const
+    void OperatorLRHxc<double, base_device::DEVICE_CPU>::grid_calculation(const int& nbands) const
     {
         ModuleBase::TITLE("OperatorLRHxc", "grid_calculation(real)");
         ModuleBase::timer::tick("OperatorLRHxc", "grid_calculation");
-        this->gint->transfer_DM2DtoGrid(this->DM_trans[iband_dm]->get_DMR_vector());     // 2d block to grid
+        this->gint->transfer_DM2DtoGrid(this->DM_trans->get_DMR_vector());     // 2d block to grid
 
         // 2. transition electron density
         // \f[ \tilde{\rho}(r)=\sum_{\mu_j, \mu_b}\tilde{\rho}_{\mu_j,\mu_b}\phi_{\mu_b}(r)\phi_{\mu_j}(r) \f]
@@ -116,7 +90,7 @@ namespace LR
     }
 
     template<>
-    void OperatorLRHxc<std::complex<double>, base_device::DEVICE_CPU>::grid_calculation(const int& nbands, const int& iband_dm) const
+    void OperatorLRHxc<std::complex<double>, base_device::DEVICE_CPU>::grid_calculation(const int& nbands) const
     {
         ModuleBase::TITLE("OperatorLRHxc", "grid_calculation(complex)");
         ModuleBase::timer::tick("OperatorLRHxc", "grid_calculation");
@@ -126,9 +100,9 @@ namespace LR
         hamilt::HContainer<double> HR_real_imag(GlobalC::ucell, &this->pmat);
         this->initialize_HR(HR_real_imag, ucell, gd);
 
-        auto dmR_to_hR = [&, this](const int& iband_dm, const char& type) -> void
+        auto dmR_to_hR = [&, this](const char& type) -> void
             {
-                LR_Util::get_DMR_real_imag_part(*this->DM_trans[iband_dm], DM_trans_real_imag, ucell.nat, type);
+                LR_Util::get_DMR_real_imag_part(*this->DM_trans, DM_trans_real_imag, ucell.nat, type);
                 // if (this->first_print)LR_Util::print_DMR(DM_trans_real_imag, ucell.nat, "DMR(2d, real)");
 
                 this->gint->transfer_DM2DtoGrid(DM_trans_real_imag.get_DMR_vector());
@@ -166,8 +140,8 @@ namespace LR
                 LR_Util::set_HR_real_imag_part(HR_real_imag, *this->hR, GlobalC::ucell.nat, type);
             };
         this->hR->set_zero();
-        dmR_to_hR(iband_dm, 'R');   //real
-        if (kv.get_nks() / this->nspin > 1) { dmR_to_hR(iband_dm, 'I'); }   //imag for multi-k
+        dmR_to_hR('R');   //real
+        if (kv.get_nks() / this->nspin > 1) { dmR_to_hR('I'); }   //imag for multi-k
         ModuleBase::timer::tick("OperatorLRHxc", "grid_calculation");
     }
 

source/module_lr/operator_casida/operator_lr_hxc.h

@@ -20,7 +20,7 @@ namespace LR
             const std::vector<int>& nocc,
             const std::vector<int>& nvirt,
             const psi::Psi<T, Device>& psi_ks_in,
-            std::vector<std::unique_ptr<elecstate::DensityMatrix<T, T>>>& DM_trans_in,
+            std::unique_ptr<elecstate::DensityMatrix<T, T>>& DM_trans_in,
             typename TGint<T>::type* gint_in,
             std::weak_ptr<PotHxcLR> pot_in,
             const UnitCell& ucell_in,
@@ -41,7 +41,6 @@ namespace LR
             this->is_first_node = true;
             this->hR = std::unique_ptr<hamilt::HContainer<T>>(new hamilt::HContainer<T>(&pmat_in));
             this->initialize_HR(*this->hR, ucell_in, gd_in);
-            this->DM_trans[0]->init_DMR(*this->hR);
         };
         ~OperatorLRHxc() { };
 
@@ -76,24 +75,8 @@ namespace LR
             hR.set_paraV(&pmat);
             if (std::is_same<T, double>::value) { hR.fix_gamma(); }
         }
-        template<typename TR>
-        void init_DM_trans(const int& nbands, std::vector<std::unique_ptr<elecstate::DensityMatrix<T, TR>>>& DM_trans)const
-        {
-            // LR_Util::print_DMR(*this->DM_trans[0], ucell.nat, "DMR[ib=" + std::to_string(0) + "]");
-            if (this->next_op != nullptr)
-            {
-                int prev_size = DM_trans.size();
-                if (prev_size > nbands) { for (int ib = nbands;ib < prev_size;++ib) { DM_trans[ib].reset(); } }
-                DM_trans.resize(nbands);
-                for (int ib = prev_size;ib < nbands;++ib)
-                {
-                    // the first dimenstion of DensityMatrix is nk=nks/nspin 
-                    DM_trans[ib] = LR_Util::make_unique<elecstate::DensityMatrix<T, TR>>(&this->pmat, this->nspin, this->kv.kvec_d, this->kv.get_nks() / nspin);
-                    DM_trans[ib]->init_DMR(*this->hR);
-                }
-            }
-        }
-        void grid_calculation(const int& nbands, const int& iband_dm)const;
+
+        void grid_calculation(const int& nbands)const;
 
         //global sizes
         const int& nspin;
@@ -109,7 +92,7 @@ namespace LR
         const psi::Psi<T, Device>& psi_ks = nullptr;
 
         /// transition density matrix
-        std::vector<std::unique_ptr<elecstate::DensityMatrix<T, T>>>& DM_trans;
+        std::unique_ptr<elecstate::DensityMatrix<T, T>>& DM_trans;
 
         /// transition hamiltonian in AO representation
         std::unique_ptr<hamilt::HContainer<T>> hR = nullptr;
@@ -127,8 +110,6 @@ namespace LR
         std::vector<double> orb_cutoff_;
         Grid_Driver& gd;
 
-        bool tdm_sym = false; ///< whether transition density matrix is symmetric
-
         /// test
         mutable bool first_print = true;
     };