diff --git a/source/source_lcao/module_deepks/deepks_pdm.cpp b/source/source_lcao/module_deepks/deepks_pdm.cpp index cad6683a72e..5e2186420dd 100644 --- a/source/source_lcao/module_deepks/deepks_pdm.cpp +++ b/source/source_lcao/module_deepks/deepks_pdm.cpp @@ -147,9 +147,8 @@ void DeePKS_domain::update_dmr(const std::vector>& k } calculated_pairs.push_back(std::make_tuple(ibt1, ibt2, dR.x, dR.y, dR.z)); - dm_pair.find_R(dR); - hamilt::BaseMatrix* dmr_ptr = dm_pair.find_matrix(dR); - dmr_ptr->set_zero(); // must reset to zero to avoid accumulation! + const int r_index = dm_pair.find_R(dR); + dm_pair.get_HR_values(r_index).set_zero(); // must reset to zero to avoid accumulation! for (int ik = 0; ik < dmk.size(); ik++) { @@ -166,11 +165,11 @@ void DeePKS_domain::update_dmr(const std::vector>& k TK* kphase_ptr = reinterpret_cast(&kphase); if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) { - dm_pair.add_from_matrix(dmk[ik].data(), pv.get_row_size(), *kphase_ptr, 1); + dm_pair.add_from_matrix(r_index, dmk[ik].data(), pv.get_row_size(), *kphase_ptr, 1); } else { - dm_pair.add_from_matrix(dmk[ik].data(), pv.get_col_size(), *kphase_ptr, 0); + dm_pair.add_from_matrix(r_index, dmk[ik].data(), pv.get_col_size(), *kphase_ptr, 0); } } } diff --git a/source/source_lcao/module_gint/gint_common.cpp b/source/source_lcao/module_gint/gint_common.cpp index 987fcbe5a0e..07aa4d78bd7 100644 --- a/source/source_lcao/module_gint/gint_common.cpp +++ b/source/source_lcao/module_gint/gint_common.cpp @@ -103,7 +103,7 @@ void compose_hr_gint(HContainer& hr_gint) { for (int icol = 0; icol < upper_mat->get_col_size(); ++icol) { - lower_mat->get_value(icol, irow) = upper_ap->get_value(irow, icol); + lower_mat->get_value(icol, irow) = upper_mat->get_value(irow, icol); } } } diff --git a/source/source_lcao/module_hcontainer/atom_pair.cpp b/source/source_lcao/module_hcontainer/atom_pair.cpp index d776c33e3cf..6877299297d 100644 --- a/source/source_lcao/module_hcontainer/atom_pair.cpp +++ b/source/source_lcao/module_hcontainer/atom_pair.cpp @@ -32,7 +32,6 @@ AtomPair::AtomPair(const int& atom_i_, const int& atom_j_, const Parallel_Orb this->col_size = this->paraV->get_col_size(atom_j); this->R_index.resize(0); this->R_index.push_back(ModuleBase::Vector3(0, 0, 0)); - this->current_R = 0; if (existed_matrix != nullptr) { BaseMatrix tmp(row_size, col_size, existed_matrix); @@ -66,7 +65,6 @@ AtomPair::AtomPair(const int& atom_i_, this->col_size = this->paraV->get_col_size(atom_j); this->R_index.resize(0); this->R_index.push_back(ModuleBase::Vector3(rx, ry, rz)); - this->current_R = 0; if (existed_matrix != nullptr) { BaseMatrix tmp(row_size, col_size, existed_matrix); @@ -98,7 +96,6 @@ AtomPair::AtomPair(const int& atom_i_, this->col_size = this->paraV->get_col_size(atom_j); this->R_index.resize(0); this->R_index.push_back(ModuleBase::Vector3(R_index)); - this->current_R = 0; if (existed_matrix != nullptr) { BaseMatrix tmp(row_size, col_size, existed_matrix); @@ -128,7 +125,6 @@ AtomPair::AtomPair(const int& atom_i_, this->col_size = col_atom_begin[atom_j + 1] - col_atom_begin[atom_j]; this->R_index.resize(0); this->R_index.push_back(ModuleBase::Vector3(0, 0, 0)); - this->current_R = 0; if (existed_matrix != nullptr) { BaseMatrix tmp(row_size, col_size, existed_matrix); @@ -160,7 +156,6 @@ AtomPair::AtomPair(const int& atom_i_, this->col_size = col_atom_begin[atom_j + 1] - col_atom_begin[atom_j]; this->R_index.resize(0); this->R_index.push_back(ModuleBase::Vector3(rx, ry, rz)); - this->current_R = 0; if (existed_matrix != nullptr) { BaseMatrix tmp(row_size, col_size, existed_matrix); @@ -183,7 +178,6 @@ template AtomPair::AtomPair(const AtomPair& other, T* data_pointer) : R_index(other.R_index), paraV(other.paraV), - current_R(other.current_R), atom_i(other.atom_i), atom_j(other.atom_j), row_ap(other.row_ap), @@ -247,7 +241,6 @@ AtomPair& AtomPair::operator=(const AtomPair& other) R_index = other.R_index; values = other.values; paraV = other.paraV; - current_R = other.current_R; atom_i = other.atom_i; atom_j = other.atom_j; row_ap = other.row_ap; @@ -264,7 +257,6 @@ AtomPair::AtomPair(AtomPair&& other) noexcept : R_index(std::move(other.R_index)), values(std::move(other.values)), paraV(other.paraV), - current_R(other.current_R), atom_i(other.atom_i), atom_j(other.atom_j), row_ap(other.row_ap), @@ -285,7 +277,6 @@ AtomPair& AtomPair::operator=(AtomPair&& other) noexcept values = std::move(other.values); paraV = other.paraV; other.paraV = nullptr; - current_R = other.current_R; atom_i = other.atom_i; atom_j = other.atom_j; row_ap = other.row_ap; @@ -388,6 +379,33 @@ bool AtomPair::identify(const int& atom_i_, const int& atom_j_) const } } +// Helper function to find R index in R_index array +template +int AtomPair::find_R(const int& rx_in, const int& ry_in, const int& rz_in) const +{ + for (int i = 0; i < this->R_index.size(); i++) + { + if (R_index[i].x == rx_in && R_index[i].y == ry_in && R_index[i].z == rz_in) + { + return i; + } + } + return -1; +} + +template +int AtomPair::find_R(const ModuleBase::Vector3& R_in) const +{ + for (int i = 0; i < this->R_index.size(); i++) + { + if (R_index[i] == R_in) + { + return i; + } + } + return -1; +} + // get_HR_values for no-const AtomPair template BaseMatrix& AtomPair::get_HR_values(int rx_in, int ry_in, int rz_in) @@ -396,7 +414,7 @@ BaseMatrix& AtomPair::get_HR_values(int rx_in, int ry_in, int rz_in) const int r_index = this->find_R(rx_in, ry_in, rz_in); if (r_index != -1) { - // if found, return this->values[current_R] + // if found, return this->values[r_index] return this->values[r_index]; } // if not found, add a new BaseMatrix for this R index @@ -411,11 +429,11 @@ BaseMatrix& AtomPair::get_HR_values(int rx_in, int ry_in, int rz_in) template const BaseMatrix& AtomPair::get_HR_values(int rx_in, int ry_in, int rz_in) const { - // if current_R is -1, R index has not been fixed, find existed R index + // find existed R index const int r_index = this->find_R(rx_in, ry_in, rz_in); if (r_index != -1) { - // if found, return this->values[current_R] + // if found, return this->values[r_index] return this->values[r_index]; } // if not found, throw a error message @@ -436,36 +454,6 @@ BaseMatrix& AtomPair::get_HR_values(const int& index) const return const_cast&>(this->values[index]); } -// find_R -template -int AtomPair::find_R(const int& rx_in, const int& ry_in, const int& rz_in) const -{ - for (int i = 0; i < this->R_index.size(); i++) - { - if (R_index[i].x == rx_in && R_index[i].y == ry_in && R_index[i].z == rz_in) - { - this->current_R = i; - return i; - } - } - return (-1); -} - -// find_R -template -int AtomPair::find_R(const ModuleBase::Vector3& R_in) const -{ - for (int i = 0; i < this->R_index.size(); i++) - { - if (R_index[i] == R_in) - { - this->current_R = i; - return i; - } - } - return (-1); -} - // find_matrix template const BaseMatrix* AtomPair::find_matrix(const int& rx_in, const int& ry_in, const int& rz_in) const @@ -551,21 +539,22 @@ void AtomPair::merge(const AtomPair& other, bool skip_R) rz = other.R_index[i].z; } const BaseMatrix& matrix_tmp = other.get_HR_values(i); + const int r_index = this->find_R(rx, ry, rz); //if not found, push_back this BaseMatrix to this->values - if (this->find_R(rx, ry, rz) == -1) + if (r_index == -1) { this->R_index.push_back(ModuleBase::Vector3(rx, ry, rz)); this->values.push_back(matrix_tmp); } //if found but not allocated, skip this BaseMatrix values - else if (this->values[current_R].get_pointer() == nullptr || matrix_tmp.get_pointer() == nullptr) + else if (this->values[r_index].get_pointer() == nullptr || matrix_tmp.get_pointer() == nullptr) { continue; } // if found and allocated, add data else { - this->values[current_R].add_array(matrix_tmp.get_pointer()); + this->values[r_index].add_array(matrix_tmp.get_pointer()); } } } @@ -595,21 +584,20 @@ void AtomPair::merge_to_gamma() } this->values.clear(); this->values.push_back(tmp); - - this->current_R = 0; } template -void AtomPair::add_to_matrix(std::complex* hk, +void AtomPair::add_from_matrix(const std::complex* hk, const int ld_hk, const std::complex& kphase, - const int hk_type) const + const int hk_type) { - const BaseMatrix& matrix = values[current_R]; + // Only works for R_index[0], which is (0,0,0) + const BaseMatrix& matrix = values[0]; T* hr_tmp = matrix.get_pointer(); - std::complex* hk_tmp = hk; - T* hk_real_pointer = nullptr; - T* hk_imag_pointer = nullptr; + const std::complex* hk_tmp = hk; + const T* hk_real_pointer = nullptr; + const T* hk_imag_pointer = nullptr; const int ld_hk_2 = ld_hk * 2; // row major if (hk_type == 0) @@ -619,8 +607,8 @@ void AtomPair::add_to_matrix(std::complex* hk, { hk_real_pointer = (T*)hk_tmp; hk_imag_pointer = hk_real_pointer+1; - BlasConnector::axpy(this->col_size, kphase.real(), hr_tmp, 1, hk_real_pointer, 2); - BlasConnector::axpy(this->col_size, kphase.imag(), hr_tmp, 1, hk_imag_pointer, 2); + BlasConnector::axpy(this->col_size, kphase.real(), hk_real_pointer, 2, hr_tmp, 1); + BlasConnector::axpy(this->col_size, -kphase.imag(), hk_imag_pointer, 2, hr_tmp, 1); hk_tmp += ld_hk; hr_tmp += this->col_size; } @@ -633,32 +621,29 @@ void AtomPair::add_to_matrix(std::complex* hk, { hk_real_pointer = (T*)hk_tmp; hk_imag_pointer = hk_real_pointer+1; - BlasConnector::axpy(this->col_size, kphase.real(), hr_tmp, 1, hk_real_pointer, ld_hk_2); - BlasConnector::axpy(this->col_size, kphase.imag(), hr_tmp, 1, hk_imag_pointer, ld_hk_2); + BlasConnector::axpy(this->col_size, kphase.real(), hk_real_pointer, ld_hk_2, hr_tmp, 1); + BlasConnector::axpy(this->col_size, -kphase.imag(), hk_imag_pointer, ld_hk_2, hr_tmp, 1); hk_tmp ++; hr_tmp += this->col_size; } } } -// add_to_matrix +// add_from_matrix template -void AtomPair::add_to_matrix(T* hk, const int ld_hk, const T& kphase, const int hk_type) const +void AtomPair::add_from_matrix(const T* hk, const int ld_hk, const T& kphase, const int hk_type) { - const BaseMatrix& matrix = values[current_R]; + // Only works for R_index[0], which is (0,0,0) + const BaseMatrix& matrix = values[0]; T* hr_tmp = matrix.get_pointer(); - T* hk_tmp = hk; + const T* hk_tmp = hk; // row major if (hk_type == 0) { hk_tmp += this->row_ap * ld_hk + this->col_ap; for (int mu = 0; mu < this->row_size; mu++) { - BlasConnector::axpy(this->col_size, kphase, hr_tmp, 1, hk_tmp, 1); - /*for (int nu = 0; nu < this->col_size; nu++) - { - hk_tmp[nu] += matrix.get_value(mu, nu) * kphase; - }*/ + BlasConnector::axpy(this->col_size, kphase, hk_tmp, 1, hr_tmp, 1); hk_tmp += ld_hk; hr_tmp += this->col_size; } @@ -669,24 +654,22 @@ void AtomPair::add_to_matrix(T* hk, const int ld_hk, const T& kphase, const i hk_tmp += this->col_ap * ld_hk + this->row_ap; for (int mu = 0; mu < this->row_size; mu++) { - BlasConnector::axpy(this->col_size, kphase, hr_tmp, 1, hk_tmp, ld_hk); - /*for (int mu = 0; mu < this->row_size; mu++) - { - hk_tmp[mu] += matrix.get_value(mu, nu) * kphase; - }*/ + BlasConnector::axpy(this->col_size, kphase, hk_tmp, ld_hk, hr_tmp, 1); ++hk_tmp; hr_tmp += this->col_size; } } } +// add_from_matrix with explicit R_index - thread-safe version template -void AtomPair::add_from_matrix(const std::complex* hk, - const int ld_hk, - const std::complex& kphase, - const int hk_type) +void AtomPair::add_from_matrix(const int R_index, + const std::complex* hk, + const int ld_hk, + const std::complex& kphase, + const int hk_type) { - const BaseMatrix& matrix = values[current_R]; + const BaseMatrix& matrix = values[R_index]; T* hr_tmp = matrix.get_pointer(); const std::complex* hk_tmp = hk; const T* hk_real_pointer = nullptr; @@ -722,11 +705,15 @@ void AtomPair::add_from_matrix(const std::complex* hk, } } -// add_to_matrix +// add_from_matrix with explicit R_index - thread-safe version template -void AtomPair::add_from_matrix(const T* hk, const int ld_hk, const T& kphase, const int hk_type) +void AtomPair::add_from_matrix(const int R_index, + const T* hk, + const int ld_hk, + const T& kphase, + const int hk_type) { - const BaseMatrix& matrix = values[current_R]; + const BaseMatrix& matrix = values[R_index]; T* hr_tmp = matrix.get_pointer(); const T* hk_tmp = hk; // row major @@ -736,10 +723,6 @@ void AtomPair::add_from_matrix(const T* hk, const int ld_hk, const T& kphase, for (int mu = 0; mu < this->row_size; mu++) { BlasConnector::axpy(this->col_size, kphase, hk_tmp, 1, hr_tmp, 1); - /*for (int nu = 0; nu < this->col_size; nu++) - { - hk_tmp[nu] += matrix.get_value(mu, nu) * kphase; - }*/ hk_tmp += ld_hk; hr_tmp += this->col_size; } @@ -751,32 +734,107 @@ void AtomPair::add_from_matrix(const T* hk, const int ld_hk, const T& kphase, for (int mu = 0; mu < this->row_size; mu++) { BlasConnector::axpy(this->col_size, kphase, hk_tmp, ld_hk, hr_tmp, 1); - /*for (int mu = 0; mu < this->row_size; mu++) - { - hk_tmp[mu] += matrix.get_value(mu, nu) * kphase; - }*/ ++hk_tmp; hr_tmp += this->col_size; } } } -// add_to_array +// add_to_matrix with explicit R_index - thread-safe version template -void AtomPair::add_to_array(T* array, const T& kphase) const +void AtomPair::add_to_matrix(const int R_index, + std::complex* hk, + const int ld_hk, + const std::complex& kphase, + const int hk_type) const { - const BaseMatrix& matrix = values[current_R]; + const BaseMatrix& matrix = values[R_index]; + T* hr_tmp = matrix.get_pointer(); + std::complex* hk_tmp = hk; + T* hk_real_pointer = nullptr; + T* hk_imag_pointer = nullptr; + const int ld_hk_2 = ld_hk * 2; + // row major + if (hk_type == 0) + { + hk_tmp += this->row_ap * ld_hk + this->col_ap; + for (int mu = 0; mu < this->row_size; mu++) + { + hk_real_pointer = (T*)hk_tmp; + hk_imag_pointer = hk_real_pointer+1; + BlasConnector::axpy(this->col_size, kphase.real(), hr_tmp, 1, hk_real_pointer, 2); + BlasConnector::axpy(this->col_size, kphase.imag(), hr_tmp, 1, hk_imag_pointer, 2); + hk_tmp += ld_hk; + hr_tmp += this->col_size; + } + } + // column major + else if (hk_type == 1) + { + hk_tmp += this->col_ap * ld_hk + this->row_ap; + for (int mu = 0; mu < this->row_size; mu++) + { + hk_real_pointer = (T*)hk_tmp; + hk_imag_pointer = hk_real_pointer+1; + BlasConnector::axpy(this->col_size, kphase.real(), hr_tmp, 1, hk_real_pointer, ld_hk_2); + BlasConnector::axpy(this->col_size, kphase.imag(), hr_tmp, 1, hk_imag_pointer, ld_hk_2); + hk_tmp ++; + hr_tmp += this->col_size; + } + } +} + +// add_to_matrix with explicit R_index - thread-safe version +template +void AtomPair::add_to_matrix(const int R_index, + T* hk, + const int ld_hk, + const T& kphase, + const int hk_type) const +{ + const BaseMatrix& matrix = values[R_index]; + T* hr_tmp = matrix.get_pointer(); + T* hk_tmp = hk; + // row major + if (hk_type == 0) + { + hk_tmp += this->row_ap * ld_hk + this->col_ap; + for (int mu = 0; mu < this->row_size; mu++) + { + BlasConnector::axpy(this->col_size, kphase, hr_tmp, 1, hk_tmp, 1); + hk_tmp += ld_hk; + hr_tmp += this->col_size; + } + } + // column major + else if (hk_type == 1) + { + hk_tmp += this->col_ap * ld_hk + this->row_ap; + for (int mu = 0; mu < this->row_size; mu++) + { + BlasConnector::axpy(this->col_size, kphase, hr_tmp, 1, hk_tmp, ld_hk); + ++hk_tmp; + hr_tmp += this->col_size; + } + } +} + +// add_to_array with explicit R_index - thread-safe version +template +void AtomPair::add_to_array(const int R_index, T* array, const T& kphase) const +{ + const BaseMatrix& matrix = values[R_index]; for (int i = 0; i < this->row_size * this->col_size; i++) { array[i] += matrix.get_pointer()[i] * kphase; } } -// add_to_array +// add_to_array with explicit R_index - thread-safe version template -void AtomPair::add_to_array(std::complex* array, const std::complex& kphase) const +void AtomPair::add_to_array(const int R_index, std::complex* array, const std::complex& kphase) const { - const BaseMatrix& matrix = values[current_R]; + const BaseMatrix& matrix = values[R_index]; for (int i = 0; i < this->row_size * this->col_size; i++) { array[i] += matrix.get_pointer()[i] * kphase; @@ -786,7 +844,6 @@ void AtomPair::add_to_array(std::complex* array, const std::complex& kp template std::tuple, T*> AtomPair::get_matrix_values(int ir) const { - if(ir<0) ir = this->current_R; return std::tuple, T*>({this->row_ap, this->row_size, this->col_ap, this->col_size}, this->values[ir].get_pointer()); } @@ -806,50 +863,39 @@ ModuleBase::Vector3 AtomPair::get_R_index(const int& index) const } } +// get_value with R_index template -ModuleBase::Vector3 AtomPair::get_R_index() const -{ - // return the ModuleBase::Vector3 of R_index[index] - return R_index[current_R]; -} - -// get_value -template -T& AtomPair::get_value(const int& i) const +T& AtomPair::get_value(const int R_index, const int& i) const { #ifdef __DEBUG assert(i < this->row_size * this->col_size); assert(i >= 0); - assert(current_R < this->values.size()); - assert(current_R >= 0); + assert(R_index < this->values.size()); + assert(R_index >= 0); #endif - return this->values[current_R].get_pointer()[i]; + return this->values[R_index].get_pointer()[i]; } -// get_value +// get_value with R_index template -T& AtomPair::get_value(const int& row, const int& col) const +T& AtomPair::get_value(const int R_index, const int& row, const int& col) const { #ifdef __DEBUG assert(row < this->row_size && row >= 0); assert(col < this->col_size && col >= 0); - assert(current_R < this->values.size()); - assert(current_R >= 0); + assert(R_index < this->values.size()); + assert(R_index >= 0); #endif - return this->values[current_R].get_value(row, col); + return this->values[R_index].get_value(row, col); } // get_pointer template T* AtomPair::get_pointer(int ir) const { - if(ir<0) - { - ir = current_R; - } #ifdef __DEBUG - assert(current_R < this->values.size()); - assert(current_R >= 0); + assert(ir < this->values.size()); + assert(ir >= 0); #endif return this->values[ir].get_pointer(); } diff --git a/source/source_lcao/module_hcontainer/atom_pair.h b/source/source_lcao/module_hcontainer/atom_pair.h index beb45247d33..20f76f9b33c 100644 --- a/source/source_lcao/module_hcontainer/atom_pair.h +++ b/source/source_lcao/module_hcontainer/atom_pair.h @@ -189,22 +189,35 @@ class AtomPair // interface for get (rx, ry, rz) of index-th R-index in this->R_index, the return should be ModuleBase::Vector3 ModuleBase::Vector3 get_R_index(const int& index) const; - // interface for get (rx, ry, rz) of current_R, the return should be ModuleBase::Vector3 - ModuleBase::Vector3 get_R_index() const; - // interface for search (rx, ry, rz) in this->R_index, if found, current_R would be set to index + // interface for search (rx, ry, rz) in this->R_index, if found, return index, else return -1 int find_R(const int& rx_in, const int& ry_in, const int& rz_in) const; int find_R(const ModuleBase::Vector3& R_in) const; - // interface for search (rx, ry, rz) in this->R_index, if found, current_R would be set to index - // and return BaseMatrix* of this->values[index] + // interface for search (rx, ry, rz) in this->R_index, if found, return BaseMatrix* of this->values[index] const BaseMatrix* find_matrix(const int& rx_in, const int& ry_in, const int& rz_in) const; BaseMatrix* find_matrix(const int& rx_in, const int& ry_in, const int& rz_in); const BaseMatrix* find_matrix(const ModuleBase::Vector3& R_in) const; BaseMatrix* find_matrix(const ModuleBase::Vector3& R_in); - // this interface will call get_value in this->values - // these four interface can be used only when R-index has been choosed (current_R >= 0) - T& get_value(const int& i) const; - T& get_value(const int& row, const int& col) const; + /** + * @brief Get value at position i in the matrix values[R_index]. + * Thread-safe version: explicitly pass R_index. + * + * @param R_index Index of the R vector in this->values + * @param i Position in the flattened matrix (row * col_size + col) + * @return T& Reference to the value + */ + T& get_value(const int R_index, const int& i) const; + + /** + * @brief Get value at position (row, col) in the matrix values[R_index]. + * Thread-safe version: explicitly pass R_index. + * + * @param R_index Index of the R vector in this->values + * @param row Row index + * @param col Column index + * @return T& Reference to the value + */ + T& get_value(const int R_index, const int& row, const int& col) const; /** * @brief get values of this->values[ir] for a whole matrix @@ -213,12 +226,13 @@ class AtomPair * std::vector(4) contains (row_begin_index, row_size, col_begin_index, col_size) * T* is pointer of values[ir].value_begin, legal index is [0, row_size*col_size) */ - std::tuple, T*> get_matrix_values(int ir = -1) const; + std::tuple, T*> get_matrix_values(int ir) const; /** * @brief get pointer of value from a submatrix + * @param ir index of this->values */ - T* get_pointer(int ir=-1) const; + T* get_pointer(int ir) const; // add another BaseMatrix to this->values with specific R index. void convert_add(const BaseMatrix& target, int rx_in, int ry_in, int rz_in); @@ -236,27 +250,40 @@ class AtomPair void merge_to_gamma(); /** - * @brief Add this->value[current_R] * kphase as a block matrix of hk. + * @brief Add this->value[R_index] * kphase as a block matrix of hk. + * Thread-safe version: explicitly pass R_index. * - * For row major dense matrix (hk_type == 0): value[current_R][i*col_size+j] -> hk[(row_ap+i) * ld_hk + col_ap + j] - * For column major dense matrix (hk_type == 1): value[current_R][i*col_size+j] -> hk[row_ap + i + (col_ap+j) * - * ld_hk] For sparse matrix (hk_type == 2): not implemented yet + * For row major dense matrix (hk_type == 0): value[R_index][i*col_size+j] -> hk[(row_ap+i) * ld_hk + col_ap + j] + * For column major dense matrix (hk_type == 1): value[R_index][i*col_size+j] -> hk[row_ap + i + (col_ap+j) * ld_hk] * + * @param R_index Index of the R vector in this->values * @param hk Pointer to the target matrix. * @param ld_hk Leading dimension of the target matrix. * @param kphase Complex scalar to be multiplied with the block matrix. * @param hk_type The type of matrix layout (default: 0). */ - void add_to_matrix(std::complex* hk, + void add_to_matrix(const int R_index, + std::complex* hk, const int ld_hk, const std::complex& kphase, const int hk_type = 0) const; /** - * @brief Add this->value[current_R] * kphase as a block matrix of hk. + * @brief Add this->value[R_index] * kphase as a block matrix of hk. + * Thread-safe version: explicitly pass R_index. * for non-collinear spin case only + * + * @param R_index Index of the R vector in this->values + * @param hk Pointer to the target matrix. + * @param ld_hk Leading dimension of the target matrix. + * @param kphase Scalar to be multiplied with the block matrix. + * @param hk_type The type of matrix layout (default: 0). */ - void add_to_matrix(T* hk, const int ld_hk, const T& kphase, const int hk_type = 0) const; + void add_to_matrix(const int R_index, + T* hk, + const int ld_hk, + const T& kphase, + const int hk_type = 0) const; void add_from_matrix(const std::complex* hk, const int ld_hk, @@ -266,15 +293,58 @@ class AtomPair void add_from_matrix(const T* hk, const int ld_hk, const T& kphase, const int hk_type = 0); /** - * @brief Add this->value[current_R] * kphase to an array. + * @brief Add data from hk to this->value[R_index]. + * Thread-safe version: explicitly pass R_index. + * + * @param R_index Index of the R vector in this->values + * @param hk Pointer to the source matrix. + * @param ld_hk Leading dimension of the source matrix. + * @param kphase Complex scalar to be multiplied with the source matrix. + * @param hk_type The type of matrix layout (default: 0). + */ + void add_from_matrix(const int R_index, + const std::complex* hk, + const int ld_hk, + const std::complex& kphase, + const int hk_type = 0); + + /** + * @brief Add data from hk to this->value[R_index]. + * Thread-safe version: explicitly pass R_index. + * for non-collinear spin case only + * + * @param R_index Index of the R vector in this->values + * @param hk Pointer to the source matrix. + * @param ld_hk Leading dimension of the source matrix. + * @param kphase Scalar to be multiplied with the source matrix. + * @param hk_type The type of matrix layout (default: 0). + */ + void add_from_matrix(const int R_index, + const T* hk, + const int ld_hk, + const T& kphase, + const int hk_type = 0); + + /** + * @brief Add this->value[R_index] * kphase to an array. + * Thread-safe version: explicitly pass R_index. * T = double or float + * + * @param R_index Index of the R vector in this->values + * @param target_array Pointer to the target array. + * @param kphase Scalar to be multiplied with the block matrix. */ - void add_to_array(std::complex* target_array, const std::complex& kphase) const; + void add_to_array(const int R_index, std::complex* target_array, const std::complex& kphase) const; /** - * @brief Add this->value[current_R] * kphase to an array. + * @brief Add this->value[R_index] * kphase to an array. + * Thread-safe version: explicitly pass R_index. * for non-collinear spin case only (T = std::complex or complex) + * + * @param R_index Index of the R vector in this->values + * @param target_array Pointer to the target array. + * @param kphase Scalar to be multiplied with the block matrix. */ - void add_to_array(T* target_array, const T& kphase) const; + void add_to_array(const int R_index, T* target_array, const T& kphase) const; // comparation function, used for sorting bool operator<(const AtomPair& other) const; @@ -309,13 +379,6 @@ class AtomPair // only for 2d-block const Parallel_Orbitals* paraV = nullptr; - // the default R index is (0, 0, 0) - // if current_R > 0, it means R index has been fixed - // if current_R == 0, it means R index refers to the first cell - // if current_R == 0 with gamma_only, it means R index refers to the center cell - // !!!!!!!!!!! BE CAREFUL, current_R IS NOT THREADING-SAFE !!!!!!!!!!!!!!!!!!!!! - mutable int current_R = 0; - // index for identifying atom I and J for this atom-pair int atom_i = -1; int atom_j = -1; diff --git a/source/source_lcao/module_hcontainer/func_folding.cpp b/source/source_lcao/module_hcontainer/func_folding.cpp index 4b276846d68..0f828382366 100644 --- a/source/source_lcao/module_hcontainer/func_folding.cpp +++ b/source/source_lcao/module_hcontainer/func_folding.cpp @@ -131,8 +131,9 @@ void folding_HR(const hamilt::HContainer& hR, // if TK==double, kphase is 1.0 double kphase = 1.0; - // Hk = HR - hR.get_atom_pair(i).add_to_matrix(hk, hk_ld , kphase, hk_type); + // Hk = HR + // Gamma-only case: only R_index=0 exists, pass it explicitly + hR.get_atom_pair(i).add_to_matrix(0, hk, hk_ld, kphase, hk_type); } } diff --git a/source/source_lcao/module_hcontainer/hcontainer.cpp b/source/source_lcao/module_hcontainer/hcontainer.cpp index e1d1b909912..a50c7ef16be 100644 --- a/source/source_lcao/module_hcontainer/hcontainer.cpp +++ b/source/source_lcao/module_hcontainer/hcontainer.cpp @@ -619,7 +619,7 @@ T* HContainer::data(int atom_i, int atom_j) const AtomPair* atom_ij = this->find_pair(atom_i, atom_j); if (atom_ij != nullptr) { - return atom_ij->get_pointer(); + return atom_ij->get_pointer(0); } else { diff --git a/source/source_lcao/module_hcontainer/output_hcontainer.cpp b/source/source_lcao/module_hcontainer/output_hcontainer.cpp index ad3a1d26dd9..41971575cdd 100644 --- a/source/source_lcao/module_hcontainer/output_hcontainer.cpp +++ b/source/source_lcao/module_hcontainer/output_hcontainer.cpp @@ -124,7 +124,9 @@ void Output_HContainer::write_single_R(int rx, int ry, int rz) for (int iap = 0; iap < this->_hcontainer->size_atom_pairs(); ++iap) { auto atom_pair = this->_hcontainer->get_atom_pair(iap); - auto tmp_matrix_info = atom_pair.get_matrix_values(); + const int r_index = atom_pair.find_R(rx, ry, rz); + if (r_index < 0) continue; + auto tmp_matrix_info = atom_pair.get_matrix_values(r_index); int* tmp_index = std::get<0>(tmp_matrix_info).data(); T* tmp_data = std::get<1>(tmp_matrix_info); for (int irow = tmp_index[0]; irow < tmp_index[0] + tmp_index[1]; ++irow) diff --git a/source/source_lcao/module_hcontainer/test/test_hcontainer.cpp b/source/source_lcao/module_hcontainer/test/test_hcontainer.cpp index ef948f87d8d..c7f4e8e90f1 100644 --- a/source/source_lcao/module_hcontainer/test/test_hcontainer.cpp +++ b/source/source_lcao/module_hcontainer/test/test_hcontainer.cpp @@ -82,7 +82,8 @@ TEST_F(HContainerTest, insert_pair) hamilt::AtomPair atom_kl(1, 0); atom_kl.set_size(2, 2); double tmp_array[4] = {1, 2, 3, 4}; - atom_kl.get_HR_values(1, 0, 0).add_array(&tmp_array[0]); + const int rx = 1, ry = 0, rz = 0; + atom_kl.get_HR_values(rx, ry, rz).add_array(&tmp_array[0]); // insert atom_kl into HR HR->insert_pair(atom_kl); // check if atom_kl is inserted into HR @@ -96,7 +97,7 @@ TEST_F(HContainerTest, insert_pair) EXPECT_EQ(HR->get_atom_pair(1, 0).get_row_size(), 2); EXPECT_EQ(HR->get_atom_pair(1, 0).get_col_size(), 2); // check if data is correct - double* data_ptr = HR->get_atom_pair(1, 0).get_HR_values(1, 0, 0).get_pointer(); + double* data_ptr = HR->get_atom_pair(1, 0).get_HR_values(rx, ry, rz).get_pointer(); EXPECT_EQ(data_ptr[0], 1); EXPECT_EQ(data_ptr[1], 2); EXPECT_EQ(data_ptr[2], 3); @@ -195,7 +196,8 @@ TEST_F(HContainerTest, fix_R) // check if R is fixed EXPECT_EQ(HR->get_current_R(), 0); // fix another R - EXPECT_EQ(HR->fix_R(1, 0, 0), false); + const int rx_another = 1, ry_another = 0, rz_another = 0; + EXPECT_EQ(HR->fix_R(rx_another, ry_another, rz_another), false); // check if R is fixed EXPECT_EQ(HR->get_current_R(), -1); // unfix R @@ -216,7 +218,8 @@ TEST_F(HContainerTest, fix_gamma) EXPECT_EQ(HR->size_R_loop(), 1); hamilt::AtomPair atom_ij(0, 1); atom_ij.set_size(2, 2); - hamilt::BaseMatrix& tmp = atom_ij.get_HR_values(1, 0, 0); + const int rx_offset = 1, ry_offset = 0, rz_offset = 0; + hamilt::BaseMatrix& tmp = atom_ij.get_HR_values(rx_offset, ry_offset, rz_offset); double tmp_array[4] = {1, 2, 3, 4}; tmp.add_array(tmp_array); // insert atom_ij into HR @@ -282,7 +285,8 @@ TEST_F(HContainerTest, size_atom_pairs) // fix to another R hamilt::AtomPair atom_ij(0, 1); atom_ij.set_size(2, 2); - hamilt::BaseMatrix& tmp = atom_ij.get_HR_values(1, 0, 0); + const int rx_another = 1, ry_another = 0, rz_another = 0; + hamilt::BaseMatrix& tmp = atom_ij.get_HR_values(rx_another, ry_another, rz_another); double tmp_array[4] = {1, 2, 3, 4}; tmp.add_array(tmp_array); // insert atom_ij into HR @@ -296,7 +300,7 @@ TEST_F(HContainerTest, size_atom_pairs) // check size_R_loop after R index unfixed size_R_loop = HR->size_R_loop(); EXPECT_EQ(size_R_loop, 2); - ok = HR->fix_R(1, 0, 0); + ok = HR->fix_R(rx_another, ry_another, rz_another); // check if R is fixed EXPECT_EQ(ok, true); EXPECT_EQ(HR->get_current_R(), 1); @@ -306,13 +310,13 @@ TEST_F(HContainerTest, size_atom_pairs) EXPECT_EQ(HR->get_atom_pair(0).get_atom_j(), 1); EXPECT_EQ(HR->get_atom_pair(0).get_row_size(), 2); EXPECT_EQ(HR->get_atom_pair(0).get_col_size(), 2); - const ModuleBase::Vector3 R_ptr = HR->get_atom_pair(0).get_R_index(); - EXPECT_EQ(R_ptr.x, 1); - EXPECT_EQ(R_ptr.y, 0); - EXPECT_EQ(R_ptr.z, 0); + const int r_index_fix = HR->get_atom_pair(0).find_R(rx_another, ry_another, rz_another); + EXPECT_EQ(HR->get_atom_pair(0).get_R_index(r_index_fix).x, rx_another); + EXPECT_EQ(HR->get_atom_pair(0).get_R_index(r_index_fix).y, ry_another); + EXPECT_EQ(HR->get_atom_pair(0).get_R_index(r_index_fix).z, rz_another); EXPECT_EQ(HR->get_atom_pair(0).get_R_index(5), ModuleBase::Vector3(-1, -1, -1)); // check if data is correct - double* data_ptr = HR->get_atom_pair(0).get_pointer(); + double* data_ptr = HR->get_atom_pair(0).get_pointer(r_index_fix); EXPECT_EQ(data_ptr[0], 1); EXPECT_EQ(data_ptr[1], 2); EXPECT_EQ(data_ptr[2], 3); @@ -337,9 +341,11 @@ TEST_F(HContainerTest, size_atom_pairs) TEST_F(HContainerTest, data) { // set up a hamilt::AtomPair - hamilt::AtomPair atom_ij(0, 1); + const int atom_i = 0, atom_j = 1; + hamilt::AtomPair atom_ij(atom_i, atom_j); atom_ij.set_size(2, 2); - hamilt::BaseMatrix& tmp = atom_ij.get_HR_values(0, 0, 0); + const int rx_center = 0, ry_center = 0, rz_center = 0; + hamilt::BaseMatrix& tmp = atom_ij.get_HR_values(rx_center, ry_center, rz_center); double tmp_array[4] = {1, 2, 3, 4}; tmp.add_array(tmp_array); EXPECT_EQ(HR->size_atom_pairs(), 9); @@ -347,17 +353,17 @@ TEST_F(HContainerTest, data) HR->insert_pair(atom_ij); // check if atom_ij is inserted into HR EXPECT_EQ(HR->size_atom_pairs(), 9); - EXPECT_EQ(HR->get_atom_pair(0, 1).get_atom_i(), 0); - EXPECT_EQ(HR->get_atom_pair(0, 1).get_atom_j(), 1); - EXPECT_EQ(HR->get_atom_pair(0, 1).get_row_size(), 2); - EXPECT_EQ(HR->get_atom_pair(0, 1).get_col_size(), 2); + EXPECT_EQ(HR->get_atom_pair(atom_i, atom_j).get_atom_i(), atom_i); + EXPECT_EQ(HR->get_atom_pair(atom_i, atom_j).get_atom_j(), atom_j); + EXPECT_EQ(HR->get_atom_pair(atom_i, atom_j).get_row_size(), 2); + EXPECT_EQ(HR->get_atom_pair(atom_i, atom_j).get_col_size(), 2); // get data pointer - double* data_ptr = HR->data(0, 1); + double* data_ptr = HR->data(atom_i, atom_j); // check if data pointer is correct - EXPECT_EQ(data_ptr, HR->get_atom_pair(0, 1).get_pointer()); - EXPECT_EQ(data_ptr, HR->get_atom_pair(0, 1).get_pointer(0)); - int r_index[3] = {0, 0, 0}; - EXPECT_EQ(data_ptr, HR->data(0, 1, r_index)); + EXPECT_EQ(data_ptr, HR->get_atom_pair(atom_i, atom_j).get_pointer(0)); + EXPECT_EQ(data_ptr, HR->get_atom_pair(atom_i, atom_j).get_pointer(0)); + int r_index[3] = {rx_center, ry_center, rz_center}; + EXPECT_EQ(data_ptr, HR->data(atom_i, atom_j, r_index)); EXPECT_EQ(HR->data(0, 10), nullptr); EXPECT_EQ(HR->data(0, 10, r_index), nullptr); // check if data is correct @@ -418,10 +424,12 @@ TEST_F(HContainerTest, atompair_funcs) } PO.nrow = 4; PO.ncol = 4; + const int natom = 2; + const int rx_shift = 1, ry_shift = 1, rz_shift = 1; hamilt::AtomPair atom_ij(0, 0, &PO, nullptr); - hamilt::AtomPair atom_ij2(0, 1, 1, 1, 1, &PO, nullptr); - hamilt::AtomPair atom_ij3(1, 0, PO.atom_begin_row.data(), PO.atom_begin_col.data(), 2, nullptr); - hamilt::AtomPair atom_ij33(1, 1, 1, 1, 1, PO.atom_begin_row.data(), PO.atom_begin_col.data(), 2, nullptr); + hamilt::AtomPair atom_ij2(0, 1, rx_shift, ry_shift, rz_shift, &PO, nullptr); + hamilt::AtomPair atom_ij3(1, 0, PO.atom_begin_row.data(), PO.atom_begin_col.data(), natom, nullptr); + hamilt::AtomPair atom_ij33(1, 1, rx_shift, ry_shift, rz_shift, PO.atom_begin_row.data(), PO.atom_begin_col.data(), natom, nullptr); EXPECT_EQ(atom_ij HR(2); - for(int atom_i = 0;atom_i<2;++atom_i) + const int matrix_size = 2; + const int rx_center = 0, ry_center = 0, rz_center = 0; + hamilt::HContainer HR(natom); + for(int atom_i = 0; atom_i < natom; ++atom_i) { - for(int atom_j = 0; atom_j<2; ++atom_j) + for(int atom_j = 0; atom_j < natom; ++atom_j) { - hamilt::AtomPair tmp(atom_i, atom_j, 0, 0, 0, PO.atom_begin_row.data(), PO.atom_begin_col.data(), 2, nullptr); + hamilt::AtomPair tmp(atom_i, atom_j, rx_center, ry_center, rz_center, PO.atom_begin_row.data(), PO.atom_begin_col.data(), natom, nullptr); tmp.allocate(nullptr, false); - double* tmp_data = tmp.get_HR_values(0, 0, 0).get_pointer(); - for(int i=0;i<4;++i) + double* tmp_data = tmp.get_HR_values(rx_center, ry_center, rz_center).get_pointer(); + for(int i = 0; i < 4; ++i) { - tmp_data[i] = atom_i*2 + atom_j*4 + i + 1; + tmp_data[i] = atom_i * 2 + atom_j * 4 + i + 1; } HR.insert_pair(tmp); } @@ -474,12 +484,25 @@ TEST_F(HContainerTest, atompair_funcs) for(int iap = 0;iap& tmp = HR.get_atom_pair(iap); - // row major case - tmp.add_to_matrix(&hk_data[0], 4, 1.0, 0); - tmp.add_to_matrix(&hk_data2[0], 4, std::complex(1.0, 0.5), 0); - // colomn major case - tmp.add_to_matrix(&hk_data1[0], 4, 1.0, 1); - tmp.add_to_matrix(&hk_data3[0], 4, std::complex(1.0, 0.5), 1); + // Find the R_index that matches the current R + int r_index = -1; + for(int j = 0; j < tmp.get_R_size(); ++j) + { + if(tmp.get_R_index(j) == ModuleBase::Vector3(rx, ry, rz)) + { + r_index = j; + break; + } + } + if(r_index >= 0) + { + // row major case - use thread-safe version with explicit R_index + tmp.add_to_matrix(r_index, &hk_data[0], 4, 1.0, 0); + tmp.add_to_matrix(r_index, &hk_data2[0], 4, std::complex(1.0, 0.5), 0); + // colomn major case + tmp.add_to_matrix(r_index, &hk_data1[0], 4, 1.0, 1); + tmp.add_to_matrix(r_index, &hk_data3[0], 4, std::complex(1.0, 0.5), 1); + } /* output for show the result std::cout<<"iap = "<& ap_in) { - auto data_ij4 = ap_in.get_matrix_values(); + // Note: these AtomPairs only contain R=(0,0,0), so r_index is always 0 + auto checkdata = [&](hamilt::AtomPair& ap_in, const int r_index = 0) { + auto data_ij4 = ap_in.get_matrix_values(r_index); int* tmp_index = std::get<0>(data_ij4).data(); double* tmp_data = std::get<1>(data_ij4); double sum_error = 0.0; @@ -625,33 +664,34 @@ TEST_F(HContainerTest, wrapper_mode) EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_row_size(), 2); EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_col_size(), 2); EXPECT_EQ(HR_wrapper.get_wrapper(), hr_data.data()); - EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_pointer(), hr_data.data()); + const int r_index_center = 0; + EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_pointer(r_index_center), hr_data.data()); for (size_t i = 0; i < hr_data.size(); i++) { hr_data[i] = i; } - EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(0, 0), 0.0); - EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(0, 1), 1.0); - EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(1, 0), 2.0); - EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(1, 1), 3.0); - EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(0, 0), 4.0); - EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(0, 1), 5.0); - EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(1, 0), 6.0); - EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(1, 1), 7.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(r_index_center, 0, 0), 0.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(r_index_center, 0, 1), 1.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(r_index_center, 1, 0), 2.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(0).get_value(r_index_center, 1, 1), 3.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(r_index_center, 0, 0), 4.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(r_index_center, 0, 1), 5.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(r_index_center, 1, 0), 6.0); + EXPECT_EQ(HR_wrapper.get_atom_pair(1).get_value(r_index_center, 1, 1), 7.0); hamilt::AtomPair atom_ij(HR->get_atom_pair(0), hr_data.data()); hamilt::BaseMatrix matrix_test = hamilt::BaseMatrix(atom_ij.get_row_size(), atom_ij.get_col_size(), hr_data.data()); - EXPECT_EQ(atom_ij.get_value(1, 1), 3.0); + EXPECT_EQ(atom_ij.get_value(r_index_center, 1, 1), 3.0); EXPECT_EQ(matrix_test.get_value(1, 1), 3.0); HR->allocate(hr_data.data(), false); - EXPECT_EQ(HR->get_atom_pair(0).get_value(1, 1), 3.0); - EXPECT_EQ(HR->get_atom_pair(1).get_value(1, 1), 7.0); - EXPECT_EQ(HR->get_atom_pair(2).get_value(1, 1), 11.0); - EXPECT_EQ(HR->get_atom_pair(3).get_value(1, 1), 15.0); + EXPECT_EQ(HR->get_atom_pair(0).get_value(r_index_center, 1, 1), 3.0); + EXPECT_EQ(HR->get_atom_pair(1).get_value(r_index_center, 1, 1), 7.0); + EXPECT_EQ(HR->get_atom_pair(2).get_value(r_index_center, 1, 1), 11.0); + EXPECT_EQ(HR->get_atom_pair(3).get_value(r_index_center, 1, 1), 15.0); HR->allocate(hr_data.data(), true); - EXPECT_EQ(HR->get_atom_pair(0).get_value(1, 1), 0.0); - EXPECT_EQ(HR->get_atom_pair(1).get_value(1, 1), 0.0); - EXPECT_EQ(HR->get_atom_pair(2).get_value(1, 1), 0.0); - EXPECT_EQ(HR->get_atom_pair(3).get_value(1, 1), 0.0); + EXPECT_EQ(HR->get_atom_pair(0).get_value(r_index_center, 1, 1), 0.0); + EXPECT_EQ(HR->get_atom_pair(1).get_value(r_index_center, 1, 1), 0.0); + EXPECT_EQ(HR->get_atom_pair(2).get_value(r_index_center, 1, 1), 0.0); + EXPECT_EQ(HR->get_atom_pair(3).get_value(r_index_center, 1, 1), 0.0); } int main(int argc, char** argv) diff --git a/source/source_lcao/module_hcontainer/test/test_hcontainer_complex.cpp b/source/source_lcao/module_hcontainer/test/test_hcontainer_complex.cpp index 954e7cfd6e3..65795bc71aa 100644 --- a/source/source_lcao/module_hcontainer/test/test_hcontainer_complex.cpp +++ b/source/source_lcao/module_hcontainer/test/test_hcontainer_complex.cpp @@ -309,13 +309,13 @@ TEST_F(HContainerTest, size_atom_pairs) EXPECT_EQ(HR->get_atom_pair(0).get_atom_j(), 1); EXPECT_EQ(HR->get_atom_pair(0).get_row_size(), 2); EXPECT_EQ(HR->get_atom_pair(0).get_col_size(), 2); - const ModuleBase::Vector3 R_ptr = HR->get_atom_pair(0).get_R_index(); - EXPECT_EQ(R_ptr.x, 1); - EXPECT_EQ(R_ptr.y, 0); - EXPECT_EQ(R_ptr.z, 0); + const int r_index_fix = HR->get_atom_pair(0).find_R(1, 0, 0); + EXPECT_EQ(HR->get_atom_pair(0).get_R_index(r_index_fix).x, 1); + EXPECT_EQ(HR->get_atom_pair(0).get_R_index(r_index_fix).y, 0); + EXPECT_EQ(HR->get_atom_pair(0).get_R_index(r_index_fix).z, 0); EXPECT_EQ(HR->get_atom_pair(0).get_R_index(5), ModuleBase::Vector3(-1, -1, -1)); // check if data is correct - std::complex* data_ptr = HR->get_atom_pair(0).get_pointer(); + std::complex* data_ptr = HR->get_atom_pair(0).get_pointer(r_index_fix); EXPECT_EQ(data_ptr[0], std::complex(1)); EXPECT_EQ(data_ptr[1], std::complex(2)); EXPECT_EQ(data_ptr[2], std::complex(3)); @@ -358,7 +358,7 @@ TEST_F(HContainerTest, data) // get data pointer std::complex* data_ptr = HR->data(0, 1); // check if data pointer is correct - EXPECT_EQ(data_ptr, HR->get_atom_pair(0, 1).get_pointer()); + EXPECT_EQ(data_ptr, HR->get_atom_pair(0, 1).get_pointer(0)); EXPECT_EQ(data_ptr, HR->get_atom_pair(0, 1).get_pointer(0)); int r_index[3] = {0, 0, 0}; EXPECT_EQ(data_ptr, HR->data(0, 1, r_index)); @@ -468,10 +468,23 @@ TEST_F(HContainerTest, atompair_funcs) for(int iap = 0;iap>& tmp = HR.get_atom_pair(iap); - // row major case - tmp.add_to_matrix(&hk_data2[0], 4, std::complex(1.0, 0.5), 0); - // colomn major case - tmp.add_to_matrix(&hk_data3[0], 4, std::complex(1.0, 0.5), 1); + // Find the R_index that matches the current R + int r_index = -1; + for(int j = 0; j < tmp.get_R_size(); ++j) + { + if(tmp.get_R_index(j) == ModuleBase::Vector3(rx, ry, rz)) + { + r_index = j; + break; + } + } + if(r_index >= 0) + { + // row major case - use thread-safe version with explicit R_index + tmp.add_to_matrix(r_index, &hk_data2[0], 4, std::complex(1.0, 0.5), 0); + // colomn major case + tmp.add_to_matrix(r_index, &hk_data3[0], 4, std::complex(1.0, 0.5), 1); + } } } HR.unfix_R(); @@ -500,7 +513,21 @@ TEST_F(HContainerTest, atompair_funcs) for(int iap = 0;iap(1.0, 0.0)); + // Find the R_index that matches the current R + int r_index = -1; + for(int j = 0; j < tmp.get_R_size(); ++j) + { + if(tmp.get_R_index(j) == ModuleBase::Vector3(rx, ry, rz)) + { + r_index = j; + break; + } + } + if(r_index >= 0) + { + // Use thread-safe version with explicit R_index + tmp.add_to_array(r_index, ptr1, std::complex(1.0, 0.0)); + } ptr1 += tmp.get_size(); } } @@ -525,9 +552,10 @@ TEST_F(HContainerTest, atompair_funcs) // construct AtomPair from existed matrix hamilt::AtomPair> atom_ij4(0, 0, &PO, test_array); - EXPECT_EQ(atom_ij4.get_value(0, 0), correct_array[0]); - EXPECT_EQ(atom_ij4.get_value(1, 1), correct_array[5]); - EXPECT_EQ(atom_ij4.get_value(0), correct_array[0]); + const int r_index_center = 0; + EXPECT_EQ(atom_ij4.get_value(r_index_center, 0, 0), correct_array[0]); + EXPECT_EQ(atom_ij4.get_value(r_index_center, 1, 1), correct_array[5]); + EXPECT_EQ(atom_ij4.get_value(r_index_center, 0), correct_array[0]); hamilt::AtomPair> atom_ij5(0, 1, 1, 1, 1, &PO, &test_array[4]); hamilt::AtomPair> atom_ij6(1, 0, PO.atom_begin_row.data(), PO.atom_begin_col.data(), 2, &test_array[8]); hamilt::AtomPair> atom_ij7(1, 1, 1, 1, 1, PO.atom_begin_row.data(), PO.atom_begin_col.data(), 2, &test_array[12]); @@ -535,8 +563,9 @@ TEST_F(HContainerTest, atompair_funcs) // so we need to set them std::ofstream ofs("test_hcontainer_complex.log"); PO.set_serial(4, 4); - auto checkdata = [&](hamilt::AtomPair>& ap_in) { - auto data_ij4 = ap_in.get_matrix_values(); + // Note: these AtomPairs only contain R=(0,0,0), so r_index is always 0 + auto checkdata = [&](hamilt::AtomPair>& ap_in, const int r_index = 0) { + auto data_ij4 = ap_in.get_matrix_values(r_index); int* tmp_index = std::get<0>(data_ij4).data(); std::complex* tmp_data = std::get<1>(data_ij4); double sum_error = 0.0; diff --git a/source/source_lcao/module_hcontainer/test/test_hcontainer_output.cpp b/source/source_lcao/module_hcontainer/test/test_hcontainer_output.cpp index 70fdcb66af2..5fa66f0b478 100644 --- a/source/source_lcao/module_hcontainer/test/test_hcontainer_output.cpp +++ b/source/source_lcao/module_hcontainer/test/test_hcontainer_output.cpp @@ -101,27 +101,29 @@ TEST_F(OutputHContainerTest, Write) for (int iap = 0; iap < HR.size_atom_pairs(); ++iap) { hamilt::AtomPair& tmp_ap = HR.get_atom_pair(iap); + const int r_index = tmp_ap.find_R(rx, ry, rz); + if (r_index < 0) continue; if (rx == 0 && ry == 1 && rz == 1) { - EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 0), 0); - EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 1), 4); - EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 0), 7); - EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 1), 0); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[0], 0); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[1], 2); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[2], 2); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[3], 2); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 0, 0), 0); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 0, 1), 4); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 1, 0), 7); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 1, 1), 0); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[0], 0); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[1], 2); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[2], 2); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[3], 2); } else if (rx == 0 && ry == 0 && rz == 0) { - EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 0), 5); - EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 1), 6); - EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 0), 0); - EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 1), 10); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[0], 2); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[1], 2); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[2], 2); - EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values())[3], 2); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 0, 0), 5); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 0, 1), 6); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 1, 0), 0); + EXPECT_DOUBLE_EQ(tmp_ap.get_value(r_index, 1, 1), 10); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[0], 2); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[1], 2); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[2], 2); + EXPECT_DOUBLE_EQ(std::get<0>(tmp_ap.get_matrix_values(r_index))[3], 2); } } HR.unfix_R(); diff --git a/source/source_lcao/module_rt/td_folding.cpp b/source/source_lcao/module_rt/td_folding.cpp index 2df49042ca2..6eac2e13da7 100644 --- a/source/source_lcao/module_rt/td_folding.cpp +++ b/source/source_lcao/module_rt/td_folding.cpp @@ -28,8 +28,7 @@ void folding_HR_td(const hamilt::HContainer& hR, std::complex kphase = std::complex(cosp, sinp); kphase *= phase_hybrid.at(r_index); - tmp.find_R(r_index); - tmp.add_to_matrix(hk, ncol, kphase, hk_type); + tmp.add_to_matrix(ir, hk, ncol, kphase, hk_type); } } } @@ -59,8 +58,7 @@ void folding_partial_HR(const UnitCell& ucell, std::complex kphase = std::complex(cosp, sinp); const ModuleBase::Vector3 dR_car = dR * ucell.latvec * ucell.lat0; - tmp.find_R(r_index); - tmp.add_to_matrix(hk, ncol, kphase * ModuleBase::IMAG_UNIT * std::complex(dR_car[ix]), hk_type); + tmp.add_to_matrix(ir, hk, ncol, kphase * ModuleBase::IMAG_UNIT * std::complex(dR_car[ix]), hk_type); } } } @@ -97,8 +95,7 @@ void folding_partial_HR_td(const UnitCell& ucell, kphase *= phase_hybrid.at(r_index); const ModuleBase::Vector3 dR_car = dR * ucell.latvec * ucell.lat0; - tmp.find_R(r_index); - tmp.add_to_matrix(hk, ncol, kphase * ModuleBase::IMAG_UNIT * std::complex(dR_car[ix]), hk_type); + tmp.add_to_matrix(ir, hk, ncol, kphase * ModuleBase::IMAG_UNIT * std::complex(dR_car[ix]), hk_type); } } } @@ -133,8 +130,7 @@ void folding_partial_dot(const hamilt::HContainer& dR, kphase *= phase_hybrid.at(r_index); ModuleBase::Vector3 dtau = ucell->cal_dtau(iat1, iat1, r_index); kphase *= Et * dtau * ucell->lat0; - tmp.find_R(r_index); - tmp.add_to_matrix(dk, ncol, kphase, hk_type); + tmp.add_to_matrix(ir, dk, ncol, kphase, hk_type); } } } diff --git a/source/source_lcao/test/test_init_dm_from_file.cpp b/source/source_lcao/test/test_init_dm_from_file.cpp index 31079f28855..b14f48a0cf7 100644 --- a/source/source_lcao/test/test_init_dm_from_file.cpp +++ b/source/source_lcao/test/test_init_dm_from_file.cpp @@ -100,9 +100,10 @@ class InitDMFileTest : public testing::Test if (ap) { int nw = ucell.atoms[0].nw; + const int r_index = 0; // R=(0,0,0) for (int k = 0; k < nw; k++) { - ap->get_pointer()[k * nw + k] = scale * (k + 1) * 0.1; + ap->get_pointer(r_index)[k * nw + k] = scale * (k + 1) * 0.1; } } } @@ -157,9 +158,10 @@ TEST_F(InitDMFileTest, Nspin1_ReadSingleFile) auto* ap = dmr0->find_pair(0, 0); ASSERT_NE(ap, nullptr); bool has_nonzero = false; + const int r_index = 0; // R=(0,0,0) for (int i = 0; i < ap->get_size(); i++) { - if (std::abs(ap->get_pointer()[i]) > 1e-15) + if (std::abs(ap->get_pointer(r_index)[i]) > 1e-15) { has_nonzero = true; break; @@ -200,10 +202,11 @@ TEST_F(InitDMFileTest, Nspin2_ReadTwoFiles) bool values_differ = false; int check_size = std::min(ap0->get_size(), ap1->get_size()); + const int r_index = 0; // R=(0,0,0) for (int i = 0; i < check_size; i++) { - double v0 = ap0->get_pointer()[i]; - double v1 = ap1->get_pointer()[i]; + double v0 = ap0->get_pointer(r_index)[i]; + double v1 = ap1->get_pointer(r_index)[i]; if (std::abs(v0) > 1e-15 && std::abs(v0 - v1) > 1e-15) { values_differ = true; @@ -271,10 +274,11 @@ TEST_F(InitDMFileTest, HR_Nspin1_ReadSingleFile) // Check diagonal has expected values (scale=2.0, value = 2.0*(k+1)*0.1) int nw = ucell.atoms[0].nw; + const int r_index = 0; // R=(0,0,0) for (int k = 0; k < nw; k++) { double expected = 2.0 * (k + 1) * 0.1; - EXPECT_NEAR(ap->get_pointer()[k * nw + k], expected, 1e-6); + EXPECT_NEAR(ap->get_pointer(r_index)[k * nw + k], expected, 1e-6); } } @@ -321,10 +325,11 @@ TEST_F(InitDMFileTest, HR_Nspin2_ReadTwoFiles) auto* ap_up = hR_up->find_pair(0, 0); ASSERT_NE(ap_up, nullptr); int nw = ucell.atoms[0].nw; + const int r_index = 0; // R=(0,0,0) for (int k = 0; k < nw; k++) { double expected_up = 1.0 * (k + 1) * 0.1; - EXPECT_NEAR(ap_up->get_pointer()[k * nw + k], expected_up, 1e-6); + EXPECT_NEAR(ap_up->get_pointer(r_index)[k * nw + k], expected_up, 1e-6); } // Verify spin-down values (scale=3.0) @@ -333,12 +338,12 @@ TEST_F(InitDMFileTest, HR_Nspin2_ReadTwoFiles) for (int k = 0; k < nw; k++) { double expected_down = 3.0 * (k + 1) * 0.1; - EXPECT_NEAR(ap_down->get_pointer()[k * nw + k], expected_down, 1e-6); + EXPECT_NEAR(ap_down->get_pointer(r_index)[k * nw + k], expected_down, 1e-6); } // Verify the two are independent (different values) - double val_up = ap_up->get_pointer()[0]; - double val_down = ap_down->get_pointer()[0]; + double val_up = ap_up->get_pointer(r_index)[0]; + double val_down = ap_down->get_pointer(r_index)[0]; EXPECT_GT(std::abs(val_up), 1e-15); EXPECT_NEAR(val_down / val_up, 3.0, 1e-6); diff --git a/source/source_lcao/test/test_output_hcontainer_consistency.cpp b/source/source_lcao/test/test_output_hcontainer_consistency.cpp index f8539220e86..fbecde2739b 100644 --- a/source/source_lcao/test/test_output_hcontainer_consistency.cpp +++ b/source/source_lcao/test/test_output_hcontainer_consistency.cpp @@ -102,9 +102,10 @@ class OutputHContainerTest : public testing::Test if (ap) { int nw = ucell.atoms[0].nw; + const int r_index = 0; // R=(0,0,0) for (int k = 0; k < nw; k++) { - ap->get_pointer()[k * nw + k] = scale * (k + 1) * 0.1; + ap->get_pointer(r_index)[k * nw + k] = scale * (k + 1) * 0.1; } } } @@ -143,10 +144,11 @@ TEST_F(OutputHContainerTest, WriteReadConsistency) ASSERT_NE(ap_w, nullptr); ASSERT_NE(ap_r, nullptr); + const int r_index = 0; // R=(0,0,0) for (int k = 0; k < nw; k++) { - EXPECT_NEAR(ap_w->get_pointer()[k * nw + k], - ap_r->get_pointer()[k * nw + k], 1e-8) + EXPECT_NEAR(ap_w->get_pointer(r_index)[k * nw + k], + ap_r->get_pointer(r_index)[k * nw + k], 1e-8) << "Mismatch at atom " << i << " orbital " << k; } } @@ -161,10 +163,11 @@ TEST_F(OutputHContainerTest, WriteReadConsistency) auto* ap_r = hc_read->find_pair(i, j); if (ap_w && ap_r) { + const int r_index = 0; // R=(0,0,0) for (int k = 0; k < nw * nw; k++) { - EXPECT_NEAR(ap_w->get_pointer()[k], - ap_r->get_pointer()[k], 1e-10); + EXPECT_NEAR(ap_w->get_pointer(r_index)[k], + ap_r->get_pointer(r_index)[k], 1e-10); } } } @@ -227,11 +230,12 @@ TEST_F(OutputHContainerTest, PrecisionParameter) auto* ap_r = hc_read->find_pair(0, 0); ASSERT_NE(ap_r, nullptr); + const int r_index = 0; // R=(0,0,0) for (int k = 0; k < nw; k++) { double expected = 1.23456789 * (k + 1) * 0.1; // 4-digit precision => ~1e-4 tolerance - EXPECT_NEAR(ap_r->get_pointer()[k * nw + k], expected, 5e-4) + EXPECT_NEAR(ap_r->get_pointer(r_index)[k * nw + k], expected, 5e-4) << "Low-precision round-trip failed at orbital " << k; } @@ -264,17 +268,19 @@ TEST_F(OutputHContainerTest, Nspin2TwoFileConsistency) reader_down.read(); int nw = ucell.atoms[0].nw; + // Note: these HContainer only contain R=(0,0,0), so r_index is always 0 + const int r_index_center = 0; for (int k = 0; k < nw; k++) { double exp_up = 1.0 * (k + 1) * 0.1; double exp_down = 3.0 * (k + 1) * 0.1; - EXPECT_NEAR(hc_read_up->find_pair(0, 0)->get_pointer()[k * nw + k], exp_up, 1e-8); - EXPECT_NEAR(hc_read_down->find_pair(0, 0)->get_pointer()[k * nw + k], exp_down, 1e-8); + EXPECT_NEAR(hc_read_up->find_pair(0, 0)->get_pointer(r_index_center)[k * nw + k], exp_up, 1e-8); + EXPECT_NEAR(hc_read_down->find_pair(0, 0)->get_pointer(r_index_center)[k * nw + k], exp_down, 1e-8); } // Verify the two are independent - EXPECT_GT(std::abs(hc_read_up->find_pair(0, 0)->get_pointer()[0] - - hc_read_down->find_pair(0, 0)->get_pointer()[0]), 1e-6); + EXPECT_GT(std::abs(hc_read_up->find_pair(0, 0)->get_pointer(r_index_center)[0] + - hc_read_down->find_pair(0, 0)->get_pointer(r_index_center)[0]), 1e-6); delete hc_up; delete hc_down;