Skip to content

Commit 17f8b7e

Browse files
maki49maki49
committed
rewrite io_cube
1 parent be48991 commit 17f8b7e

File tree

4 files changed

+245
-256
lines changed

4 files changed

+245
-256
lines changed

‎source/module_io/cube_io.h‎

Lines changed: 48 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -26,42 +26,46 @@ void write_cube(
2626
const int precision = 11,
2727
const int out_fermi = 1); // mohan add 2007-10-17
2828

29+
struct CubeInfo
30+
{
31+
CubeInfo(
32+
const std::vector<std::string>& comment,
33+
const int natom,
34+
const std::vector<double>& cel_pos,
35+
const std::vector<int>& nvoxel,
36+
const std::vector<std::vector<double>>& axis_vecs,
37+
const std::vector<int>& atom_type,
38+
const std::vector<double>& atom_charge,
39+
const std::vector<std::vector<double>>& atom_pos,
40+
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+
};
2948

30-
// when MPI:
31-
// read file as order (ixy,iz) to data[ixy*nz+iz]
32-
// when serial:
33-
// read file as order (ixy,iz) to data[iz*nxy+ixy]
34-
void read_cube_core_match(
35-
std::ifstream &ifs,
36-
const Parallel_Grid& pgrid,
37-
const bool flag_read_rank,
38-
double*const data,
39-
const int nxy,
40-
const int nz);
41-
42-
void read_cube_core_mismatch(
43-
std::ifstream &ifs,
44-
const Parallel_Grid& pgrid,
45-
const bool flag_read_rank,
46-
double*const data,
47-
const int nx,
48-
const int ny,
49-
const int nz,
50-
const int nx_read,
51-
const int ny_read,
52-
const int nz_read);
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+
};
5360

54-
// when MPI:
55-
// write data[ixy*nplane+iz] to file as order (ixy,iz)
56-
// when serial:
57-
// write data[iz*nxy+ixy] to file as order (ixy,iz)
58-
void write_cube_core(
59-
const Parallel_Grid& pgrid,
60-
std::ofstream& ofs_cube,
61-
const double*const data,
62-
const int nxy,
63-
const int nz,
64-
const int n_data_newline);
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);
6569

6670
/**
6771
* @brief The trilinear interpolation method
@@ -84,28 +88,23 @@ void write_cube_core(
8488
* directions, divided by the grid spacing. Here, it is assumed that the grid spacing is equal and can be
8589
* omitted during computation.
8690
*
87-
* @param ifs the ifstream used to read charge density
91+
* @param data_in the input data of size nxyz_read
8892
* @param nx_read nx read from file
8993
* @param ny_read ny read from file
9094
* @param nz_read nz read from file
9195
* @param nx the dimension of grids along x
9296
* @param ny the dimension of grids along y
9397
* @param nz the dimension of grids along z
94-
* @param data the interpolated results
98+
* @param data_out the interpolated results of size nxyz
9599
*/
96-
void trilinear_interpolate(std::ifstream& ifs,
97-
const int& nx_read,
98-
const int& ny_read,
99-
const int& nz_read,
100-
const int& nx,
101-
const int& ny,
102-
const int& nz,
103-
#ifdef __MPI
104-
std::vector<std::vector<double>> &data
105-
#else
106-
double* data
107-
#endif
108-
);
100+
void trilinear_interpolate(const double* const data_in,
101+
const int& nx_read,
102+
const int& ny_read,
103+
const int& nz_read,
104+
const int& nx,
105+
const int& ny,
106+
const int& nz,
107+
double* data_out);
109108
}
110109

111110
#endif

‎source/module_io/read_cube.cpp‎

Lines changed: 94 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,19 @@
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+
}
518

