Skip to content

Commit 5f168ad

Browse files
maki49maki49
committed
remove CubeInfo and xyz-zxy permutation
1 parent c461e9f commit 5f168ad

File tree

7 files changed

+126
-124
lines changed

7 files changed

+126
-124
lines changed

‎source/module_elecstate/module_charge/symmetry_rho.cpp‎

Lines changed: 2 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -127,24 +127,8 @@ void Symmetry_rho::psymm(double* rho_part,
127127
}
128128

129129
// (3)
130-
const int ncxy = rho_basis->nx * rho_basis->ny;
131-
std::vector<double> zpiece(ncxy);
132-
for (int iz = 0; iz < rho_basis->nz; iz++)
133-
{
134-
ModuleBase::GlobalFunc::ZEROS(zpiece.data(), ncxy);
135-
if (GlobalV::MY_RANK == 0)
136-
{
137-
for (int ix = 0; ix < rho_basis->nx; ix++)
138-
{
139-
for (int iy = 0; iy < rho_basis->ny; iy++)
140-
{
141-
const int ir = ix * rho_basis->ny + iy;
142-
zpiece[ir] = rhotot[ix * rho_basis->ny * rho_basis->nz + iy * rho_basis->nz + iz];
143-
}
144-
}
145-
}
146-
Pgrid.zpiece_to_all(zpiece.data(), iz, rho_part);
147-
}
130+
Pgrid.bcast(rhotot.data(), rho_part);
131+
148132
#endif
149133
return;
150134
}

‎source/module_hamilt_pw/hamilt_pwdft/parallel_grid.cpp‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,28 @@ void Parallel_Grid::z_distribution()
185185

186186

187187
#ifdef __MPI
188-
void Parallel_Grid::zpiece_to_all(double *zpiece, const int &iz, double *rho) const
188+
void Parallel_Grid::bcast(const double* const data_global, double* data_local)const
189+
{
190+
std::vector<double> zpiece(ncxy);
191+
for (int iz = 0; iz < this->ncz; ++iz)
192+
{
193+
ModuleBase::GlobalFunc::ZEROS(zpiece.data(), ncxy);
194+
if (GlobalV::MY_RANK == 0)
195+
{
196+
for (int ix = 0; ix < ncx; ix++)
197+
{
198+
for (int iy = 0; iy < ncy; iy++)
199+
{
200+
const int ixy = ix * ncy + iy;
201+
zpiece[ixy] = data_global[ixy * ncz + iz];
202+
}
203+
}
204+
}
205+
this->zpiece_to_all(zpiece.data(), iz, data_local);
206+
}
207+
}
208+
209+
void Parallel_Grid::zpiece_to_all(double* zpiece, const int& iz, double* rho) const
189210
{
190211
if(PARAM.inp.esolver_type == "sdft")
191212
{

‎source/module_hamilt_pw/hamilt_pwdft/parallel_grid.h‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,12 @@ class Parallel_Grid
2727
const int &nczp, const int &nrxx, const int &nbz, const int &bz); //LiuXh add 20180606
2828

2929
#ifdef __MPI
30-
void zpiece_to_all(double *zpiece, const int &iz, double *rho) const;
31-
void zpiece_to_stogroup(double *zpiece, const int &iz, double *rho) const; //qainrui add for sto-dft 2021-7-21
32-
30+
void zpiece_to_stogroup(double* zpiece, const int& iz, double* rho) const; //qainrui add for sto-dft 2021-7-21
31+
void zpiece_to_all(double* zpiece, const int& iz, double* rho) const;
32+
33+
/// @brief Broadcast data from root to all processors. The index order is [x][y][z].
34+
void bcast(const double* const data_global, double* data_local)const;
35+
/// @brief Reduce data from all processors to root. The index order is [x][y][z].
3336
void reduce(double* rhotot, const double* constrhoin)const;
3437
#endif
3538

‎source/module_io/cube_io.h‎

Lines changed: 16 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,20 @@ namespace ModuleIO
2626
const int precision = 11,
2727
const int out_fermi = 1); // mohan add 2007-10-17
2828

