diff --git a/source/source_estate/occ_matrix.cpp b/source/source_estate/occ_matrix.cpp index 349d4a896f..1294e42328 100644 --- a/source/source_estate/occ_matrix.cpp +++ b/source/source_estate/occ_matrix.cpp @@ -3,6 +3,8 @@ #include "source_base/timer.h" #include "source_cell/unitcell.h" +#include + void OccupationMatrix::init(const UnitCell& cell, const std::vector& l_channel, const int nspin, @@ -30,53 +32,61 @@ void OccupationMatrix::init(const UnitCell& cell, continue; } - // only the first radial channel (n=0) of each l is stored - for (int l = 0; l <= cell.atoms[it].nwl; l++) - { - if (nspin == 1 || nspin == 2) - { - occ_[iat][l].resize(2); - occ_save_[iat][l].resize(2); + allocate_atom_matrices(cell, iat, it); + build_corr_iwt(cell, iat, ia); + } + } +} - occ_[iat][l][0].create(2 * l + 1, 2 * l + 1); - occ_[iat][l][1].create(2 * l + 1, 2 * l + 1); +void OccupationMatrix::allocate_atom_matrices(const UnitCell& cell, + const int iat, + const int it) +{ + const int nchan = (nspin_ == 4) ? 1 : 2; + const int nwl = cell.atoms[it].nwl; - occ_save_[iat][l][0].create(2 * l + 1, 2 * l + 1); - occ_save_[iat][l][1].create(2 * l + 1, 2 * l + 1); - } - else if (nspin == 4) - { - occ_[iat][l].resize(1); - occ_save_[iat][l].resize(1); + for (int l = 0; l <= nwl; l++) + { + const int dim = (nspin_ == 4) ? (2 * l + 1) * npol_ : 2 * l + 1; - occ_[iat][l][0].create((2 * l + 1) * npol, (2 * l + 1) * npol); - occ_save_[iat][l][0].create((2 * l + 1) * npol, (2 * l + 1) * npol); - } - } + occ_[iat][l].resize(nchan); + occ_save_[iat][l].resize(nchan); - for (int L = 0; L <= cell.atoms[it].nwl; L++) - { - corr_iwt_[iat][L].resize(2 * L + 1); - for (int m = 0; m < 2 * L + 1; m++) - { - corr_iwt_[iat][L][m].resize(npol); - } - } + for (int is = 0; is < nchan; is++) + { + occ_[iat][l][is].create(dim, dim); + occ_save_[iat][l][is].create(dim, dim); + } + } +} - for (int iw = 0; iw < cell.atoms[it].nw * npol; iw++) - { - const int iw0 = iw / npol; - const int ipol = iw % npol; - const int iwt = cell.itiaiw2iwt(it, ia, iw); - const int l = cell.atoms[it].iw2l[iw0]; - const int n = cell.atoms[it].iw2n[iw0]; - const int m = cell.atoms[it].iw2m[iw0]; - - if (n == 0) - { - corr_iwt_[iat][l][m][ipol] = iwt; - } - } +void OccupationMatrix::build_corr_iwt(const UnitCell& cell, + const int iat, + const int ia) +{ + const int it = cell.iat2it[iat]; + const int nwl = cell.atoms[it].nwl; + for (int L = 0; L <= nwl; L++) + { + corr_iwt_[iat][L].resize(2 * L + 1); + for (int m = 0; m < 2 * L + 1; m++) + { + corr_iwt_[iat][L][m].resize(npol_); + } + } + + for (int iw = 0; iw < cell.atoms[it].nw * npol_; iw++) + { + const int iw0 = iw / npol_; + const int ipol = iw % npol_; + const int iwt = cell.itiaiw2iwt(it, ia, iw); + const int l = cell.atoms[it].iw2l[iw0]; + const int n = cell.atoms[it].iw2n[iw0]; + const int m = cell.atoms[it].iw2m[iw0]; + + if (n == 0) + { + corr_iwt_[iat][l][m][ipol] = iwt; } } } @@ -265,29 +275,18 @@ void OccupationMatrix::write_save_to_flat(const UnitCell& cell, for (int I = 0; I < cell.atoms[T].na; I++) { const int iat = cell.itia2iat(T, I); - const int size = occ_save_[iat][target_l][0].nr * occ_save_[iat][target_l][0].nc; + const ModuleBase::matrix& save0 = occ_save_[iat][target_l][0]; + const int size = save0.nr * save0.nc; - if (nspin_ == 4) - { - for (int mm = 0; mm < size; mm++) - { - uom_save[index[iat] + mm] = occ_save_[iat][target_l][0].c[mm]; - } - } - else if (nspin_ == 1 || nspin_ == 2) + // all supported spin modes share the channel-0 layout + std::copy(save0.c, save0.c + size, uom_save.begin() + index[iat]); + + if (nspin_ == 2) { - for (int mm = 0; mm < size; mm++) - { - uom_save[index[iat] + mm] = occ_save_[iat][target_l][0].c[mm]; - } - if (nspin_ == 2) - { - const int half_size = uom_save.size() / 2; - for (int mm = 0; mm < size; mm++) - { - uom_save[half_size + index[iat] + mm] = occ_save_[iat][target_l][1].c[mm]; - } - } + const ModuleBase::matrix& save1 = occ_save_[iat][target_l][1]; + const int half_size = uom_save.size() / 2; + std::copy(save1.c, save1.c + size, + uom_save.begin() + half_size + index[iat]); } } } diff --git a/source/source_estate/occ_matrix.h b/source/source_estate/occ_matrix.h index 8a66d0de35..44d5f83785 100644 --- a/source/source_estate/occ_matrix.h +++ b/source/source_estate/occ_matrix.h @@ -106,6 +106,13 @@ class OccupationMatrix int npol() const { return npol_; } private: + /// allocate occ_ and occ_save_ matrices for one atom + void allocate_atom_matrices(const UnitCell& cell, int iat, int it); + + /// build corr_iwt_ mapping from the first radial channel to the + /// global orbital index for one atom + void build_corr_iwt(const UnitCell& cell, int iat, int ia); + std::vector>> occ_; std::vector>> occ_save_; std::vector>>> corr_iwt_; diff --git a/source/source_estate/test/CMakeLists.txt b/source/source_estate/test/CMakeLists.txt index a39b330c35..9ffbeb4e14 100644 --- a/source/source_estate/test/CMakeLists.txt +++ b/source/source_estate/test/CMakeLists.txt @@ -123,6 +123,12 @@ AddTest( SOURCES test_occ_mixer.cpp ../occ_mixer.cpp ../occ_matrix.cpp ) +AddTest( + TARGET MODULE_ESTATE_occ_matrix + LIBS parameter base device cell_info + SOURCES test_occ_matrix.cpp ../occ_matrix.cpp +) + AddTest( TARGET MODULE_ESTATE_occ_comput LIBS parameter base device diff --git a/source/source_estate/test/test_occ_matrix.cpp b/source/source_estate/test/test_occ_matrix.cpp new file mode 100644 index 0000000000..c08c07e5d4 --- /dev/null +++ b/source/source_estate/test/test_occ_matrix.cpp @@ -0,0 +1,317 @@ +#include "source_estate/occ_matrix.h" + +#include "source_cell/atom_spec.h" +#include "source_cell/magnetism.h" +#include "source_cell/unitcell.h" +#include "gtest/gtest.h" + +#include + +// UnitCell's constructor/destructor reference Magnetism; provide the +// minimal mock definitions (the real magnetism.cpp pulls heavy deps). +Magnetism::Magnetism() +{ + this->tot_mag = 0.0; + this->abs_mag = 0.0; +} +Magnetism::~Magnetism() +{ +} + +/*********************************************************************** + * Unit tests for OccupationMatrix. + * + * Covered: + * 1. init() layout for nspin=1, 2, 4 (channel count and matrix dims). + * 2. init() skips uncorrelated atom types (l_channel=-1). + * 3. init() builds corr_iwt with atom offsets and spinor strides from + * the first radial channel (n=0) only. + * 4. write_save_to_flat() for nspin=1/2/4, including the empty-buffer + * no-op path. + * + * A minimal UnitCell is built directly in each fixture; no private + * helpers or access-control tricks are used. + ***********************************************************************/ + +class OccMatrixTest : public ::testing::Test +{ + protected: + void SetUp() override + { + } + + /// Build a minimal UnitCell with `ntype` types. Each type has one + /// atom. `l_channels` gives the correlated l for each type (-1 for + /// uncorrelated). `nwl` and `l_nchi` are set to match `l_channels`. + void make_cell(const std::vector& l_channels, + const std::vector& nwls, + const std::vector>& l_nchis, + int npol) + { + cell.ntype = static_cast(l_channels.size()); + cell.nat = cell.ntype; + atoms_storage.resize(cell.ntype); + cell.atoms = atoms_storage.data(); + cell.iat2it.resize(cell.nat); + cell.iat2ia.resize(cell.nat); + cell.itia2iat.create(cell.ntype, cell.nat); + + for (int it = 0; it < cell.ntype; it++) + { + cell.atoms[it].na = 1; + cell.atoms[it].nwl = nwls[it]; + cell.atoms[it].l_nchi = l_nchis[it]; + cell.atoms[it].nw = 0; + for (int l = 0; l <= nwls[it]; l++) + { + cell.atoms[it].nw += (2 * l + 1) * l_nchis[it][l]; + } + cell.atoms[it].set_index(); + cell.iat2it[it] = it; + cell.iat2ia[it] = 0; + cell.itia2iat(it, 0) = it; + } + cell.set_iat2iwt(npol); + } + + UnitCell cell; + std::vector atoms_storage; +}; + +// ---------------------------------------------------------------------- +// init: nspin=1 layout +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, InitCollinearNspin1) +{ + make_cell({2}, {2}, {{0, 0, 1}}, 1); + OccupationMatrix occmat; + occmat.init(cell, {2}, 1, 1); + + EXPECT_EQ(occmat.nspin(), 1); + EXPECT_EQ(occmat.npol(), 1); + + const int iat = 0; + const int l = 2; + EXPECT_EQ(occmat.data()[iat][l].size(), 2u); + EXPECT_EQ(occmat.data()[iat][l][0].nr, 5); + EXPECT_EQ(occmat.data()[iat][l][0].nc, 5); + EXPECT_EQ(occmat.data()[iat][l][1].nr, 5); + EXPECT_EQ(occmat.data()[iat][l][1].nc, 5); + + EXPECT_EQ(occmat.data_save()[iat][l].size(), 2u); + EXPECT_EQ(occmat.data_save()[iat][l][0].nr, 5); + EXPECT_EQ(occmat.data_save()[iat][l][0].nc, 5); + EXPECT_EQ(occmat.data_save()[iat][l][1].nr, 5); + EXPECT_EQ(occmat.data_save()[iat][l][1].nc, 5); +} + +// ---------------------------------------------------------------------- +// init: nspin=2 layout +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, InitCollinearNspin2) +{ + make_cell({2}, {2}, {{0, 0, 1}}, 1); + OccupationMatrix occmat; + occmat.init(cell, {2}, 2, 1); + + EXPECT_EQ(occmat.nspin(), 2); + EXPECT_EQ(occmat.npol(), 1); + + const int iat = 0; + const int l = 2; + EXPECT_EQ(occmat.data()[iat][l].size(), 2u); + EXPECT_EQ(occmat.data()[iat][l][0].nr, 5); + EXPECT_EQ(occmat.data()[iat][l][0].nc, 5); + EXPECT_EQ(occmat.data()[iat][l][1].nr, 5); + EXPECT_EQ(occmat.data()[iat][l][1].nc, 5); + + EXPECT_EQ(occmat.data_save()[iat][l].size(), 2u); + EXPECT_EQ(occmat.data_save()[iat][l][0].nr, 5); + EXPECT_EQ(occmat.data_save()[iat][l][0].nc, 5); + EXPECT_EQ(occmat.data_save()[iat][l][1].nr, 5); + EXPECT_EQ(occmat.data_save()[iat][l][1].nc, 5); +} + +// ---------------------------------------------------------------------- +// init: nspin=4 layout +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, InitNoncollinearLayout) +{ + make_cell({2}, {2}, {{0, 0, 1}}, 2); + OccupationMatrix occmat; + occmat.init(cell, {2}, 4, 2); + + EXPECT_EQ(occmat.nspin(), 4); + EXPECT_EQ(occmat.npol(), 2); + + const int iat = 0; + const int l = 2; + EXPECT_EQ(occmat.data()[iat][l].size(), 1u); + EXPECT_EQ(occmat.data()[iat][l][0].nr, 10); + EXPECT_EQ(occmat.data()[iat][l][0].nc, 10); + + EXPECT_EQ(occmat.data_save()[iat][l].size(), 1u); + EXPECT_EQ(occmat.data_save()[iat][l][0].nr, 10); + EXPECT_EQ(occmat.data_save()[iat][l][0].nc, 10); +} + +// ---------------------------------------------------------------------- +// init: uncorrelated type is skipped and does not pollute the rest +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, InitSkipsUncorrelatedType) +{ + make_cell({2, -1}, {2, 2}, {{0, 0, 1}, {0, 0, 1}}, 1); + OccupationMatrix occmat; + occmat.init(cell, {2, -1}, 1, 1); + + // correlated type 0: allocated + EXPECT_EQ(occmat.data()[0][2].size(), 2u); + EXPECT_EQ(occmat.data()[0][2][0].nr, 5); + EXPECT_EQ(occmat.data_save()[0][2].size(), 2u); + + // uncorrelated type 1: no valid occupation blocks + EXPECT_EQ(occmat.data()[1][2].size(), 0u); + EXPECT_EQ(occmat.data_save()[1][2].size(), 0u); +} + +// ---------------------------------------------------------------------- +// init: corr_iwt points to the first radial channel (n=0) +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, InitBuildsCorrIwtFromFirstRadialChannel) +{ + // one type, one atom, l=2 with two radial channels (n=0 and n=1) + make_cell({2}, {2}, {{0, 0, 2}}, 1); + OccupationMatrix occmat; + occmat.init(cell, {2}, 1, 1); + + const int iat = 0; + const int l = 2; + const int m_size = 5; + + // global indices: atom 0 starts at 0, nw = 2*5 = 10 + // n=0: iw = 0..4, n=1: iw = 5..9 + for (int m = 0; m < m_size; m++) + { + EXPECT_EQ(occmat.corr_iwt(iat, l, m, 0), m); + } +} + +TEST_F(OccMatrixTest, InitBuildsCorrIwtForLaterNoncollinearAtom) +{ + // Atom 0 contributes one s orbital with two polarizations, so atom 1 + // starts at global index 2. Its d orbitals are interleaved by ipol. + make_cell({-1, 2}, {0, 2}, {{1}, {0, 0, 1}}, 2); + OccupationMatrix occmat; + occmat.init(cell, {-1, 2}, 4, 2); + + const int iat = 1; + const int l = 2; + for (int m = 0; m < 5; m++) + { + EXPECT_EQ(occmat.corr_iwt(iat, l, m, 0), 2 + 2 * m); + EXPECT_EQ(occmat.corr_iwt(iat, l, m, 1), 2 + 2 * m + 1); + } +} + +// ---------------------------------------------------------------------- +// write_save_to_flat: nspin=1 +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, WriteSaveToFlatNspin1) +{ + make_cell({2}, {2}, {{0, 0, 1}}, 1); + OccupationMatrix occmat; + occmat.init(cell, {2}, 1, 1); + + const int block = 25; + std::vector index = {0}; + std::vector uom_save(block, -1.0); + + for (int m = 0; m < block; m++) + { + occmat.data_save()[0][2][0].c[m] = 10.0 + m; + } + + occmat.write_save_to_flat(cell, {2}, index, uom_save); + + for (int m = 0; m < block; m++) + { + EXPECT_DOUBLE_EQ(uom_save[index[0] + m], 10.0 + m); + } +} + +// ---------------------------------------------------------------------- +// write_save_to_flat: nspin=2 split layout +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, WriteSaveToFlatNspin2SplitLayout) +{ + // two atoms, both correlated d + make_cell({2, 2}, {2, 2}, {{0, 0, 1}, {0, 0, 1}}, 1); + OccupationMatrix occmat; + occmat.init(cell, {2, 2}, 2, 1); + + const int block = 25; + const int half = cell.nat * block; + std::vector index = {0, block}; + std::vector uom_save(2 * half, -1.0); + + // four distinct data sets + for (int m = 0; m < block; m++) + { + occmat.data_save()[0][2][0].c[m] = 1000.0 + m; // atom0 up + occmat.data_save()[1][2][0].c[m] = 2000.0 + m; // atom1 up + occmat.data_save()[0][2][1].c[m] = 3000.0 + m; // atom0 down + occmat.data_save()[1][2][1].c[m] = 4000.0 + m; // atom1 down + } + + occmat.write_save_to_flat(cell, {2, 2}, index, uom_save); + + for (int m = 0; m < block; m++) + { + EXPECT_DOUBLE_EQ(uom_save[0 * block + m], 1000.0 + m); + EXPECT_DOUBLE_EQ(uom_save[1 * block + m], 2000.0 + m); + EXPECT_DOUBLE_EQ(uom_save[half + 0 * block + m], 3000.0 + m); + EXPECT_DOUBLE_EQ(uom_save[half + 1 * block + m], 4000.0 + m); + } +} + +// ---------------------------------------------------------------------- +// write_save_to_flat: nspin=4 full Pauli block +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, WriteSaveToFlatNspin4) +{ + make_cell({2}, {2}, {{0, 0, 1}}, 2); + OccupationMatrix occmat; + occmat.init(cell, {2}, 4, 2); + + const int block = 100; // 10 x 10 + std::vector index = {0}; + std::vector uom_save(block, -1.0); + + for (int m = 0; m < block; m++) + { + occmat.data_save()[0][2][0].c[m] = 500.0 + m; + } + + occmat.write_save_to_flat(cell, {2}, index, uom_save); + + for (int m = 0; m < block; m++) + { + EXPECT_DOUBLE_EQ(uom_save[m], 500.0 + m); + } +} + +// ---------------------------------------------------------------------- +// write_save_to_flat: empty buffer is a no-op +// ---------------------------------------------------------------------- +TEST_F(OccMatrixTest, WriteSaveToFlatEmptyBufferIsNoOp) +{ + make_cell({2}, {2}, {{0, 0, 1}}, 1); + OccupationMatrix occmat; + occmat.init(cell, {2}, 1, 1); + + std::vector index = {0}; + std::vector uom_save; + + occmat.write_save_to_flat(cell, {2}, index, uom_save); + EXPECT_TRUE(uom_save.empty()); +}