Refactor class DensityMatrix. Make the code easier to read

PeizeLin · PeizeLin · commit 8b530979356a · 2026-01-06T00:14:16.000+08:00
diff --git a/source/source_estate/module_dm/density_matrix.cpp b/source/source_estate/module_dm/density_matrix.cpp
@@ -60,7 +60,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR(const int ik_in)
     using TR = double;
 
     // To check whether DMR has been initialized
-    assert(!this->_DMR.empty() && "DMR has not been initialized!");
+    assert(this->_DMR.size()==this->_nspin && "DMR has not been initialized!");
 
     ModuleBase::timer::tick("DensityMatrix", "cal_DMR");
     const int ld_hk = this->_paraV->nrow;
@@ -120,7 +120,6 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR(const int ik_in)
             for(int ik = 0; ik < this->_nk; ++ik)
             {
                 if(ik_in >= 0 && ik_in != ik) { continue; }
-
                 // copy column-major DMK to row-major DMK_mat_trans (for the purpose of computational efficiency)
                 const std::complex<double>*const DMK_mat_ptr
                     = this->_DMK[ik + ik_begin].data()
@@ -185,10 +184,10 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR(const int ik_in)
                             tmp[3] = tmp_DMR_mat[icol + step_trace[3]];
                             // transfer to Pauli matrix and save the real part
                             // save them back to the target_mat
-                            target_DMR_mat[icol + step_trace[0]] = tmp[0].real() + tmp[3].real();  // (rho_upup + rho_downdown).real()
-                            target_DMR_mat[icol + step_trace[1]] = tmp[1].real() + tmp[2].real();  // (rho_updown + rho_downup).real()
-                            target_DMR_mat[icol + step_trace[2]] = -tmp[1].imag() + tmp[2].imag(); // (i * (rho_updown - rho_downup)).real()
-                            target_DMR_mat[icol + step_trace[3]] = tmp[0].real() - tmp[3].real();  // (rho_upup - rho_downdown).real()
+                            target_DMR_mat[icol + step_trace[0]] = tmp[0].real() + tmp[3].real();  // rho_0 = (rho_upup + rho_downdown).real()
+                            target_DMR_mat[icol + step_trace[1]] = tmp[1].real() + tmp[2].real();  // rho_x = (rho_updown + rho_downup).real()
+                            target_DMR_mat[icol + step_trace[2]] = -tmp[1].imag() + tmp[2].imag(); // rho_y = (i * (rho_updown - rho_downup)).real()
+                            target_DMR_mat[icol + step_trace[3]] = tmp[0].real() - tmp[3].real();  // rho_z = (rho_upup - rho_downdown).real()
                         }
                         tmp_DMR_mat += col_size * 2;
                         target_DMR_mat += col_size * 2;