29-
struct CubeInfo
30-
{
31-
CubeInfo(
29+
/// read the full data from a cube file
30+
bool read_cube(const std::string& file,
31+
std::vector<std::string>& comment,
32+
int& natom,
33+
std::vector<double>& cel_pos,
34+
std::vector<int>& nvoxel,
35+
std::vector<std::vector<double>>& axis_vecs,
36+
std::vector<int>& atom_type,
37+
std::vector<double>& atom_charge,
38+
std::vector<std::vector<double>>& atom_pos,
39+
std::vector<double>& data);
40+
41+
/// write a cube file
42+
void write_cube(const std::string& file,
3243
const std::vector<std::string>& comment,
3344
const int natom,
3445
const std::vector<double>& cel_pos,
@@ -38,34 +49,8 @@ struct CubeInfo
3849
const std::vector<double>& atom_charge,
3950
const std::vector<std::vector<double>>& atom_pos,
4051
const std::vector<double>& data,
41-
const bool valid)
42-
: comment(comment), natom(natom), cel_pos(cel_pos),
43-
nvoxel(nvoxel), axis_vecs(axis_vecs),
44-
atom_type(atom_type), atom_charge(atom_charge), atom_pos(atom_pos),
45-
data(data), valid(valid)
46-
{
47-
};
48-
49-
const std::vector<std::string> comment = {};
50-
const int natom = 0;
51-
const std::vector<double> cel_pos = {};
52-
const std::vector<int> nvoxel = {};
53-
const std::vector<std::vector<double>> axis_vecs = {};
54-
const std::vector<int> atom_type = {};
55-
const std::vector<double> atom_charge = {};
56-
const std::vector<std::vector<double>> atom_pos = {};
57-
const std::vector<double> data = {};
58-
const bool valid = false;
59-
};
60-
61-
/// read the full data from a cube file
62-
CubeInfo read_cube(const std::string& file);
63-
/// write a cube file
64-
void write_cube(const std::string& file, const CubeInfo& info, const int precision, const int ndata_line = 6);
65-
/// change the index order: [x][y][z] (.cube file) -> [z][x][y] (ABACUS)
66-
void xyz2zxy(const double* const xyz, const int nx, const int ny, const int nz, double* const zxy);
67-
/// change the index order: [z][x][y] (ABACUS) -> [x][y][z] (.cube file)
68-
void zxy2xyz(const double* const zxy, const int nx, const int ny, const int nz, double* const xyz);
52+
const int precision,
53+
const int ndata_line = 6);
6954

