37 uint32_t j = idx__ & 0xFFF;
39 uint32_t i = idx__ >> 12;
41 RTE_ASSERT(j < (uint32_t)
z_columns_[i].z.size());
50 PROFILE(
"fft::Gvec::find_z_columns");
53 non_zero_columns.
zero();
57 auto add_new_column = [&](
int i,
int j) {
59 if (non_zero_columns(i, j)) {
63 std::vector<int> zcol;
66 int zmax = fft_box__[2] - 1;
69 zmax = fft_box__.
limits(2).second;
72 for (
int iz = 0; iz <= zmax; iz++) {
78 if (vgk.length() <= Gmax__) {
86 num_gvec_ +=
static_cast<int>(zcol.size());
88 non_zero_columns(i, j) = 1;
92 if (mi >= fft_box__.
limits(0).first && mi <= fft_box__.
limits(0).second &&
93 mj >= fft_box__.
limits(1).first && mj <= fft_box__.
limits(1).second) {
94 non_zero_columns(mi, mj) = 1;
100 PROFILE_START(
"fft::Gvec::find_z_columns|add");
106 add_new_column(i, j);
112 for (
int j = fft_box__.
limits(1).first; j <= fft_box__.
limits(1).second; j++) {
113 add_new_column(i, j);
116 PROFILE_STOP(
"fft::Gvec::find_z_columns|add");
120 for (
size_t i = 0; i <
z_columns_.size(); i++) {
128 PROFILE_START(
"fft::Gvec::find_z_columns|sym");
131 auto lat_sym = sirius::find_lat_sym(this->unit_cell_lattice_vectors(), this->
sym_tol_);
133 std::fill(non_zero_columns.at(sddk::memory_t::host), non_zero_columns.at(sddk::memory_t::host) + non_zero_columns.
size(), -1);
134 for (
int i = 0; i < static_cast<int>(
z_columns_.size()); i++) {
138 std::vector<z_column_descriptor> z_columns_tmp;
140 auto remove_incomplete = [
this, &z_columns_tmp, &lat_sym, &non_zero_columns](
int i)
142 int z_min = (2 << 20);
143 int z_max = -(2 << 20);
144 for (
int iz = 0; iz < static_cast<int>(
z_columns_[i].z.size()); iz++) {
149 for (
int iz = 0; iz < static_cast<int>(
z_columns_[i].z.size()); iz++) {
150 bool found_for_all_sym{
true};
154 for (
auto& R: lat_sym) {
156 auto G1 = r3::dot(G, R);
158 if (G1[0] == 0 && G1[1] == 0) {
159 G1[2] = std::abs(G1[2]);
162 int i1 = non_zero_columns(G1[0], G1[1]);
165 i1 = non_zero_columns(G1[0], G1[1]);
168 s <<
"index of z-column is not found" << std::endl
169 <<
" G : " << G << std::endl
176 found_for_all_sym = found_for_all_sym && found;
179 if (found_for_all_sym) {
188 for (
int i = 0; i < static_cast<int>(
z_columns_.size()); i++) {
189 remove_incomplete(i);
195 num_gvec_ +=
static_cast<int>(zc.z.size());
198 PROFILE_STOP(
"fft::Gvec::find_z_columns|sym");
200 PROFILE_START(
"fft::Gvec::find_z_columns|sort");
205 PROFILE_STOP(
"fft::Gvec::find_z_columns|sort");
213 std::vector<std::vector<z_column_descriptor>> zcols_local(comm().size());
217 for (
int rank = 0; rank < comm().
size(); rank++) {
221 zcols_local[rank].push_back(
z_columns_[icol]);
232 std::vector<int> ranks;
233 for (
int i = n; i < static_cast<int>(
z_columns_.size()); i++) {
236 ranks.resize(comm().
size());
237 std::iota(ranks.begin(), ranks.end(), 0);
240 auto rank_with_min_gvec = std::min_element(ranks.begin(), ranks.end(), [
this](
const int& a,
const int& b) {
241 return gvec_distr_.counts[a] < gvec_distr_.counts[b];
245 zcols_local[*rank_with_min_gvec].push_back(
z_columns_[i]);
251 ranks.erase(rank_with_min_gvec);
258 for (
int rank = 0; rank < comm().
size(); rank++) {
264 for (
int rank = 0; rank < comm().
size(); rank++) {
268 RTE_THROW(
"wrong number of G-vectors");
281 PROFILE(
"fft::Gvec::find_gvec_shells");
283 auto lat_sym = sirius::find_lat_sym(this->unit_cell_lattice_vectors(), this->
sym_tol_);
294 auto G = gvec<index_domain_t::global>(ig);
295 for (
auto& R: lat_sym) {
296 auto G1 = r3::dot(G, R);
302 RTE_THROW(
"symmetry-related G-vector is not found");
312 s <<
"Error in G-vector shell index" << std::endl
313 <<
" G : " << G << std::endl
314 <<
" rotated G : " << G1 << std::endl
316 <<
" rotated G shell index : " <<
gvec_shell_[ig1] << std::endl
317 <<
" length of G : " << gc.length() << std::endl
318 <<
" length of rotated G : " << gc1.length() << std::endl
319 <<
" length difference : " << std::abs(gc.length() - gc1.length());
329 RTE_THROW(
"wrong G-vector shell");
339 auto g = gvec_cart<index_domain_t::global>(ig).length();
349 std::map<int, std::vector<int>> gshmap;
350 for (
int igloc = 0; igloc < this->
count(); igloc++) {
352 if (gshmap.count(igsh) == 0) {
353 gshmap[igsh] = std::vector<int>();
355 gshmap[igsh].push_back(igloc);
361 for (
auto it = gshmap.begin(); it != gshmap.end(); ++it) {
362 int igsh = it->first;
364 for (
auto igloc: it->second) {
377 for (
int igloc = 0; igloc <
count(); igloc++) {
378 int ig =
offset() + igloc;
380 for (
int x : {0, 1, 2}) {
381 gvec_(x, igloc) = G[x];
399 for (
int igloc = 0; igloc <
count(); igloc++) {
402 for (
int x : {0, 1, 2}) {
410 gvec_tp_(igloc, 0) = gs[1];
411 gvec_tp_(igloc, 1) = gs[2];
414 gkvec_tp_(igloc, 0) = gks[1];
415 gkvec_tp_(igloc, 1) = gks[2];
422 PROFILE(
"fft::Gvec::init");
430 std::fill(gvec_index_by_xy_.at(sddk::memory_t::host), gvec_index_by_xy_.at(sddk::memory_t::host) + gvec_index_by_xy_.
size(),
436 for (
size_t i = 0; i <
z_columns_.size(); i++) {
441 for (
size_t j = 0; j <
z_columns_[i].z.size(); j++) {
446 RTE_THROW(
"wrong G-vector count");
450 auto gv = gvec<index_domain_t::global>(ig);
453 s <<
"wrong G-vector index: ig=" << ig <<
" gv=" << gv <<
" index_by_gvec(gv)=" <<
index_by_gvec(gv);
460 if (g0[0] || g0[1] || g0[2]) {
461 RTE_THROW(
"first G-vector is not zero");
477 if (ig >= 0 && ig <
count()) {
481 s <<
"local G-vector index is not found" << std::endl
482 <<
" G-vector: " << G << std::endl
483 <<
" G-vector index in base distribution : " <<
gvec_base_->
offset() + igloc << std::endl
485 <<
" G-vector index in new distribution : " <<
index_by_gvec(G) << std::endl
486 <<
" offset in G-vector index for this rank: " <<
offset() << std::endl
487 <<
" local number of G-vectors for this rank: " <<
count();
499 auto v = g1__ - g2__;
508 s <<
"wrong index of G-G' vector" << std::endl
509 <<
" G: " << g1__ << std::endl
510 <<
" G': " << g2__ << std::endl
511 <<
" G - G': " << v << std::endl
515 return std::make_pair(idx,
conj);
521 if (reduced() && G__[0] == 0 && G__[1] == 0 && G__[2] < 0) {
524 int ig0 = gvec_index_by_xy_(0, G__[0], G__[1]);
529 int icol = gvec_index_by_xy_(1, G__[0], G__[1]) & 0xFFFFF;
535 int col_size = gvec_index_by_xy_(1, G__[0], G__[1]) >> 20;
553 int offs = (z0 >= 0) ? z0 : z0 + col_size;
556 RTE_ASSERT(ig >= 0 && ig <
num_gvec());
564 if (comm__.
rank() == source__) {
565 ::sirius::fft::serialize(s, gv_src__);
568 s.send_recv(comm__, source__, dest__);
570 Gvec gv(gv_src__.comm());
572 if (comm__.
rank() == dest__) {
573 ::sirius::fft::deserialize(s, gv);
578void Gvec_fft::build_fft_distr()
632 for (
int x : {0, 1, 2}) {
644 , comm_fft_(comm_fft__)
645 , comm_ortho_fft_(comm_ortho_fft__)
649 s <<
"wrong size of communicators" << std::endl
652 <<
" gvec_.comm().size() = " <<
gvec_.comm().
size();
666Gvec_shells::Gvec_shells(Gvec
const& gvec__)
667 : comm_(gvec__.comm())
670 PROFILE(
"fft::Gvec_shells");
672 a2a_send_ = mpi::block_data_descriptor(comm_.size());
673 a2a_recv_ = mpi::block_data_descriptor(comm_.size());
680 for (
int igloc = 0; igloc <
gvec_.
count(); igloc++) {
683 a2a_send_.counts[spl_num_gsh_.location(igsh).ib]++;
685 a2a_send_.calc_offsets();
688 RTE_THROW(
"wrong number of G-vectors");
691 for (
int r = 0; r < comm_.size(); r++) {
697 if (spl_num_gsh_.location(igsh).ib == comm_.rank()) {
698 a2a_recv_.counts[r]++;
702 a2a_recv_.calc_offsets();
704 int ng = gvec_count_remapped();
705 comm_.allreduce(&ng, 1);
707 RTE_THROW(
"wrong number of G-vectors");
711 gvec_remapped_ = sddk::mdarray<int, 2>(3, gvec_count_remapped(), sddk::memory_t::host,
"gvec_remapped_");
712 gvec_shell_remapped_ = sddk::mdarray<int, 1>(gvec_count_remapped(), sddk::memory_t::host,
"gvec_shell_remapped_");
713 std::vector<int> counts(comm_.size(), 0);
714 for (
int r = 0; r < comm_.size(); r++) {
719 if (spl_num_gsh_.location(igsh).ib == comm_.rank()) {
720 for (
int x : {0, 1, 2}) {
721 gvec_remapped_(x, a2a_recv_.offsets[r] + counts[r]) = G[x];
723 gvec_shell_remapped_(a2a_recv_.offsets[r] + counts[r]) = igsh;
728 for (
int ig = 0; ig < gvec_count_remapped(); ig++) {
729 idx_gvec_[gvec_remapped(ig)] = ig;
732 for (
int igloc = 0; igloc < this->gvec_count_remapped(); igloc++) {
733 auto G = this->gvec_remapped(igloc);
734 if (this->index_by_gvec(G) != igloc) {
735 RTE_THROW(
"Wrong remapped index of G-vector");
737 int igsh = this->gvec_shell_remapped(igloc);
738 if (igsh != this->
gvec().shell(G)) {
739 RTE_THROW(
"Wrong remapped shell of G-vector");
744void serialize(serializer& s__, Gvec
const& gv__)
746 serialize(s__, gv__.vk_);
747 serialize(s__, gv__.Gmax_);
748 serialize(s__, gv__.lattice_vectors_);
749 serialize(s__, gv__.reduce_gvec_);
750 serialize(s__, gv__.bare_gvec_);
751 serialize(s__, gv__.num_gvec_);
752 serialize(s__, gv__.num_gvec_shells_);
753 serialize(s__, gv__.gvec_full_index_);
754 serialize(s__, gv__.gvec_shell_);
755 serialize(s__, gv__.gvec_shell_len_);
756 serialize(s__, gv__.gvec_index_by_xy_);
757 serialize(s__, gv__.z_columns_);
758 serialize(s__, gv__.gvec_distr_);
759 serialize(s__, gv__.zcol_distr_);
760 serialize(s__, gv__.gvec_base_mapping_);
761 serialize(s__, gv__.offset_);
762 serialize(s__, gv__.count_);
765void deserialize(serializer& s__, Gvec& gv__)
767 deserialize(s__, gv__.vk_);
768 deserialize(s__, gv__.Gmax_);
769 deserialize(s__, gv__.lattice_vectors_);
770 deserialize(s__, gv__.reduce_gvec_);
771 deserialize(s__, gv__.bare_gvec_);
772 deserialize(s__, gv__.num_gvec_);
773 deserialize(s__, gv__.num_gvec_shells_);
774 deserialize(s__, gv__.gvec_full_index_);
775 deserialize(s__, gv__.gvec_shell_);
776 deserialize(s__, gv__.gvec_shell_len_);
777 deserialize(s__, gv__.gvec_index_by_xy_);
778 deserialize(s__, gv__.z_columns_);
779 deserialize(s__, gv__.gvec_distr_);
780 deserialize(s__, gv__.zcol_distr_);
781 deserialize(s__, gv__.gvec_base_mapping_);
782 deserialize(s__, gv__.offset_);
783 deserialize(s__, gv__.count_);
Helper class to create FFT grids of given sizes and compute indices in space- and frequency domains.
const std::pair< int, int > & limits(int idim__) const
Limits of a given dimension.
int freq_by_coord(int x__) const
Get frequency by coordinate.
mpi::Communicator const & comm_fft() const
Return FFT communicator.
int count() const
Local number of G-vectors for FFT-friendly distribution for this rank.
mpi::block_data_descriptor gvec_fft_slab_
Distribution of G-vectors inside FFT-friendly "fat" slab.
mpi::Communicator const & comm_ortho_fft() const
Return a communicator that is orthogonal to the FFT communicator.
mpi::Communicator const & comm_ortho_fft_
Communicator which is orthogonal to FFT communicator.
sddk::mdarray< int, 2 > rank_map_
Mapping of MPI ranks used to split G-vectors to a 2D grid.
sddk::mdarray< double, 2 > gkvec_cart_array_
Cartesian coordinaes of a local set of G+k-vectors.
mpi::Communicator const & comm_fft_
Communicator for the FFT.
void update_gkvec_cart()
Update Cartesian coordinates after a change in lattice vectors.
sddk::mdarray< int, 2 > gvec_array_
Lattice coordinates of a local set of G-vectors.
mpi::block_data_descriptor gvec_distr_fft_
Distribution of G-vectors for FFT.
Gvec const & gvec_
Reference to the G-vector instance.
Gvec const & gvec() const
Return the original (not reshuffled) G-vector class.
void pile_gvec()
Stack together the G-vector slabs to make a larger ("fat") slab for a FFT driver.
int num_zcol_local_
Local number of z-columns.
A set of G-vectors for FFTs and G+k basis functions.
int offset_
Offset in the global index for the local part of G-vectors.
int gvec_count(int rank__) const
Number of G-vectors for a fine-grained distribution.
int num_zcol() const
Return global number of z-columns.
sddk::mdarray< double, 1 > gvec_shell_len_
Radii (or lengths) of G-vector shells in a.u.^-1.
mpi::block_data_descriptor zcol_distr_
Fine-grained distribution of z-columns.
sddk::mdarray< uint32_t, 1 > gvec_full_index_
Mapping between G-vector index [0:num_gvec_) and a full index.
sddk::mdarray< double, 2 > gvec_cart_
Cartiesian coordinaes of a local set of G-vectors.
int num_gvec() const
Return the total number of G-vectors within the cutoff.
std::vector< double > gvec_shell_len_local_
Radii of G-vector shells in the local index counting [0, num_gvec_shells_local)
int offset() const
Offset (in the global index) of G-vectors for a fine-grained distribution for a current MPI rank.
sddk::mdarray< int, 2 > gvec_
Lattice coordinates of a local set of G-vectors.
std::vector< z_column_descriptor > z_columns_
Global list of non-zero z-columns.
sddk::mdarray< int, 1 > gvec_shell_
Index of the shell to which the given G-vector belongs.
sddk::mdarray< int, 1 > gvec_base_mapping_
Mapping between current and base G-vector sets.
void distribute_z_columns()
Distribute z-columns between MPI ranks.
sddk::mdarray< double, 1 > gvec_len_
Length of the local fraction of G-vectors.
r3::vector< double > vk_
k-vector of G+k.
int zcol_offset(int rank__) const
Offset in the global index of z-columns for a given rank.
int count() const
Number of G-vectors for a fine-grained distribution for the current MPI rank.
void find_gvec_shells()
Find a list of G-vector shells.
int gvec_offset(int rank__) const
Offset (in the global index) of G-vectors for a fine-grained distribution.
bool bare_gvec_
True if this a list of G-vectors without k-point shift.
int count_
Local number of G-vectors.
int num_shells() const
Return number of G-vector shells.
double sym_tol_
Symmetry tolerance of the real-space lattice.
void init_gvec_cart_local()
Initialize Cartesian coordinates of the local fraction of G-vectors.
int num_gvec_shells_local_
Local number of G-vector shells for the local number of G-vectors.
int zcol_count(int rank__) const
Number of z-columns for a fine-grained distribution.
mpi::block_data_descriptor gvec_distr_
Fine-grained distribution of G-vectors.
r3::vector< int > gvec(int ig__) const
Return G vector in fractional coordinates.
sddk::mdarray< double, 2 > gkvec_cart_
Cartesian coordinaes of a local set of G+k-vectors.
bool reduce_gvec_
Indicates that G-vectors are reduced by inversion symmetry.
int shell(int ig__) const
Return index of the G-vector shell by the G-vector index.
int index_by_gvec(r3::vector< int > const &G__) const
Return a global G-vector index in the range [0, num_gvec) by the G-vector.
int num_zcol_local_
Local number of z-columns.
int num_gvec_shells_
Number of G-vector shells (groups of G-vectors with the same length).
Gvec const * gvec_base_
Set of G-vectors on which the current G-vector distribution can be based.
double Gmax_
Cutoff for |G+k| vectors.
double shell_len(int igs__) const
Return length of the G-vector shell.
sddk::mdarray< double, 2 > gkvec_
Lattice coordinates of a local set of G+k-vectors.
auto const & gvec_local() const
Return local list of G-vectors.
r3::vector< int > gvec_by_full_index(uint32_t idx__) const
Return corresponding G-vector for an index in the range [0, num_gvec).
std::vector< int > gvec_shell_idx_local_
Mapping between local index of G-vector and local G-shell index.
void init(fft::Grid const &fft_grid)
Initialize everything.
r3::matrix< double > lattice_vectors_
Reciprocal lattice vectors.
int num_gvec_
Total number of G-vectors.
void init_gvec_local()
Initialize lattice coordinates of the local fraction of G-vectors.
void find_z_columns(double Gmax__, fft::Grid const &fft_box__)
Find z-columns of G-vectors inside a sphere with Gmax radius.
MPI communicator wrapper.
int size() const
Size of the communicator (number of ranks).
void allreduce(T *buffer__, int count__) const
Perform the in-place (the output buffer is used as the input buffer) all-to-all reduction.
int rank() const
Rank of MPI process inside communicator.
Simple implementation of 3d vector.
void zero(memory_t mem__, size_t idx0__, size_t n__)
Zero n elements starting from idx0.
size_t size() const
Return total size (number of elements) of the array.
Serialize and deserialize objects.
Declaration and implementation of Gvec class.
Crystal lattice functions.
auto spherical_coordinates(vector< double > vc)
Transform Cartesian coordinates [x,y,z] to spherical coordinates [r,theta,phi].
Namespace of the SIRIUS library.
strong_type< int, struct __block_id_tag > block_id
ID of the block.
strong_type< int, struct __n_blocks_tag > n_blocks
Number of blocks to which the global index is split.
auto conj(double x__)
Return complex conjugate of a number. For a real value this is the number itself.
Serializer for simple data structures.
Descriptor of the z-column (x,y fixed, z varying) of the G-vectors.
int y
Y-coordinate (can be negative and positive).
int x
X-coordinate (can be negative and positive).