@@ -207,7 +206,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_td(const UnitCell& uce
     ModuleBase::TITLE("DensityMatrix", "cal_DMR_td");
     using TR = double;
     // To check whether DMR has been initialized
-    assert(!this->_DMR.empty() && "DMR has not been initialized!");
+    assert(this->_DMR.size()==this->_nspin && "DMR has not been initialized!");
 
     ModuleBase::timer::tick("DensityMatrix", "cal_DMR_td");
     const int ld_hk = this->_paraV->nrow;
@@ -249,10 +248,9 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_td(const UnitCell& uce
                 }
                 #endif
                 target_DMR_mat_vec[iR] = target_mat->get_pointer();
-                double arg_td = 0.0;
                 //cal tddft phase for hybrid gauge
-                ModuleBase::Vector3<double> dtau = ucell.cal_dtau(iat1, iat2, R_index);
-                arg_td = At * dtau * ucell.lat0;
+                const ModuleBase::Vector3<double> dtau = ucell.cal_dtau(iat1, iat2, R_index);
+                const double arg_td = At * dtau * ucell.lat0;
                 for(int ik = 0; ik < this->_nk; ++ik)
                 {
                     if(ik_in >= 0 && ik_in != ik) { continue; }
@@ -271,7 +269,6 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_td(const UnitCell& uce
             for(int ik = 0; ik < this->_nk; ++ik)
             {
                 if(ik_in >= 0 && ik_in != ik) { continue; }
-
                 // copy column-major DMK to row-major DMK_mat_trans (for the purpose of computational efficiency)
                 const std::complex<double>*const DMK_mat_ptr
                     = this->_DMK[ik + ik_begin].data()
@@ -336,10 +333,10 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_td(const UnitCell& uce
                             tmp[3] = tmp_DMR_mat[icol + step_trace[3]];
                             // transfer to Pauli matrix and save the real part
                             // save them back to the target_mat
-                            target_DMR_mat[icol + step_trace[0]] = tmp[0].real() + tmp[3].real();  // (rho_upup + rho_downdown).real()
-                            target_DMR_mat[icol + step_trace[1]] = tmp[1].real() + tmp[2].real();  // (rho_updown + rho_downup).real()
-                            target_DMR_mat[icol + step_trace[2]] = -tmp[1].imag() + tmp[2].imag(); // (i * (rho_updown - rho_downup)).real()
-                            target_DMR_mat[icol + step_trace[3]] = tmp[0].real() - tmp[3].real();  // (rho_upup - rho_downdown).real()
+                            target_DMR_mat[icol + step_trace[0]] = tmp[0].real() + tmp[3].real();  // rho_0 = (rho_upup + rho_downdown).real()
+                            target_DMR_mat[icol + step_trace[1]] = tmp[1].real() + tmp[2].real();  // rho_x = (rho_updown + rho_downup).real()
+                            target_DMR_mat[icol + step_trace[2]] = -tmp[1].imag() + tmp[2].imag(); // rho_y = (i * (rho_updown - rho_downup)).real()
+                            target_DMR_mat[icol + step_trace[3]] = tmp[0].real() - tmp[3].real();  // rho_z = (rho_upup - rho_downdown).real()
                         }
                         tmp_DMR_mat += col_size * 2;
                         target_DMR_mat += col_size * 2;
@@ -353,23 +350,26 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_td(const UnitCell& uce
 
 // calculate DMR from DMK using blas for multi-k calculation
 template <>
-void DensityMatrix<double, double>::cal_DMR_full(hamilt::HContainer<std::complex<double>>* dmR_out)const{}
+void DensityMatrix<double, double>::cal_DMR_full(hamilt::HContainer<std::complex<double>>* dmR_out, const int ik_in)const{}
 template <>
-void DensityMatrix<std::complex<double>, double>::cal_DMR_full(hamilt::HContainer<std::complex<double>>* dmR_out)const
+void DensityMatrix<std::complex<double>, double>::cal_DMR_full(
+    hamilt::HContainer<std::complex<double>>* dmR_out,
+    const int ik_in) const
 {
     ModuleBase::TITLE("DensityMatrix", "cal_DMR_full");
+    using TR = std::complex<double>;
 
     ModuleBase::timer::tick("DensityMatrix", "cal_DMR_full");
     const int ld_hk = this->_paraV->nrow;
-    hamilt::HContainer<std::complex<double>>* target_DMR = dmR_out;
+    hamilt::HContainer<TR>* target_DMR = dmR_out;
     // set zero since this function is called in every scf step
     target_DMR->set_zero();
     #ifdef _OPENMP
     #pragma omp parallel for schedule(dynamic)
     #endif
     for (int i = 0; i < target_DMR->size_atom_pairs(); ++i)
     {
-        hamilt::AtomPair<std::complex<double>>& target_ap = target_DMR->get_atom_pair(i);
+        hamilt::AtomPair<TR>& target_ap = target_DMR->get_atom_pair(i);
         const int iat1 = target_ap.get_atom_i();
         const int iat2 = target_ap.get_atom_j();
         // get global indexes of whole matrix for each atom in this process
@@ -379,14 +379,15 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_full(hamilt::HContaine
         const int col_size = this->_paraV->get_col_size(iat2);
         const int mat_size = row_size * col_size;
         const int R_size = target_ap.get_R_size();
+        assert(row_ap != -1 && col_ap != -1 && "Atom-pair not belong this process");
 
         // calculate kphase and target_mat_ptr
         std::vector<std::vector<std::complex<double>>> kphase_vec(this->_nk, std::vector<std::complex<double>>(R_size));
-        std::vector<std::complex<double>*> target_DMR_mat_vec(R_size);
+        std::vector<TR*> target_DMR_mat_vec(R_size);
         for(int iR = 0; iR < R_size; ++iR)
         {
             const ModuleBase::Vector3<int> R_index = target_ap.get_R_index(iR);
-            hamilt::BaseMatrix<std::complex<double>>*const target_mat = target_ap.find_matrix(R_index);
+            hamilt::BaseMatrix<TR>*const target_mat = target_ap.find_matrix(R_index);
             #ifdef __DEBUG
             if (target_mat == nullptr)
             {
@@ -397,6 +398,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_full(hamilt::HContaine
             target_DMR_mat_vec[iR] = target_mat->get_pointer();
             for(int ik = 0; ik < this->_nk; ++ik)
             {
+                if(ik_in >= 0 && ik_in != ik) { continue; }
                 // cal k_phase
                 // if TK==std::complex<double>, kphase is e^{ikR}
                 const ModuleBase::Vector3<double> dR(R_index[0], R_index[1], R_index[2]);
@@ -410,6 +412,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_full(hamilt::HContaine
         std::vector<std::complex<double>> DMK_mat_trans(mat_size);
         for(int ik = 0; ik < this->_nk; ++ik)
         {
+            if(ik_in >= 0 && ik_in != ik) { continue; }
             // copy column-major DMK to row-major DMK_mat_trans (for the purpose of computational efficiency)
             const std::complex<double>*const DMK_mat_ptr
                 = this->_DMK[ik].data()
@@ -422,12 +425,11 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR_full(hamilt::HContaine
             for(int iR = 0; iR < R_size; ++iR)
             {
                 const std::complex<double> kphase = kphase_vec[ik][iR];
-                std::complex<double>* target_DMR_mat = target_DMR_mat_vec[iR];
                 BlasConnector::axpy(mat_size,
                                     kphase,
                                     DMK_mat_trans.data(),
                                     1,
-                                    target_DMR_mat,
+                                    target_DMR_mat_vec[iR],
                                     1);
             }
         }
diff --git a/source/source_estate/module_dm/density_matrix.h b/source/source_estate/module_dm/density_matrix.h
@@ -179,13 +179,15 @@ class DensityMatrix
 
     /**
      * @brief calculate density matrix DMR from dm(k) using blas::axpy
+     * @param ik_in
      * if ik_in < 0, calculate all k-points
      * if ik_in >= 0, calculate only one k-point without summing over k-points
      */
     void cal_DMR(const int ik_in = -1);
 
     /**
      * @brief calculate density matrix DMR with additional vector potential phase, used for hybrid gauge tddft
+     * @param ik_in
      * if ik_in < 0, calculate all k-points
      * if ik_in >= 0, calculate only one k-point without summing over k-points
      */
@@ -195,8 +197,11 @@ class DensityMatrix
      * @brief calculate complex density matrix DMR with both real and imaginary part for noncollinear-spin calculation
      * the stored dm(k) has been used to calculate the passin DMR
      * @param dmR_out pointer of HContainer object to store the calculated complex DMR
+     * @param ik_in
+     * if ik_in < 0, calculate all k-points
+     * if ik_in >= 0, calculate only one k-point without summing over k-points
      */
-    void cal_DMR_full(hamilt::HContainer<std::complex<double>>* dmR_out) const;
+    void cal_DMR_full(hamilt::HContainer<std::complex<double>>* dmR_out, const int ik_in = -1) const;
 
     /**
      * @brief (Only nspin=2) switch DMR to total density matrix or magnetization density matrix
