fix add_HexxR bug in nspin=4 (#4893)

deepmodeling · Aug 7, 2024 · 52462a7 · 52462a7
1 parent 58eed06
commit 52462a7
Showing 1 changed file with 34 additions and 30 deletions.
diff --git a/source/module_ri/RI_2D_Comm.hpp b/source/module_ri/RI_2D_Comm.hpp
@@ -29,7 +29,7 @@ inline RI::Tensor<std::complex<double>> tensor_conj(const RI::Tensor<std::comple
     return r;
 }
 template<typename Tdata, typename Tmatrix>
-auto RI_2D_Comm::split_m2D_ktoR(const K_Vectors& kv, const std::vector<const Tmatrix*>& mks_2D, const Parallel_2D& pv, const int nspin)
+auto RI_2D_Comm::split_m2D_ktoR(const K_Vectors& kv, const std::vector<const Tmatrix*>& mks_2D, const Parallel_2D& pv, const int nspin)
 -> std::vector<std::map<TA,std::map<TAC,RI::Tensor<Tdata>>>>
 {
 	ModuleBase::TITLE("RI_2D_Comm","split_m2D_ktoR");
@@ -228,43 +228,47 @@ void RI_2D_Comm::add_HexxR(
 {
     ModuleBase::TITLE("RI_2D_Comm", "add_HexxR");
     ModuleBase::timer::tick("RI_2D_Comm", "add_HexxR");
-    for (const auto& Hs_tmpA : Hs[GlobalV::NSPIN == 2 ? current_spin : 0])
+    const std::map<int, std::vector<int>> is_list = { {1,{0}}, {2,{current_spin}}, {4,{0,1,2,3}} };
+    for (const int is_hs : is_list.at(GlobalV::NSPIN))
     {
-        const TA& iat0 = Hs_tmpA.first;
-        for (const auto& Hs_tmpB : Hs_tmpA.second)
+        int is0_b = 0, is1_b = 0;
+        std::tie(is0_b, is1_b) = RI_2D_Comm::split_is_block(is_hs);
+        for (const auto& Hs_tmpA : Hs[is_hs])
         {
-            const TA& iat1 = Hs_tmpB.first.first;
-            const TC& cell = Hs_tmpB.first.second;
-            const Abfs::Vector3_Order<int> R = RI_Util::array3_to_Vector3(
-                (cell_nearest ?
-                    cell_nearest->get_cell_nearest_discrete(iat0, iat1, cell)
-                    : cell));
-            hamilt::BaseMatrix<TR>* HlocR = hR.find_matrix(iat0, iat1, R.x, R.y, R.z);
-            if (HlocR == nullptr)
-            { // add R to HContainer
-                hamilt::AtomPair<TR> tmp(iat0, iat1, R.x, R.y, R.z, &pv);
-                hR.insert_pair(tmp);
-                HlocR = hR.find_matrix(iat0, iat1, R.x, R.y, R.z);
-            }
-            auto row_indexes = pv.get_indexes_row(iat0);
-            auto col_indexes = pv.get_indexes_col(iat1);
-            const RI::Tensor<Tdata>& HexxR = (Tdata)alpha * Hs_tmpB.second;
-            for (int lw0 = 0;lw0 < row_indexes.size();lw0 += npol)
-                for (int lw1 = 0;lw1 < col_indexes.size();lw1 += npol)
+            const TA& iat0 = Hs_tmpA.first;
+            for (const auto& Hs_tmpB : Hs_tmpA.second)
+            {
+                const TA& iat1 = Hs_tmpB.first.first;
+                const TC& cell = Hs_tmpB.first.second;
+                const Abfs::Vector3_Order<int> R = RI_Util::array3_to_Vector3(
+                    (cell_nearest ?
+                        cell_nearest->get_cell_nearest_discrete(iat0, iat1, cell)
+                        : cell));
+                hamilt::BaseMatrix<TR>* HlocR = hR.find_matrix(iat0, iat1, R.x, R.y, R.z);
+                if (HlocR == nullptr)
+                { // add R to HContainer
+                    hamilt::AtomPair<TR> tmp(iat0, iat1, R.x, R.y, R.z, &pv);
+                    hR.insert_pair(tmp);
+                    HlocR = hR.find_matrix(iat0, iat1, R.x, R.y, R.z);
+                }
+                auto row_indexes = pv.get_indexes_row(iat0);
+                auto col_indexes = pv.get_indexes_col(iat1);
+                const RI::Tensor<Tdata>& HexxR = (Tdata)alpha * Hs_tmpB.second;
+                for (int lw0_b = 0;lw0_b < row_indexes.size();lw0_b += npol)    // block
                 {
-                    const int& gw0 = row_indexes[lw0] / npol;
-                    const int& gw1 = col_indexes[lw1] / npol;
-                    // std::cout << "gw0=" << gw0 << ", gw1=" << gw1 << ", lw0=" << lw0 << ", lw1=" << lw1 << std::endl;
-                    HlocR->add_element(lw0, lw1, RI::Global_Func::convert<TR>(HexxR(gw0, gw1)));
-                    if (npol == 2)
+                    const int& gw0 = row_indexes[lw0_b] / npol;
+                    const int& lw0 = (npol == 2) ? (lw0_b + is0_b) : lw0_b;
+                    for (int lw1_b = 0;lw1_b < col_indexes.size();lw1_b += npol)
                     {
-                        HlocR->add_element(lw0, lw1 + 1, RI::Global_Func::convert<TR>(Hs[1].at(iat0).at({ iat1, cell })(gw0, gw1)) * alpha);
-                        HlocR->add_element(lw0 + 1, lw1, RI::Global_Func::convert<TR>(Hs[2].at(iat0).at({ iat1, cell })(gw0, gw1)) * alpha);
-                        HlocR->add_element(lw0 + 1, lw1 + 1, RI::Global_Func::convert<TR>(Hs[3].at(iat0).at({ iat1, cell })(gw0, gw1)) * alpha);
+                        const int& gw1 = col_indexes[lw1_b] / npol;
+                        const int& lw1 = (npol == 2) ? (lw1_b + is1_b) : lw1_b;
+                        HlocR->add_element(lw0, lw1, RI::Global_Func::convert<TR>(HexxR(gw0, gw1)));
                     }
                 }
+            }
         }
     }
+
     ModuleBase::timer::tick("RI_2D_Comm", "add_HexxR");
 }