Skip to content

Commit e76341d

Browse files
authored
Merge pull request #510 from sunliang98/planewave
Add distributeg_method2 and corresponding unit tests.
2 parents 41d38a0 + dfdf315 commit e76341d

File tree

9 files changed

+1202
-560
lines changed

9 files changed

+1202
-560
lines changed

source/module_pw/pw_basis.h

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ class PW_Basis
8888

8989
//distribute plane waves to different processors
9090
void distribution_method1(); // x varies fast
91-
// void distribution_method2(); // y varies fast
91+
void distribution_method2(); // sticks sorted according to ixy
92+
// void distribution_method3(); // y varies fast
9293

9394
void collect_local_pw();
9495

@@ -118,16 +119,19 @@ class PW_Basis
118119
int* st_length2D, // the number of planewaves that belong to the stick located on (x, y).
119120
int* st_bottom2D // the z-coordinate of the bottom of stick on (x, y).
120121
);
122+
void get_ig2isz_is2ixy(
123+
int* st_bottom, // minimum z of stick, stored in 1d array with tot_nst elements.
124+
int* st_length // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
125+
);
126+
// for distributeg_method1
121127
void collect_st(
122128
int* st_length2D, // the number of planewaves that belong to the stick located on (x, y), stored in 2d x-y plane.
123129
int* st_bottom2D, // the z-coordinate of the bottom of stick on (x, y), stored in 2d x-y plane.
124130
int* st_i, // x or x + nx (if x < 0) of stick.
125131
int* st_j, // y or y + ny (if y < 0) of stick.
126132
int* st_length // number of planewaves in stick, stored in 1d array with tot_nst elements.
127133
);
128-
// for distributeg_method1
129134
void divide_sticks(
130-
const int tot_npw, // total number of planewaves.
131135
int* st_i, // x or x + nx (if x < 0) of stick.
132136
int* st_j, // y or y + ny (if y < 0) of stick.
133137
int* st_length, // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
@@ -138,11 +142,15 @@ class PW_Basis
138142
int* st_i, // x or x + nx (if x < 0) of stick.
139143
int* st_j // y or y + ny (if y < 0) of stick.
140144
);
141-
void get_ig2isz_is2ixy(
142-
int* st_bottom, // minimum z of stick, stored in 1d array with tot_nst elements.
143-
int* st_length // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
144-
);
145145
// for distributeg_method2
146+
void divide_sticks2(
147+
int* nst_per // number of sticks on each core.
148+
);
149+
void create_maps(
150+
int* st_length2D, // the number of planewaves that belong to the stick located on (x, y), stored in 2d x-y plane.
151+
int* npw_per // number of planewaves on each core.
152+
);
153+
// for distributeg_method3
146154
// void divide_sticks2(
147155
// const int tot_npw, // total number of planewaves.
148156
// int* st_i, // x or x + nx (if x < 0) of stick.

source/module_pw/pw_distributeg.cpp

Lines changed: 125 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "pw_basis.h"
22
#include "../module_base/tool_quit.h"
3+
#include "../module_base/global_function.h"
34
#include <iostream>
45

