Skip to content

Commit

Permalink
UT for rearrange_smat
Browse files Browse the repository at this point in the history
  • Loading branch information
maki49 committed Aug 6, 2023
1 parent 707d7bc commit 77e38a8
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 27 deletions.
29 changes: 15 additions & 14 deletions source/module_ri/exx_symmetry.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,34 +11,33 @@
namespace ExxSym
{
std::vector<std::vector<std::complex<double>>> rearange_smat(
const int ikibz,
const std::vector<std::complex<double>> sloc_ikibz,
const int nbasis,
const Parallel_2D& p2d,
std::vector<std::vector<int>>& isym_iat_rotiat,
std::vector<std::map<int, ModuleBase::Vector3<double>>> kstars,
const UnitCell& ucell)
std::map<int, ModuleBase::Vector3<double>> kstar_ibz,
const UnitCell& ucell,
const bool col_inside,
std::ofstream& ofs_running)
{
ModuleBase::TITLE("ExxSym", "rearange_smat");
// add title
std::vector<std::vector<std::complex<double>>> kvd_sloc; //result

//1. get the full smat
//get the full smat
// note: the globalfunc-major means outside, while scalapack's major means inside
bool col_inside = !ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER();
// bool col_inside = !ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER();
//sfull is symmetric but sloc not, so row/col major should be considered
std::vector<std::complex<double>> sfull_ikbz = get_full_smat(sloc_ikibz, nbasis, p2d,
col_inside);

//2. get invmap: iat1 to iat0
std::vector<int> isym_rotiat_iat = ModuleSymmetry::Symmetry::invmap(isym_iat_rotiat[0].data(), ucell.nat);

for (auto& isym_kvecd : kstars[ikibz])
for (auto& isym_kvecd : kstar_ibz)
{
int isym = isym_kvecd.first;
//get invmap : iat1 to iat0
std::vector<int> rotiat_iat = ModuleSymmetry::Symmetry::invmap(isym_iat_rotiat[isym].data(), ucell.nat);
// 1 symmetry operation - may more than one kvec_d ?? check it !!!
// ModuleBase::Vector3<double> kvds = isym_kvecd.second;
// set rearanged sloc_ik
std::vector<std::complex<double>> sloc_ik(p2d.get_local_size(), 0);
std::vector<std::complex<double>> sloc_ik(p2d.get_local_size());
for (int mu = 0;mu < p2d.get_row_size();++mu)
{
int gr = p2d.local2global_row(mu);
Expand All @@ -47,7 +46,7 @@ namespace ExxSym
int gc = p2d.local2global_col(nu);
int iat1 = ucell.iwt2iat[gc];
int iw = ucell.iwt2iw[gc];//orb-index of iat1, the same as iat0
int iat0 = isym_rotiat_iat[iat1]; //g(iat0)=iat1
int iat0 = rotiat_iat[iat1]; //g(iat0)=iat1
assert(ucell.iat2it[iat1] == ucell.iat2it[iat0]);//check for the same it
int gc0 = ucell.iat2iwt[iat0] + iw;
if (col_inside)
Expand All @@ -67,6 +66,7 @@ namespace ExxSym
const Parallel_2D& p2d,
const bool col_inside)
{
ModuleBase::TITLE("ExxSym", "get_full_mat");
#ifdef __MPI
std::vector<std::complex<double>> fullmat(nbasis * nbasis, 0);
if (col_inside)
Expand All @@ -92,6 +92,7 @@ namespace ExxSym
const int& nbasis,
const int& nbands)
{
ModuleBase::TITLE("ExxSym", "restore_psik_lapack");
int nkstar = sloc_ik.size();
psi::Psi<std::complex<double>, psi::DEVICE_CPU> c_k(nkstar, nbands, nbasis);
psi::Psi<std::complex<double>, psi::DEVICE_CPU> tmpSc(1, nbands, nbasis);
Expand Down Expand Up @@ -146,6 +147,7 @@ namespace ExxSym
const int& nbands,
const Parallel_Orbitals& pv)
{
ModuleBase::TITLE("ExxSym", "restore_psik_scalapack");
int nkstar = sloc_ik.size();
psi::Psi<std::complex<double>, psi::DEVICE_CPU> c_k(nkstar, pv.ncol_bands, pv.get_row_size());
psi::Psi<std::complex<double>, psi::DEVICE_CPU> tmpSc(1, pv.ncol_bands, pv.get_row_size());
Expand Down Expand Up @@ -197,5 +199,4 @@ namespace ExxSym
return c_k;
}
#endif

}
13 changes: 7 additions & 6 deletions source/module_ri/exx_symmetry.h
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
#pragma once
#include <vector>
#include "module_basis/module_ao/parallel_orbitals.h"
#include "module_cell/unitcell.h"
#include "module_base/vector3.h"
#include "module_psi/psi.h"

#include "module_cell/unitcell.h"
namespace ExxSym
{
/// @brief Rearrange the $\nu$ index of overlap matrices of ibz-kpoints $S_{\mu,\nu}(gk)$ according to the symmetry operations
Expand All @@ -12,17 +12,18 @@ namespace ExxSym
/// @param nbasis [in] global number of basis
/// @param p2d [in]2d parallel info
/// @param isym_iat_rotiat [in] atom index map corresponding to each symmetry operation
/// @param kstars [in] symmetry-equal k points to each ibz-kpont: [iksibz][isym][kvec_d]
/// @param kstar_ibz [in] symmetry-equal k points to current ibz-kpont: [isym][kvec_d]
/// @param ucell [in] unitcell
/// @return local S(k) for each k in kstars[ikibz]
std::vector<std::vector<std::complex<double>>> rearange_smat(
const int ikibz,
const std::vector<std::complex<double>> sloc_ikibz,
const int nbasis,
const Parallel_2D& p2d,
std::vector<std::vector<int>>& isym_iat_rotiat,
std::vector<std::map<int, ModuleBase::Vector3<double>>> kstars,
const UnitCell& ucell);
std::map<int, ModuleBase::Vector3<double>> kstar_ibz,
const UnitCell& ucell,
const bool col_inside,
std::ofstream& ofs_running);

/// @brief restore c_k from c_gk: $c_k=\tilde{S}^{-1}(k)S(gk)c_{gk}$
/// @param ikibz [in] ibz-kpoint index
Expand Down
109 changes: 102 additions & 7 deletions source/module_ri/test/exx_symmetry_test.cpp
Original file line number Diff line number Diff line change
@@ -1,18 +1,32 @@
#include "gtest/gtest.h"
#include "module_cell/module_symmetry/symmetry.h"
#include "../exx_symmetry.h"
#include "module_cell/setup_nonlocal.h"
#include <map>
#include <tuple>


Magnetism::Magnetism() {}
Magnetism::~Magnetism() {}
InfoNonlocal::InfoNonlocal() {}
InfoNonlocal::~InfoNonlocal() {}
UnitCell::UnitCell() {
iat2it = nullptr;
iwt2iat = nullptr;
iwt2iw = nullptr;
}
UnitCell::~UnitCell()
{
delete[] iat2it;
delete[] iwt2iat;
delete[] iwt2iw;
}
class SymExxTest : public testing::Test
{
protected:
int dsize = 1;
int my_rank = 0;
std::ofstream ofs_running;
Parallel_Orbitals pv;

UnitCell ucell;
int nbasis = 0;
//cases
std::map<std::vector<int>, std::vector<int>> invmap_cases = {
{{3, 2, 1, 0}, {3, 2, 1, 0}},
Expand Down Expand Up @@ -76,6 +90,49 @@ class SymExxTest : public testing::Test
else
sl[j * lr + i] = sg[pv.local2global_col(j) * gr + pv.local2global_row(i)];
}
static std::vector<int> invmap(const int* map, const int& size)
{
std::vector<int> invf(size);
for (size_t i = 0; i < size; ++i) invf[map[i]] = i;
return invf;
}
static std::vector<int> mapmul(const int* map1, const int* map2, const int& size)
{
std::vector<int> f2f1(size); // f1 first
for (size_t i = 0; i < size; ++i) f2f1[i] = map2[map1[i]];
return f2f1;
}
void setup_cell_index(UnitCell& ucell, std::vector<std::pair<int, int>>& na_nw)
{
ucell.ntype = na_nw.size();
// set nat and nbasis
ucell.nat = 0;
this->nbasis = 0;
for (auto& aw : na_nw) { ucell.nat += aw.first; this->nbasis += aw.first * aw.second; }

ucell.iat2it = new int[ucell.nat];
ucell.iat2iwt.resize(ucell.nat);
ucell.iwt2iat = new int[this->nbasis];
ucell.iwt2iw = new int[this->nbasis];
// set index maps
int iat = 0;
int iwt = 0;
for (int it = 0; it < ucell.ntype; ++it)
for (int ia = 0; ia < na_nw[it].first; ++ia)
{
ucell.iat2it[iat] = it;
ucell.iat2iwt[iat] = iwt;
for (int iw = 0; iw < na_nw[it].second; ++iw)
{
ucell.iwt2iat[iwt] = iat;
ucell.iwt2iw[iwt] = iw;
++iwt;
}
++iat;
}
assert(iat == ucell.nat);
assert(iwt == this->nbasis);
}
#ifdef __MPI
void SetUp() override
{
Expand All @@ -96,7 +153,7 @@ TEST_F(SymExxTest, invmap)
{
for (auto c : invmap_cases)
{
std::vector<int> invf = ModuleSymmetry::Symmetry::invmap(c.first.data(), c.first.size());
std::vector<int> invf = invmap(c.first.data(), c.first.size());
EXPECT_EQ(invf, c.second);
}
}
Expand All @@ -105,11 +162,50 @@ TEST_F(SymExxTest, mapmul)
{
for (auto c : mapmul_cases)
{
std::vector<int> f2f1 = ModuleSymmetry::Symmetry::mapmul(std::get<0>(c).data(), std::get<1>(c).data(), std::get<0>(c).size());
std::vector<int> f2f1 = mapmul(std::get<0>(c).data(), std::get<1>(c).data(), std::get<0>(c).size());
EXPECT_EQ(f2f1, std::get<2>(c));
}
}

TEST_F(SymExxTest, rearange_smat)
{
// case
std::vector<std::pair<int, int>> atoms = { {2, 1}, {3, 3} };
std::vector<std::vector<int>> isym_iat_rotiat = { {0, 1, 2, 3, 4}, {0,1,3,4,2}, {0,1,4,2,3}, {1,0,2,3,4}, {1,0,3,4,2}, {1,0,4,2,3} };
ModuleBase::Vector3<double> kvd(0, 0, 0); //only one ibzkpt and its coordinate is important in this function
std::map<int, ModuleBase::Vector3<double>> kstar_ibz = { {0, kvd}, {1, kvd}, {2, kvd}, {3, kvd}, {4, kvd}, {5, kvd} };

// set ucell and 2d division
UnitCell ucell;
this->setup_cell_index(ucell, atoms);

EXPECT_EQ(ucell.nat, 5);
EXPECT_EQ(this->nbasis, 11);
int ikibz = 0;
this->set2d(nbasis, nbasis);
// generate global symmitric matrix and copy to local
std::vector<std::complex<double>> sfull_gk(nbasis * nbasis);
this->set_int_posisym(sfull_gk.data(), nbasis, 0);
std::vector<std::complex<double>> sloc_gk(pv.get_local_size());
this->copy_from_global(sfull_gk.data(), sloc_gk.data(), nbasis, nbasis, pv.get_row_size(), pv.get_col_size(), false, false);

// run (row-major)
std::vector<std::vector<std::complex<double>>> sloc_ks = ExxSym::rearange_smat(sloc_gk, nbasis, pv, isym_iat_rotiat, kstar_ibz, ucell, false, ofs_running);
// check
for (int isym = 0;isym < isym_iat_rotiat.size();++isym)
{
std::vector<std::complex<double>> sfull_ik = ExxSym::get_full_smat(sloc_ks[isym], nbasis, pv, false);
for (int i = 0;i < pv.get_row_size();i++)
{
int iwt0 = pv.local2global_row(i);
int iat0 = ucell.iwt2iat[iwt0];
int iw = ucell.iwt2iw[iwt0];
int iat1 = isym_iat_rotiat[isym][iat0];
int iwt1 = ucell.iat2iwt[iat1] + iw;
EXPECT_EQ(sfull_gk[iwt0 * nbasis + iwt0], sfull_ik[iwt1 * nbasis + iwt0]);
}
}
}
#ifdef __MPI
TEST_F(SymExxTest, restore_psik)
{
Expand All @@ -126,7 +222,6 @@ TEST_F(SymExxTest, restore_psik)
// generate global symmitric matrix and copy to local
std::vector<std::complex<double>> sfull_gk(nbasis * nbasis);
this->set_int_posisym(sfull_gk.data(), nbasis, nkstar);

std::vector<std::complex<double>> sloc_gk(pv.get_local_size());
this->copy_from_global(sfull_gk.data(), sloc_gk.data(), nbasis, nbasis, pv.get_row_size(), pv.get_col_size(), false, false);

Expand Down

0 comments on commit 77e38a8

Please sign in to comment.