619
bool ModuleIO::read_cube(
720
const Parallel_Grid& pgrid,
@@ -12,133 +25,66 @@ bool ModuleIO::read_cube(
1225
const int natom)
1326
{
1427
ModuleBase::TITLE("ModuleIO", "read_cube");
28+
29+
// check if the file exists
1530
std::ifstream ifs(fn.c_str());
1631
if (!ifs)
1732
{
18-
std::string tmp_warning_info = "!!! Couldn't find the charge file of ";
19-
tmp_warning_info += fn;
33+
std::string tmp_warning_info = "!!! Couldn't find the file: " + fn;
2034
ofs_running << tmp_warning_info << std::endl;
2135
return false;
2236
}
2337
else
2438
{
25-
ofs_running << " Find the file, try to read charge from file." << std::endl;
39+
ofs_running << " Find the file " << fn << " , try to read it." << std::endl;
2640
}
2741

28-
// skip the first 3 lines
29-
for (int i = 0;i < 3;++i) { ifs.ignore(std::numeric_limits<std::streamsize>::max(), '\n'); }
30-
31-
int nx_read = 0;
32-
int ny_read = 0;
33-
int nz_read = 0;
34-
std::string temp;
35-
36-
ifs >> nx_read;
37-
ifs >> temp >> temp >> temp;
38-
ifs >> ny_read;
39-
ifs >> temp >> temp >> temp;
40-
ifs >> nz_read;
41-
ifs >> temp >> temp >> temp;
42-
42+
// read the full grid data
4343
const int& nx = pgrid.nx;
4444
const int& ny = pgrid.ny;
4545
const int& nz = pgrid.nz;
46+
const int& nxyz = nx * ny * nz;
47+
std::vector<double> data_zxy_full(nxyz, 0.0); // [iz][ix][iy]
48+
if (my_rank == 0)
49+
{
50+
const CubeInfo& cube_info = ModuleIO::read_cube(fn);
4651

47-
// skip this line and the next natom lines
48-
for (int i = 0;i < natom + 1;++i) { ifs.ignore(std::numeric_limits<std::streamsize>::max(), '\n'); }
49-
50-
#ifdef __MPI
51-
if(nx == nx_read && ny == ny_read && nz == nz_read)
52-
ModuleIO::read_cube_core_match(ifs, pgrid, (my_rank == 0), data, nx * ny, nz);
53-
else
54-
ModuleIO::read_cube_core_mismatch(ifs, pgrid, (my_rank == 0), data, nx, ny, nz, nx_read, ny_read, nz_read);
55-
#else
56-
57-
if(nx == nx_read && ny == ny_read && nz == nz_read)
58-
ModuleIO::read_cube_core_match(ifs, pgrid, (my_rank == 0), data, nx * ny, nz);
59-
else
60-
ModuleIO::read_cube_core_mismatch(ifs, pgrid, (my_rank == 0), data, nx, ny, nz, nx_read, ny_read, nz_read);
61-
#endif
62-
63-
return true;
64-
}
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];
6555

66-
void ModuleIO::read_cube_core_match(
67-
std::ifstream& ifs,
68-
const Parallel_Grid& pgrid,
69-
const bool flag_read_rank,
70-
double*const data,
71-
const int nxy,
72-
const int nz)
73-
{
74-
#ifdef __MPI
75-
if (flag_read_rank)
76-
{
77-
std::vector<std::vector<double>> read_rho(nz, std::vector<double>(nxy));
78-
for (int ixy = 0; ixy < nxy; ixy++)
79-
for (int iz = 0; iz < nz; iz++)
80-
ifs >> read_rho[iz][ixy];
81-
for (int iz = 0; iz < nz; iz++)
82-
pgrid.zpiece_to_all(read_rho[iz].data(), iz, data);
83-
}
84-
else
85-
{
86-
std::vector<double> zpiece(nxy);
87-
for (int iz = 0; iz < nz; iz++)
88-
pgrid.zpiece_to_all(zpiece.data(), iz, data);
56+
// if mismatch, trilinear interpolate
57+
if (nx == nx_read && ny == ny_read && nz == nz_read)
58+
{
59+
ModuleIO::xyz2zxy(cube_info.data.data(), nx, ny, nz, data_zxy_full.data());
60+
}
61+
else
62+
{
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());
66+
}
8967
}
90-
#else
91-
for (int ixy = 0; ixy < nxy; ixy++)
92-
for (int iz = 0; iz < nz; iz++)
93-
ifs >> data[iz * nxy + ixy];
94-
#endif
95-
}
9668

97-
void ModuleIO::read_cube_core_mismatch(
98-
std::ifstream &ifs,
99-
const Parallel_Grid& pgrid,
100-
const bool flag_read_rank,
101-
double*const data,
102-
const int nx,
103-
const int ny,
104-
const int nz,
105-
const int nx_read,
106-
const int ny_read,
107-
const int nz_read)
108-
{
109-
#ifdef __MPI
69+
// distribute
70+
#ifdef __MPI
11071
const int nxy = nx * ny;
111-
if (flag_read_rank)
112-
{
113-
std::vector<std::vector<double>> read_rho(nz, std::vector<double>(nxy));
114-
ModuleIO::trilinear_interpolate(ifs, nx_read, ny_read, nz_read, nx, ny, nz, read_rho);
115-
for (int iz = 0; iz < nz; iz++)
116-
pgrid.zpiece_to_all(read_rho[iz].data(), iz, data);
117-
}
118-
else
119-
{
120-
std::vector<double> zpiece(nxy);
121-
for (int iz = 0; iz < nz; iz++)
122-
pgrid.zpiece_to_all(zpiece.data(), iz, data);
123-
}
72+
for (int iz = 0;iz < nz;++iz) { pgrid.zpiece_to_all(data_zxy_full.data() + iz * nxy, iz, data); }
12473
#else
125-
ModuleIO::trilinear_interpolate(ifs, nx_read, ny_read, nz_read, nx, ny, nz, data);
74+
std::memcpy(data, data_zxy_full.data(), nxyz * sizeof(double));
12675
#endif
76+
return true;
12777
}
12878