56
namespace ModulePW
@@ -15,14 +16,134 @@ void PW_Basis::distribute_g()
1516
{
1617
this->distribution_method1();
1718
}
18-
// else if(this->distribution_type == 2)
19-
// {
20-
// this->distribution_method2();
21-
// }
19+
else if(this->distribution_type == 2)
20+
{
21+
this->distribution_method2();
22+
}
2223
else
2324
{
2425
ModuleBase::WARNING_QUIT("divide", "No such division type.");
2526
}
2627
return;
2728
}
29+
30+
//
31+
// (1) We count the total number of planewaves (tot_npw) and sticks (this->nstot) here.
32+
// 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,
33+
// so that we can scan a much smaller area in step(2).
34+
// known: nx, ny, nz, ggecut, GGT
35+
// output: tot_npw, this->nstot, st_length2D, st_bottom2D
36+
//
37+
void PW_Basis::count_pw_st(
38+
int &tot_npw, // total number of planewaves.
39+
int* st_length2D, // the number of planewaves that belong to the stick located on (x, y).
40+
int* st_bottom2D // the z-coordinate of the bottom of stick on (x, y).
41+
)
42+
{
43+
int ibox[3] = {0, 0, 0}; // an auxiliary vector, determine the boundary of the scanning area.
44+
ibox[0] = int(this->nx / 2) + 1; // scan x from -ibox[0] to ibox[0].
45+
ibox[1] = int(this->ny / 2) + 1; // scan y from -ibox[1] to ibox[1].
46+
ibox[2] = int(this->nz / 2) + 1; // scan z from -ibox[2] to ibox[2].
47+
48+
ModuleBase::Vector3<double> f;
49+
50+
int iy_start = -ibox[1]; // determine the scaning area along x-direct, if gamma-only, only positive axis is used.
51+
int iy_end = ibox[1];
52+
if (this->gamma_only)
53+
{
54+
iy_start = 0;
55+
iy_end = this->ny - 1;
56+
}
57+
for (int ix = -ibox[0]; ix <= ibox[0]; ++ix)
58+
{
59+
for (int iy = iy_start; iy <= iy_end; ++iy)
60+
{
61+
// we shift all sticks to the first quadrant in x-y plane here.
62+
// (ix, iy, iz) is the direct coordinates of planewaves.
63+
// x and y is the coordinates of shifted sticks in x-y plane.
64+
// for example, if ny = nx = 10, we will shift the stick on (-1, 2) to (9, 2),
65+
// so that its index in st_length and st_bottom is 9 * 10 + 2 = 92.
66+
int x = ix;
67+
int y = iy;
68+
if (x < 0) x += this->nx;
69+
if (y < 0) y += this->ny;
70+
int index = x * this->ny + y;
71+
72+
int length = 0; // number of planewave in stick (x, y).
73+
for (int iz = -ibox[2]; iz <= ibox[2]; ++iz)
74+
{
75+
f.x = ix;
76+
f.y = iy;
77+
f.z = iz;
78+
double modulus = f * (this->GGT * f);
79+
if (modulus <= this->ggecut)
80+
{
81+
if (length == 0) st_bottom2D[index] = iz; // length == 0 means this point is the bottom of stick (x, y).
82+
++tot_npw;
83+
++length;
84+
}
85+
}
86+
if (length > 0)
87+
{
88+
st_length2D[index] = length;
89+
++this->nstot;
90+
}
91+
}
92+
}
93+
return;
94+
}
95+
96+
//
97+
// (5) Construct ig2isz, and is2ixy.
98+
// is2ixy contains the x-coordinate and y-coordinate of sticks on current core.
99+
// ig2isz contains the z-coordinate of planewaves on current core.
100+
// We will scan all the sticks and find the planewaves on them, then store the information into ig2isz and is2ixy.
101+
// known: this->nstot, st_bottom2D, st_length2D
102+
// output: ig2isz, is2ixy
103+
//
104+
void PW_Basis::get_ig2isz_is2ixy(
105+
int* st_bottom2D, // minimum z of stick, stored in 1d array with this->nstot elements.
106+
int* st_length2D // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
107+
)
108+
{
109+
if (this->npw == 0)
110+
{
111+
this->ig2isz = new int[1]; // map ig to the z coordinate of this planewave.
112+
this->ig2isz[0] = 0;
113+
this->is2ixy = new int[1]; // map is (index of sticks) to ixy (ix + iy * nx).
114+
this->is2ixy[0] = -1;
115+
return;
116+
}
117+
118+
this->ig2isz = new int[this->npw]; // map ig to the z coordinate of this planewave.
119+
ModuleBase::GlobalFunc::ZEROS(this->ig2isz, this->npw);
120+
this->is2ixy = new int[this->nst]; // map is (index of sticks) to ixy (ix + iy * nx).
121+
for (int is = 0; is < this->nst; ++is)
122+
{
123+
this->is2ixy[is] = -1;
124+
}
125+
126+
int st_move = 0; // this is the st_move^th stick on current core.
127+
int pw_filled = 0; // how many current core's planewaves have been found.
128+
for (int ixy = 0; ixy < this->nxy; ++ixy)
129+
{
130+
if (this->ixy2ip[ixy] == this->poolrank)
131+
{
132+
int zstart = st_bottom2D[ixy];
133+
for (int iz = zstart; iz < zstart + st_length2D[ixy]; ++iz)
134+
{
135+
int z = iz;
136+
if (z < 0) z += this->nz;
137+
this->ig2isz[pw_filled] = st_move * this->nz + z;
138+
pw_filled++;
139+
// assert(pw_filled <= this->npw);
140+
}
141+
this->is2ixy[st_move] = ixy;
142+
st_move++;
143+
// assert(st_move <= this->nst);
144+
}
145+
if (st_move == this->nst && pw_filled == this->npw) break;
146+
}
147+
return;
148+
}
28149
}

source/module_pw/pw_distributeg_method1.cpp

Lines changed: 12 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -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 << "\nThe 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
//
361303
void 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

Comments
 (0)