@@ -10,6 +10,14 @@ namespace ModulePW
1010//
1111// Distribute planewaves in reciprocal space to coreors.
1212// Firstly, divide the sphere in reciprocal space into sticks, which are vertical to x-y plane.
13+ //
14+ // Example
15+ // | ---- ixy increasing ---> | ---- ixy increasing ---> |...
16+ // index of sticks 0, 1, 2, ..., nst_per[0]-1, nst_per[0], ..., nst_per[1], ...
17+ // |___________________________|____________________________|___
18+ // ip 0 1 ...
19+ // npw approximate equal to npw approximate equal to...
20+ //
1321// Secondly, distribute these sticks to coreors.
1422// Known: G, GT, GGT, ny, nx, nz, poolnproc, poolrank, ggecut
1523// output: ig2isz[ig], istot2ixy[is], ixy2istot[nxy], is2ixy[is], ixy2ip[ixy], startnsz_per[ip], nst_per[ip], nst
@@ -49,6 +57,8 @@ void PW_Basis::distribution_method1()
4957 std::cout << ' \n ' ;
5058 }
5159 // ------------------------------------------------------------
60+ #ifdef __MPI
61+ // Parallel line
5262
5363 // (2) Collect the x, y indexs, length, bottom of the sticks.
5464 int * st_i = new int [this ->nstot ]; // x or x + nx (if x < 0) of stick.
@@ -61,8 +71,6 @@ void PW_Basis::distribution_method1()
6171 std::cout << " \n The second step done\n " ;
6272 // ------------------------------------------------------------
6373
64- #ifdef __MPI
65- // Parallel line
6674 // (3) Distribute sticks to cores.
6775 int *npw_per = new int [this ->poolnproc ]; // number of planewaves on each core.
6876 int *nst_per = new int [this ->poolnproc ]; // number of sticks on each core.
@@ -78,7 +86,7 @@ void PW_Basis::distribution_method1()
7886 {
7987 this ->ixy2ip [ixy] = -1 ; // meaning this stick has not been distributed or there is no stick on (x, y).
8088 }
81- this ->divide_sticks (tot_npw, st_i, st_j, st_length, npw_per, nst_per);
89+ this ->divide_sticks (st_i, st_j, st_length, npw_per, nst_per);
8290 delete[] st_length;
8391
8492 // for test -----------------------------------------------------------------------------
@@ -181,72 +189,6 @@ void PW_Basis::distribution_method1()
181189 return ;
182190}
183191
184- //
185- // (1) We count the total number of planewaves (tot_npw) and sticks (this->nstot) here.
186- // Meanwhile, we record the number of planewaves on (x, y) in st_length2D, and store the smallest z-coordinate of each stick in st_bottom2D,
187- // so that we can scan a much smaller area in step(2).
188- // known: nx, ny, nz, ggecut, GGT
189- // output: tot_npw, this->nstot, st_length2D, st_bottom2D
190- //
191- void PW_Basis::count_pw_st (
192- int &tot_npw, // total number of planewaves.
193- int * st_length2D, // the number of planewaves that belong to the stick located on (x, y).
194- int * st_bottom2D // the z-coordinate of the bottom of stick on (x, y).
195- )
196- {
197- int ibox[3 ] = {0 , 0 , 0 }; // an auxiliary vector, determine the boundary of the scanning area.
198- ibox[0 ] = int (this ->nx / 2 ) + 1 ; // scan x from -ibox[0] to ibox[0].
199- ibox[1 ] = int (this ->ny / 2 ) + 1 ; // scan y from -ibox[1] to ibox[1].
200- ibox[2 ] = int (this ->nz / 2 ) + 1 ; // scan z from -ibox[2] to ibox[2].
201-
202- ModuleBase::Vector3<double > f;
203-
204- int iy_start = -ibox[1 ]; // determine the scaning area along x-direct, if gamma-only, only positive axis is used.
205- int iy_end = ibox[1 ];
206- if (this ->gamma_only )
207- {
208- iy_start = 0 ;
209- iy_end = this ->ny ;
210- }
211- for (int ix = -ibox[0 ]; ix <= ibox[0 ]; ++ix)
212- {
213- for (int iy = iy_start; iy <= iy_end; ++iy)
214- {
215- // we shift all sticks to the first quadrant in x-y plane here.
216- // (ix, iy, iz) is the direct coordinates of planewaves.
217- // x and y is the coordinates of shifted sticks in x-y plane.
218- // for example, if ny = nx = 10, we will shift the stick on (-1, 2) to (9, 2),
219- // so that its index in st_length and st_bottom is 9 * 10 + 2 = 92.
220- int x = ix;
221- int y = iy;
222- if (x < 0 ) x += this ->nx ;
223- if (y < 0 ) y += this ->ny ;
224- int index = x * this ->ny + y;
225-
226- int length = 0 ; // number of planewave in stick (x, y).
227- for (int iz = -ibox[2 ]; iz <= ibox[2 ]; ++iz)
228- {
229- f.x = ix;
230- f.y = iy;
231- f.z = iz;
232- double modulus = f * (this ->GGT * f);
233- if (modulus <= this ->ggecut )
234- {
235- if (length == 0 ) st_bottom2D[index] = iz; // length == 0 means this point is the bottom of stick (x, y).
236- ++tot_npw;
237- ++length;
238- }
239- }
240- if (length > 0 )
241- {
242- st_length2D[index] = length;
243- ++this ->nstot ;
244- }
245- }
246- }
247- return ;
248- }
249-
250192//
251193// (2) Collect the x, y indexs, length of the sticks.
252194// Firstly, we scan the area and construct temp_st_*.
@@ -281,7 +223,7 @@ void PW_Basis::collect_st(
281223 if (this ->gamma_only )
282224 {
283225 iy_start = 0 ;
284- iy_end = this ->ny ;
226+ iy_end = this ->ny - 1 ;
285227 }
286228 for (int ix = -ibox[0 ]; ix <= ibox[0 ]; ++ix)
287229 {
@@ -359,7 +301,6 @@ void PW_Basis::collect_st(
359301// output: npw_per, nst_per, this->nstnz_per, this->ixy2ip, this->startnsz_per
360302//
361303void PW_Basis::divide_sticks (
362- const int tot_npw, // total number of planewaves.
363304 int * st_i, // x or x + nx (if x < 0) of stick.
364305 int * st_j, // y or y + ny (if y < 0) of stick.
365306 int * st_length, // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
@@ -458,58 +399,4 @@ void PW_Basis::get_istot2ixy(
458399 return ;
459400}
460401
461- //
462- // (5) Construct ig2isz, and is2ixy.
463- // is2ixy contains the x-coordinate and y-coordinate of sticks on current core.
464- // ig2isz contains the z-coordinate of planewaves on current core.
465- // We will scan all the sticks and find the planewaves on them, then store the information into ig2isz and is2ixy.
466- // known: this->nstot, st_bottom2D, st_length2D
467- // output: ig2isz, is2ixy
468- //
469- void PW_Basis::get_ig2isz_is2ixy (
470- int * st_bottom2D, // minimum z of stick, stored in 1d array with this->nstot elements.
471- int * st_length2D // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
472- )
473- {
474- if (this ->npw == 0 )
475- {
476- this ->ig2isz = new int [1 ]; // map ig to the z coordinate of this planewave.
477- this ->ig2isz [0 ] = 0 ;
478- this ->is2ixy = new int [1 ]; // map is (index of sticks) to ixy (ix + iy * nx).
479- this ->is2ixy [0 ] = -1 ;
480- return ;
481- }
482-
483- this ->ig2isz = new int [this ->npw ]; // map ig to the z coordinate of this planewave.
484- ModuleBase::GlobalFunc::ZEROS (this ->ig2isz , this ->npw );
485- this ->is2ixy = new int [this ->nst ]; // map is (index of sticks) to ixy (ix + iy * nx).
486- for (int is = 0 ; is < this ->nst ; ++is)
487- {
488- this ->is2ixy [is] = -1 ;
489- }
490-
491- int st_move = 0 ; // this is the st_move^th stick on current core.
492- int pw_filled = 0 ; // how many current core's planewaves have been found.
493- for (int ixy = 0 ; ixy < this ->nxy ; ++ixy)
494- {
495- if (this ->ixy2ip [ixy] == this ->poolrank )
496- {
497- int zstart = st_bottom2D[ixy];
498- for (int iz = zstart; iz < zstart + st_length2D[ixy]; ++iz)
499- {
500- int z = iz;
501- if (z < 0 ) z += this ->nz ;
502- this ->ig2isz [pw_filled] = st_move * this ->nz + z;
503- pw_filled++;
504- // assert(pw_filled <= this->npw);
505- }
506- this ->is2ixy [st_move] = ixy;
507- st_move++;
508- // assert(st_move <= this->nst);
509- }
510- if (st_move == this ->nst && pw_filled == this ->npw ) break ;
511- }
512- return ;
513- }
514-
515402}
0 commit comments