129-
void ModuleIO::trilinear_interpolate(std::ifstream& ifs,
130-
const int& nx_read,
131-
const int& ny_read,
132-
const int& nz_read,
133-
const int& nx,
134-
const int& ny,
135-
const int& nz,
136-
#ifdef __MPI
137-
std::vector<std::vector<double>> &data
138-
#else
139-
double* data
140-
#endif
141-
)
79+
void ModuleIO::trilinear_interpolate(
80+
const double* const data_in,
81+
const int& nx_read,
82+
const int& ny_read,
83+
const int& nz_read,
84+
const int& nx,
85+
const int& ny,
86+
const int& nz,
87+
double* data_out)
14288
{
14389
ModuleBase::TITLE("ModuleIO", "trilinear_interpolate");
14490

@@ -153,7 +99,7 @@ void ModuleIO::trilinear_interpolate(std::ifstream& ifs,
15399
{
154100
for (int iz = 0; iz < nz_read; iz++)
155101
{
156-
ifs >> read_rho[iz][ix * ny_read + iy];
102+
read_rho[iz][ix * ny_read + iy] = data_in[(ix * ny_read + iy) * nz_read + iz];
157103
}
158104
}
159105
}
@@ -189,11 +135,7 @@ void ModuleIO::trilinear_interpolate(std::ifstream& ifs,
189135
+ read_rho[highz][lowx * ny_read + highy] * (1 - dx) * dy * dz
190136
+ read_rho[highz][highx * ny_read + highy] * dx * dy * dz;
191137

192-
#ifdef __MPI
193-
data[iz][ix * ny + iy] = result;
194-
#else
195-
data[iz * nx * ny + ix * ny + iy] = result;
196-
#endif
138+
data_out[(ix * ny + iy) * nz + iz] = result; // x > y > z order, consistent with the cube file
197139
}
198140
}
199141
}
@@ -203,4 +145,45 @@ void ModuleIO::trilinear_interpolate(std::ifstream& ifs,
203145
delete[] read_rho[iz];
204146
}
205147
delete[] read_rho;
148+
}
149+
150+
ModuleIO::CubeInfo ModuleIO::read_cube(const std::string& file)
151+
{
152+
std::ifstream ifs(file);
153+
154+
if (!ifs) { return ModuleIO::CubeInfo({}, 0, {}, {}, {}, {}, {}, {}, {}, false); }
155+
156+
std::vector<std::string> comment(2);
157+
for (auto& c : comment) { std::getline(ifs, c); }
158+
159+
int natom;
160+
ifs >> natom;
161+
std::vector<double> cell_pos(3);
162+
for (auto& cp : cell_pos) { ifs >> cp; }
163+
164+
std::vector<int> nvoxel(3);
165+
std::vector<std::vector<double>> axis_vecs(3);
166+
for (int i = 0;i < 3;++i)
167+
{
168+
std::vector<double> vec(3);
169+
ifs >> nvoxel[i] >> vec[0] >> vec[1] >> vec[2];
170+
axis_vecs.push_back(vec);
171+
}
172+
173+
std::vector<int> itype(natom);
174+
std::vector<double> charge(natom);
175+
std::vector<std::vector<double>> atom_pos(natom);
176+
for (int i = 0;i < natom;++i)
177+
{
178+
std::vector<double> apos(3);
179+
ifs >> itype[i] >> charge[i] >> apos[0] >> apos[1] >> apos[2];
180+
atom_pos.push_back(apos);
181+
}
182+
183+
const int nxyz = nvoxel[0] * nvoxel[1] * nvoxel[2];
184+
std::vector<double> data(nxyz);
185+
for (int i = 0;i < nxyz;++i) { ifs >> data[i]; }
186+
187+
ifs.close();
188+
return ModuleIO::CubeInfo(comment, natom, cell_pos, nvoxel, axis_vecs, itype, charge, atom_pos, std::move(data), true);
206189
}

0 commit comments

Comments
 (0)