7055
/**
7156
* @brief The trilinear interpolation method

‎source/module_io/read_cube.cpp‎

Lines changed: 40 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,6 @@
22
#include <limits>
33
#include "module_hamilt_pw/hamilt_pwdft/parallel_grid.h"
44
// #include "module_base/global_variable.h" // GlobalV reference removed
5-
void ModuleIO::xyz2zxy(const double* const xyz, const int nx, const int ny, const int nz, double* const zxy)
6-
{
7-
for (int ix = 0; ix < nx; ix++)
8-
{
9-
for (int iy = 0; iy < ny; iy++)
10-
{
11-
for (int iz = 0; iz < nz; iz++)
12-
{
13-
zxy[(iz * nx + ix) * ny + iy] = xyz[(ix * ny + iy) * nz + iz];
14-
}
15-
}
16-
}
17-
}
185

196
bool ModuleIO::read_grid(
207
const Parallel_Grid& pgrid,
@@ -44,34 +31,40 @@ bool ModuleIO::read_grid(
4431
const int& ny = pgrid.ny;
4532
const int& nz = pgrid.nz;
4633
const int& nxyz = nx * ny * nz;
47-
std::vector<double> data_zxy_full(nxyz, 0.0); // [iz][ix][iy]
34+
std::vector<double> data_xyz_full(nxyz, 0.0);
4835
if (my_rank == 0)
4936
{
50-
const CubeInfo& cube_info = ModuleIO::read_cube(fn);
51-
52-
const int& nx_read = cube_info.nvoxel[0];
53-
const int& ny_read = cube_info.nvoxel[1];
54-
const int& nz_read = cube_info.nvoxel[2];
37+
std::vector<std::string> comment;
38+
int natom = 0;
39+
std::vector<double> cel_pos;
40+
std::vector<int> nvoxel;
41+
std::vector<std::vector<double>> axis_vecs;
42+
std::vector<int> atom_type;
43+
std::vector<double> atom_charge;
44+
std::vector<std::vector<double>> atom_pos;
45+
std::vector<double> data_read;
46+
bool valid = ModuleIO::read_cube(fn, comment, natom, cel_pos, nvoxel, axis_vecs, atom_type, atom_charge, atom_pos, data_read);
47+
48+
const int& nx_read = nvoxel[0];
49+
const int& ny_read = nvoxel[1];
50+
const int& nz_read = nvoxel[2];
5551

5652
// if mismatch, trilinear interpolate
5753
if (nx == nx_read && ny == ny_read && nz == nz_read)
5854
{
59-
ModuleIO::xyz2zxy(cube_info.data.data(), nx, ny, nz, data_zxy_full.data());
55+
std::memcpy(data_xyz_full.data(), data_read.data(), nxyz * sizeof(double));
6056
}
6157
else
6258
{
63-
std::vector<double> data_xyz_full(nxyz, 0.0);
64-
trilinear_interpolate(cube_info.data.data(), nx_read, ny_read, nz_read, nx, ny, nz, data_xyz_full.data());
65-
ModuleIO::xyz2zxy(data_xyz_full.data(), nx, ny, nz, data_zxy_full.data());
59+
trilinear_interpolate(data_read.data(), nx_read, ny_read, nz_read, nx, ny, nz, data_xyz_full.data());
6660
}
6761
}
6862

6963
// distribute
7064
#ifdef __MPI
71-
const int nxy = nx * ny;
72-
for (int iz = 0;iz < nz;++iz) { pgrid.zpiece_to_all(data_zxy_full.data() + iz * nxy, iz, data); }
65+
pgrid.bcast(data_xyz_full.data(), data);
7366
#else
74-
std::memcpy(data, data_zxy_full.data(), nxyz * sizeof(double));
67+
std::memcpy(data, data_xyz_full.data(), nxyz * sizeof(double));
7568
#endif
7669
return true;
7770
}
@@ -147,43 +140,51 @@ void ModuleIO::trilinear_interpolate(
147140
delete[] read_rho;
148141
}
149142

150-
ModuleIO::CubeInfo ModuleIO::read_cube(const std::string& file)
143+
bool ModuleIO::read_cube(const std::string& file,
144+
std::vector<std::string>& comment,
145+
int& natom,
146+
std::vector<double>& cell_pos,
147+
std::vector<int>& nvoxel,
148+
std::vector<std::vector<double>>& axis_vecs,
149+
std::vector<int>& atom_type,
150+
std::vector<double>& atom_charge,
151+
std::vector<std::vector<double>>& atom_pos,
152+
std::vector<double>& data)
151153
{
152154
std::ifstream ifs(file);
153155

154-
if (!ifs) { return ModuleIO::CubeInfo({}, 0, {}, {}, {}, {}, {}, {}, {}, false); }
156+
if (!ifs) { return false; }
155157

156-
std::vector<std::string> comment(2);
158+
comment.resize(2);
157159
for (auto& c : comment) { std::getline(ifs, c); }
158160

159-
int natom;
160161
ifs >> natom;
161-
std::vector<double> cell_pos(3);
162+
cell_pos.resize(3);
162163
for (auto& cp : cell_pos) { ifs >> cp; }
163164

164-
std::vector<int> nvoxel(3);
165-
std::vector<std::vector<double>> axis_vecs(3);
165+
nvoxel.resize(3);
166+
axis_vecs.resize(0);
166167
for (int i = 0;i < 3;++i)
167168
{
168169
std::vector<double> vec(3);
169170
ifs >> nvoxel[i] >> vec[0] >> vec[1] >> vec[2];
170171
axis_vecs.push_back(vec);
171172
}
172173

173-
std::vector<int> itype(natom);
174-
std::vector<double> charge(natom);
175-
std::vector<std::vector<double>> atom_pos(natom);
174+
atom_type.resize(natom);
175+
atom_charge.resize(natom);
176+
atom_pos.resize(0);
176177
for (int i = 0;i < natom;++i)
177178
{
178179
std::vector<double> apos(3);
179-
ifs >> itype[i] >> charge[i] >> apos[0] >> apos[1] >> apos[2];
180+
ifs >> atom_type[i] >> atom_charge[i] >> apos[0] >> apos[1] >> apos[2];
180181
atom_pos.push_back(apos);
181182
}
182183

183184
const int nxyz = nvoxel[0] * nvoxel[1] * nvoxel[2];
184-
std::vector<double> data(nxyz);
185+
data.resize(nxyz);
185186
for (int i = 0;i < nxyz;++i) { ifs >> data[i]; }
186187

187188
ifs.close();
188-
return ModuleIO::CubeInfo(comment, natom, cell_pos, nvoxel, axis_vecs, itype, charge, atom_pos, std::move(data), true);
189+
return true;
189190
}

‎source/module_io/test_serial/rho_io_test.cpp‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,11 +123,27 @@ TEST_F(RhoIOTest, TrilinearInterpolate)
123123
}
124124
}
125125
}
126+
127+
// The old implementation is inconsistent: ifdef MPI, [x][y][z]; else, [z][x][y].
128+
// Now we use [x][y][z] for both MPI and non-MPI, so here we need to chage the index order.
129+
auto permute_xyz2zxy = [&](const double* const xyz, double* const zxy) -> void
130+
{
131+
for (int ix = 0; ix < nx; ix++)
132+
{
133+
for (int iy = 0; iy < ny; iy++)
134+
{
135+
for (int iz = 0; iz < nz; iz++)
136+
{
137+
zxy[(iz * nx + ix) * ny + iy] = xyz[(ix * ny + iy) * nz + iz];
138+
}
139+
}
140+
}
141+
};
126142
const int nxyz = nx * ny * nz;
127143
std::vector<double> data_xyz(nxyz);
128144
std::vector<double> data(nxyz); // z > x > y
129145
ModuleIO::trilinear_interpolate(data_read.data(), nx_read, ny_read, nz_read, nx, ny, nz, data_xyz.data());
130-
ModuleIO::xyz2zxy(data_xyz.data(), nx, ny, nz, data.data());
146+
permute_xyz2zxy(data_xyz.data(), data.data());
131147
EXPECT_DOUBLE_EQ(data[0], 0.0010824725010374092);
132148
EXPECT_DOUBLE_EQ(data[10], 0.058649850374240906);
133149
EXPECT_DOUBLE_EQ(data[100], 0.018931708073604996);

0 commit comments

Comments
 (0)