Fix read/write HexxR at nspin=4 (#5142)

deepmodeling · Sep 20, 2024 · 75f1e6d · 75f1e6d
1 parent c4a324d
commit 75f1e6d
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/source/module_io/restart_exx_csr.hpp b/source/module_io/restart_exx_csr.hpp
@@ -40,9 +40,10 @@ namespace ModuleIO
                 const SparseMatrix<Tdata>& matrix = csr.getMatrix(iR);
                 for (auto& ijv : matrix.getElements())
                 {
-                    const int& i = ijv.first.first;
-                    const int& j = ijv.first.second;
-                    Hexxs.at(is).at(ucell.iwt2iat[i]).at({ ucell.iwt2iat[j], { R[0], R[1], R[2] } })(ucell.iwt2iw[i], ucell.iwt2iw[j]) = ijv.second;
+                    const int& npol = ucell.get_npol();
+                    const int& i = ijv.first.first * npol;
+                    const int& j = ijv.first.second * npol;
+                    Hexxs.at(is).at(ucell.iwt2iat[i]).at({ ucell.iwt2iat[j], { R[0], R[1], R[2] } })(ucell.iwt2iw[i] / npol, ucell.iwt2iw[j] / npol) = ijv.second;
                 }
             }
         }
@@ -64,17 +65,17 @@ namespace ModuleIO
                 int iat2 = a2R_data.first.first;
                 int nw1 = ucell.atoms[ucell.iat2it[iat1]].nw;
                 int nw2 = ucell.atoms[ucell.iat2it[iat2]].nw;
-                int start1 = ucell.atoms[ucell.iat2it[iat1]].stapos_wf + ucell.iat2ia[iat1] * nw1;
-                int start2 = ucell.atoms[ucell.iat2it[iat2]].stapos_wf + ucell.iat2ia[iat2] * nw2;
+                int start1 = ucell.atoms[ucell.iat2it[iat1]].stapos_wf / ucell.get_npol() + ucell.iat2ia[iat1] * nw1;
+                int start2 = ucell.atoms[ucell.iat2it[iat2]].stapos_wf / ucell.get_npol() + ucell.iat2ia[iat2] * nw2;
 
                 const TC& R = a2R_data.first.second;
                 auto& matrix = a2R_data.second;
                 Abfs::Vector3_Order<int> dR(R[0], R[1], R[2]);
                 for (int i = 0;i < nw1;++i) {
                     for (int j = 0;j < nw2;++j) {
                         target[dR][start1 + i][start2 + j] = ((std::abs(matrix(i, j)) > sparse_threshold) ? matrix(i, j) : static_cast<Tdata>(0));
-}
-}
+                    }
+                }
             }
         }
         return target;