From cc2032bbcab8eb390fa486320f439e9b1ecca3a1 Mon Sep 17 00:00:00 2001 From: PeizeLin <78645006+PeizeLin@users.noreply.github.com> Date: Sun, 29 Sep 2024 11:57:16 +0800 Subject: [PATCH] Refactor `ModuleIO::read/write_cube()` (#5150) * Refactor ModuleIO::read/write_cube() * Refactor ModuleIO::write_cube_core() * Refactor ModuleIO::read_cube_core() * add annotation in ModuleIO::read/write_cube_core() --- .../hamilt_pwdft/parallel_grid.cpp | 4 +- .../hamilt_pwdft/parallel_grid.h | 4 +- source/module_io/cube_io.h | 102 ++++++--- source/module_io/read_cube.cpp | 202 +++++++++--------- source/module_io/write_cube.cpp | 123 ++++++----- 5 files changed, 249 insertions(+), 186 deletions(-) diff --git a/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.cpp b/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.cpp index e4c7fca0ea..d215f2f391 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.cpp @@ -182,7 +182,7 @@ void Parallel_Grid::z_distribution(void) #ifdef __MPI -void Parallel_Grid::zpiece_to_all(double *zpiece, const int &iz, double *rho) +void Parallel_Grid::zpiece_to_all(double *zpiece, const int &iz, double *rho) const { if(PARAM.inp.esolver_type == "sdft") { @@ -256,7 +256,7 @@ void Parallel_Grid::zpiece_to_all(double *zpiece, const int &iz, double *rho) #endif #ifdef __MPI -void Parallel_Grid::zpiece_to_stogroup(double *zpiece, const int &iz, double *rho) +void Parallel_Grid::zpiece_to_stogroup(double *zpiece, const int &iz, double *rho) const { assert(allocate); //TITLE("Parallel_Grid","zpiece_to_all"); diff --git a/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.h b/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.h index 222c304baa..d38997c264 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.h +++ b/source/module_hamilt_pw/hamilt_pwdft/parallel_grid.h @@ -21,8 +21,8 @@ class Parallel_Grid const int &nczp, const int &nrxx, const int &nbz, const int &bz); //LiuXh add 20180606 #ifdef __MPI - void zpiece_to_all(double *zpiece, const int &iz, double *rho); - void zpiece_to_stogroup(double *zpiece, const int &iz, double *rho); //qainrui add for sto-dft 2021-7-21 + void zpiece_to_all(double *zpiece, const int &iz, double *rho) const; + void zpiece_to_stogroup(double *zpiece, const int &iz, double *rho) const; //qainrui add for sto-dft 2021-7-21 void reduce_to_fullrho(double *rhotot, double *rhoin); #endif diff --git a/source/module_io/cube_io.h b/source/module_io/cube_io.h index c5606329c2..a6a296b5dd 100644 --- a/source/module_io/cube_io.h +++ b/source/module_io/cube_io.h @@ -10,43 +10,89 @@ namespace ModuleIO { bool read_cube( #ifdef __MPI - Parallel_Grid* Pgrid, + const Parallel_Grid*const Pgrid, #endif - int my_rank, - std::string esolver_type, - int rank_in_stogroup, - const int& is, + const int my_rank, + const std::string esolver_type, + const int rank_in_stogroup, + const int is, std::ofstream& ofs_running, - const int& nspin, + const int nspin, const std::string& fn, - double* data, - const int& nx, - const int& ny, - const int& nz, + double*const data, + const int nx, + const int ny, + const int nz, double& ef, - const UnitCell* ucell, + const UnitCell*const ucell, int& prenspin, - const bool& warning_flag = true); + const bool warning_flag = true); void write_cube( #ifdef __MPI - const int& bz, - const int& nbz, - const int& nplane, - const int& startz_current, + const int bz, + const int nbz, + const int nplane, + const int startz_current, #endif - const double* data, - const int& is, - const int& nspin, - const int& iter, + const double*const data, + const int is, + const int nspin, + const int iter, const std::string& fn, - const int& nx, - const int& ny, - const int& nz, - const double& ef, - const UnitCell* ucell, - const int& precision = 11, - const int& out_fermi = 1); // mohan add 2007-10-17 + const int nx, + const int ny, + const int nz, + const double ef, + const UnitCell*const ucell, + const int precision = 11, + const int out_fermi = 1); // mohan add 2007-10-17 + + +// when MPI: +// read file as order (ixy,iz) to data[ixy*nz+iz] +// when serial: +// read file as order (ixy,iz) to data[iz*nxy+ixy] +void read_cube_core_match( + std::ifstream &ifs, +#ifdef __MPI + const Parallel_Grid*const Pgrid, + const bool flag_read_rank, +#endif + double*const data, + const int nxy, + const int nz); + +void read_cube_core_mismatch( + std::ifstream &ifs, +#ifdef __MPI + const Parallel_Grid*const Pgrid, + const bool flag_read_rank, +#endif + double*const data, + const int nx, + const int ny, + const int nz, + const int nx_read, + const int ny_read, + const int nz_read); + +// when MPI: +// write data[ixy*nplane+iz] to file as order (ixy,iz) +// when serial: +// write data[iz*nxy+ixy] to file as order (ixy,iz) +void write_cube_core( + std::ofstream &ofs_cube, +#ifdef __MPI + const int bz, + const int nbz, + const int nplane, + const int startz_current, +#endif + const double*const data, + const int nxy, + const int nz, + const int n_data_newline); /** * @brief The trilinear interpolation method @@ -86,7 +132,7 @@ void write_cube( const int& ny, const int& nz, #ifdef __MPI - double** data + std::vector> &data #else double* data #endif diff --git a/source/module_io/read_cube.cpp b/source/module_io/read_cube.cpp index 286a5e7f2d..94f2bc1ed5 100644 --- a/source/module_io/read_cube.cpp +++ b/source/module_io/read_cube.cpp @@ -3,44 +3,44 @@ bool ModuleIO::read_cube( #ifdef __MPI - Parallel_Grid* Pgrid, + const Parallel_Grid*const Pgrid, #endif - int my_rank, - std::string esolver_type, - int rank_in_stogroup, - const int& is, + const int my_rank, + const std::string esolver_type, + const int rank_in_stogroup, + const int is, std::ofstream& ofs_running, - const int& nspin, + const int nspin, const std::string& fn, - double* data, - const int& nx, - const int& ny, - const int& nz, + double*const data, + const int nx, + const int ny, + const int nz, double& ef, - const UnitCell* ucell, + const UnitCell*const ucell, int& prenspin, - const bool& warning_flag) + const bool warning_flag) { ModuleBase::TITLE("ModuleIO","read_cube"); std::ifstream ifs(fn.c_str()); - if (!ifs) - { - std::string tmp_warning_info = "!!! Couldn't find the charge file of "; - tmp_warning_info += fn; - ofs_running << tmp_warning_info << std::endl; - return false; - } - else - { - ofs_running << " Find the file, try to read charge from file." << std::endl; - } + if (!ifs) + { + std::string tmp_warning_info = "!!! Couldn't find the charge file of "; + tmp_warning_info += fn; + ofs_running << tmp_warning_info << std::endl; + return false; + } + else + { + ofs_running << " Find the file, try to read charge from file." << std::endl; + } - bool quit=false; + bool quit=false; - ifs.ignore(300, '\n'); // skip the header + ifs.ignore(300, '\n'); // skip the header - if(nspin != 4) - { + if(nspin != 4) + { int v_in; ifs >> v_in; if (v_in != nspin) @@ -49,16 +49,16 @@ bool ModuleIO::read_cube( return false; } } - else - { - ifs >> prenspin; - } - ifs.ignore(150, ')'); + else + { + ifs >> prenspin; + } + ifs.ignore(150, ')'); - ifs >> ef; - ofs_running << " read in fermi energy = " << ef << std::endl; + ifs >> ef; + ofs_running << " read in fermi energy = " << ef << std::endl; - ifs.ignore(150, '\n'); + ifs.ignore(150, '\n'); ModuleBase::CHECK_INT(ifs, ucell->nat); ifs.ignore(150, '\n'); @@ -93,8 +93,6 @@ bool ModuleIO::read_cube( ifs >> temp >> temp >> temp; } - const bool same = (nx == nx_read && ny == ny_read && nz == nz_read) ? true : false; - for (int it = 0; it < ucell->ntype; it++) { for (int ia = 0; ia < ucell->atoms[it].na; ia++) @@ -114,86 +112,88 @@ bool ModuleIO::read_cube( } } + const bool flag_read_rank = (my_rank == 0 || (esolver_type == "sdft" && rank_in_stogroup == 0)); #ifdef __MPI - const int nxy = nx * ny; - double* zpiece = nullptr; - double** read_rho = nullptr; - if (my_rank == 0 || (esolver_type == "sdft" && rank_in_stogroup == 0)) - { - read_rho = new double*[nz]; - for (int iz = 0; iz < nz; iz++) - { - read_rho[iz] = new double[nxy]; - } - if (same) - { - for (int ix = 0; ix < nx; ix++) - { - for (int iy = 0; iy < ny; iy++) - { - for (int iz = 0; iz < nz; iz++) - { - ifs >> read_rho[iz][ix * ny + iy]; - } - } - } - } - else - { - ModuleIO::trilinear_interpolate(ifs, nx_read, ny_read, nz_read, nx, ny, nz, read_rho); - } - } + if(nx == nx_read && ny == ny_read && nz == nz_read) + ModuleIO::read_cube_core_match(ifs, Pgrid, flag_read_rank, data, nx*ny, nz); else - { - zpiece = new double[nxy]; - ModuleBase::GlobalFunc::ZEROS(zpiece, nxy); - } + ModuleIO::read_cube_core_mismatch(ifs, Pgrid, flag_read_rank, data, nx, ny, nz, nx_read, ny_read, nz_read); +#else + ofs_running << " Read SPIN = " << is + 1 << " charge now." << std::endl; + if(nx == nx_read && ny == ny_read && nz == nz_read) + ModuleIO::read_cube_core_match(ifs, data, nx*ny, nz); + else + ModuleIO::read_cube_core_mismatch(ifs, data, nx, ny, nz, nx_read, ny_read, nz_read); +#endif - for (int iz = 0; iz < nz; iz++) - { - if (my_rank == 0 || (esolver_type == "sdft" && rank_in_stogroup == 0)) - { - zpiece = read_rho[iz]; - } - Pgrid->zpiece_to_all(zpiece, iz, data); - } // iz + return true; +} - if (my_rank == 0 || (esolver_type == "sdft" && rank_in_stogroup == 0)) +void ModuleIO::read_cube_core_match( + std::ifstream &ifs, +#ifdef __MPI + const Parallel_Grid*const Pgrid, + const bool flag_read_rank, +#endif + double*const data, + const int nxy, + const int nz) +{ +#ifdef __MPI + if (flag_read_rank) { + std::vector> read_rho(nz, std::vector(nxy)); + for (int ixy = 0; ixy < nxy; ixy++) + for (int iz = 0; iz < nz; iz++) + ifs >> read_rho[iz][ixy]; for (int iz = 0; iz < nz; iz++) - { - delete[] read_rho[iz]; - } - delete[] read_rho; + Pgrid->zpiece_to_all(read_rho[iz].data(), iz, data); } else { - delete[] zpiece; + std::vector zpiece(nxy); + for (int iz = 0; iz < nz; iz++) + Pgrid->zpiece_to_all(zpiece.data(), iz, data); } #else - ofs_running << " Read SPIN = " << is + 1 << " charge now." << std::endl; - if (same) + for (int ixy = 0; ixy < nxy; ixy++) + for (int iz = 0; iz < nz; iz++) + ifs >> data[iz * nxy + ixy]; +#endif +} + +void ModuleIO::read_cube_core_mismatch( + std::ifstream &ifs, +#ifdef __MPI + const Parallel_Grid*const Pgrid, + const bool flag_read_rank, +#endif + double*const data, + const int nx, + const int ny, + const int nz, + const int nx_read, + const int ny_read, + const int nz_read) +{ +#ifdef __MPI + const int nxy = nx * ny; + if (flag_read_rank) { - for (int i = 0; i < nx; i++) - { - for (int j = 0; j < ny; j++) - { - for (int k = 0; k < nz; k++) - { - ifs >> data[k * nx * ny + i * ny + j]; - } - } - } + std::vector> read_rho(nz, std::vector(nxy)); + ModuleIO::trilinear_interpolate(ifs, nx_read, ny_read, nz_read, nx, ny, nz, read_rho); + for (int iz = 0; iz < nz; iz++) + Pgrid->zpiece_to_all(read_rho[iz].data(), iz, data); } else { - ModuleIO::trilinear_interpolate(ifs, nx_read, ny_read, nz_read, nx, ny, nz, data); + std::vector zpiece(nxy); + for (int iz = 0; iz < nz; iz++) + Pgrid->zpiece_to_all(zpiece.data(), iz, data); } +#else + ModuleIO::trilinear_interpolate(ifs, nx_read, ny_read, nz_read, nx, ny, nz, data); #endif - - if (my_rank == 0 || (esolver_type == "sdft" && rank_in_stogroup == 0)) - ifs.close(); - return true; } void ModuleIO::trilinear_interpolate(std::ifstream& ifs, @@ -204,7 +204,7 @@ void ModuleIO::trilinear_interpolate(std::ifstream& ifs, const int& ny, const int& nz, #ifdef __MPI - double** data + std::vector> &data #else double* data #endif diff --git a/source/module_io/write_cube.cpp b/source/module_io/write_cube.cpp index 0cc283c6fc..2573740218 100644 --- a/source/module_io/write_cube.cpp +++ b/source/module_io/write_cube.cpp @@ -5,23 +5,23 @@ void ModuleIO::write_cube( #ifdef __MPI - const int& bz, - const int& nbz, - const int& nplane, - const int& startz_current, + const int bz, + const int nbz, + const int nplane, + const int startz_current, #endif - const double* data, - const int& is, - const int& nspin, - const int& iter, + const double*const data, + const int is, + const int nspin, + const int iter, const std::string& fn, - const int& nx, - const int& ny, - const int& nz, - const double& ef, - const UnitCell* ucell, - const int& precision, - const int& out_fermi) + const int nx, + const int ny, + const int nz, + const double ef, + const UnitCell*const ucell, + const int precision, + const int out_fermi) { ModuleBase::TITLE("ModuleIO", "write_cube"); @@ -127,8 +127,43 @@ void ModuleIO::write_cube( ofs_cube << std::scientific; } +#ifdef __MPI + ModuleIO::write_cube_core(ofs_cube, bz, nbz, nplane, startz_current, data, nx*ny, nz, 6); +#else + ModuleIO::write_cube_core(ofs_cube, data, nx*ny, nz, 6); +#endif + + if (my_rank == 0) + { + end = time(NULL); + ModuleBase::GlobalFunc::OUT_TIME("write_cube", start, end); + + /// for cube file + ofs_cube.close(); + } + + return; +} + + +void ModuleIO::write_cube_core( + std::ofstream &ofs_cube, +#ifdef __MPI + const int bz, + const int nbz, + const int nplane, + const int startz_current, +#endif + const double*const data, + const int nxy, + const int nz, + const int n_data_newline) +{ + ModuleBase::TITLE("ModuleIO", "write_cube_core"); + #ifdef __MPI + const int my_rank = GlobalV::MY_RANK; const int my_pool = GlobalV::MY_POOL; const int rank_in_pool = GlobalV::RANK_IN_POOL; const int nproc_in_pool = GlobalV::NPROC_IN_POOL; @@ -137,7 +172,7 @@ void ModuleIO::write_cube( if (my_pool == 0) { /// for cube file - const int nxyz = nx * ny * nz; + const int nxyz = nxy * nz; std::vector data_cube(nxyz, 0.0); // num_z: how many planes on processor 'ip' @@ -177,7 +212,6 @@ void ModuleIO::write_cube( } int count = 0; - const int nxy = nx * ny; std::vector zpiece(nxy, 0.0); // save the rho one z by one z. @@ -192,21 +226,21 @@ void ModuleIO::write_cube( // case 1: the first part of rho in processor 0. if (which_ip[iz] == 0 && rank_in_pool == 0) { - for (int ir = 0; ir < nxy; ir++) + for (int ixy = 0; ixy < nxy; ixy++) { // mohan change to rho_save on 2012-02-10 // because this can make our next restart calculation lead // to the same scf_thr as the one saved. - zpiece[ir] = data[ir * nplane + iz - startz_current]; + zpiece[ixy] = data[ixy * nplane + iz - startz_current]; } } // case 2: > first part rho: send the rho to // processor 0. else if (which_ip[iz] == rank_in_pool) { - for (int ir = 0; ir < nxy; ir++) + for (int ixy = 0; ixy < nxy; ixy++) { - zpiece[ir] = data[ir * nplane + iz - startz_current]; + zpiece[ixy] = data[ixy * nplane + iz - startz_current]; } MPI_Send(zpiece.data(), nxy, MPI_DOUBLE, 0, tag, POOL_WORLD); } @@ -221,9 +255,9 @@ void ModuleIO::write_cube( if (my_rank == 0) { /// for cube file - for (int ir = 0; ir < nxy; ir++) + for (int ixy = 0; ixy < nxy; ixy++) { - data_cube[ir + iz * nxy] = zpiece[ir]; + data_cube[ixy * nz + iz] = zpiece[ixy]; } /// for cube file } @@ -232,52 +266,35 @@ void ModuleIO::write_cube( // for cube file if (my_rank == 0) { - for (int ix = 0; ix < nx; ix++) + for (int ixy = 0; ixy < nxy; ixy++) { - for (int iy = 0; iy < ny; iy++) + for (int iz = 0; iz < nz; iz++) { - for (int iz = 0; iz < nz; iz++) + ofs_cube << " " << data_cube[ixy * nz + iz]; + if ((iz % n_data_newline == n_data_newline-1) && (iz != nz - 1)) { - ofs_cube << " " << data_cube[iz * nx * ny + ix * ny + iy]; - if (iz % 6 == 5 && iz != nz - 1) - { - ofs_cube << "\n"; - } + ofs_cube << "\n"; } - ofs_cube << "\n"; } + ofs_cube << "\n"; } } /// for cube file } MPI_Barrier(MPI_COMM_WORLD); #else - for (int i = 0; i < nx; i++) + for (int ixy = 0; ixy < nxy; ixy++) { - for (int j = 0; j < ny; j++) + for (int iz = 0; iz < nz; iz++) { - for (int k = 0; k < nz; k++) + ofs_cube << " " << data[iz * nxy + ixy]; + // ++count_cube; + if ((iz % n_data_newline == n_data_newline-1) && (iz != nz - 1)) { - ofs_cube << " " << data[k * nx * ny + i * ny + j]; - // ++count_cube; - if (k % 6 == 5 && k != nz - 1) - { - ofs_cube << "\n"; - } + ofs_cube << "\n"; } - ofs_cube << "\n"; } + ofs_cube << "\n"; } #endif - - if (my_rank == 0) - { - end = time(NULL); - ModuleBase::GlobalFunc::OUT_TIME("write_cube", start, end); - - /// for cube file - ofs_cube.close(); - } - - return; }