Refactor IO cube (Useful Information: excellent example for discussions and actions among ABACUS developers)

source/module_elecstate/module_charge/symmetry_rho.cpp

@@ -127,24 +127,8 @@ void Symmetry_rho::psymm(double* rho_part,
     }
 
     // (3)
-    const int ncxy = rho_basis->nx * rho_basis->ny;
-    std::vector<double> zpiece(ncxy);
-    for (int iz = 0; iz < rho_basis->nz; iz++)
-    {
-        ModuleBase::GlobalFunc::ZEROS(zpiece.data(), ncxy);
-        if (GlobalV::MY_RANK == 0)
-        {
-            for (int ix = 0; ix < rho_basis->nx; ix++)
-            {
-                for (int iy = 0; iy < rho_basis->ny; iy++)
-                {
-                    const int ir = ix * rho_basis->ny + iy;
-                    zpiece[ir] = rhotot[ix * rho_basis->ny * rho_basis->nz + iy * rho_basis->nz + iz];
-                }
-            }
-        }
-        Pgrid.zpiece_to_all(zpiece.data(), iz, rho_part);
-    }
+Pgrid.bcast(rhotot.data(), rho_part);
+
 #endif
     return;
 }

source/module_hamilt_pw/hamilt_pwdft/parallel_grid.cpp

@@ -185,7 +185,28 @@ void Parallel_Grid::z_distribution()
 
 
 #ifdef __MPI
-void Parallel_Grid::zpiece_to_all(double *zpiece, const int &iz, double *rho) const
+void Parallel_Grid::bcast(const double* const data_global, double* data_local)const
+{
+    std::vector<double> zpiece(ncxy);
+    for (int iz = 0; iz < this->ncz; ++iz)
+    {
+        ModuleBase::GlobalFunc::ZEROS(zpiece.data(), ncxy);
+        if (GlobalV::MY_RANK == 0)
+        {
+            for (int ix = 0; ix < ncx; ix++)
+            {
+                for (int iy = 0; iy < ncy; iy++)
+                {
+                    const int ixy = ix * ncy + iy;
+                    zpiece[ixy] = data_global[ixy * ncz + iz];
+                }
+            }
+        }
+        this->zpiece_to_all(zpiece.data(), iz, data_local);
+    }
+}
+
+void Parallel_Grid::zpiece_to_all(double* zpiece, const int& iz, double* rho) const
 {
 	if(PARAM.inp.esolver_type == "sdft")
 	{

source/module_hamilt_pw/hamilt_pwdft/parallel_grid.h

@@ -27,9 +27,12 @@ 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) const;
-	void zpiece_to_stogroup(double *zpiece, const int &iz, double *rho) const; //qainrui add for sto-dft 2021-7-21
-
+    void zpiece_to_stogroup(double* zpiece, const int& iz, double* rho) const; //qainrui add for sto-dft 2021-7-21
+    void zpiece_to_all(double* zpiece, const int& iz, double* rho) const;
+
+    /// @brief  Broadcast data from root to all processors. The index order is [x][y][z].
+    void bcast(const double* const data_global, double* data_local)const;
+    /// @brief  Reduce data from all processors to root. The index order is [x][y][z].
     void reduce(double* rhotot, const double* constrhoin)const;
 #endif
 

source/module_io/cube_io.h

@@ -26,9 +26,20 @@ namespace ModuleIO
     const int precision = 11,
     const int out_fermi = 1); // mohan add 2007-10-17
 
-struct CubeInfo
-{
-    CubeInfo(
+    /// read the full data from a cube file 
+    bool read_cube(const std::string& file,
+        std::vector<std::string>& comment,
+        int& natom,
+        std::vector<double>& cel_pos,
+        std::vector<int>& nvoxel,
+        std::vector<std::vector<double>>& axis_vecs,
+        std::vector<int>& atom_type,
+        std::vector<double>& atom_charge,
+        std::vector<std::vector<double>>& atom_pos,
+        std::vector<double>& data);
+
+    /// write a cube file
+    void write_cube(const std::string& file,
         const std::vector<std::string>& comment,
         const int natom,
         const std::vector<double>& cel_pos,
@@ -38,34 +49,8 @@ struct CubeInfo
         const std::vector<double>& atom_charge,
         const std::vector<std::vector<double>>& atom_pos,
         const std::vector<double>& data,
-        const bool valid)
-        : comment(comment), natom(natom), cel_pos(cel_pos),
-        nvoxel(nvoxel), axis_vecs(axis_vecs),
-        atom_type(atom_type), atom_charge(atom_charge), atom_pos(atom_pos),
-        data(data), valid(valid)
-    {
-    };
-
-    const std::vector<std::string> comment = {};
-    const int natom = 0;
-    const std::vector<double> cel_pos = {};
-    const std::vector<int> nvoxel = {};
-    const std::vector<std::vector<double>> axis_vecs = {};
-    const std::vector<int> atom_type = {};
-    const std::vector<double> atom_charge = {};
-    const std::vector<std::vector<double>> atom_pos = {};
-    const std::vector<double> data = {};
-    const bool valid = false;
-};
-
-/// read the full data from a cube file 
-CubeInfo read_cube(const std::string& file);
-/// write a cube file
-void write_cube(const std::string& file, const CubeInfo& info, const int precision, const int ndata_line = 6);
-/// change the index order: [x][y][z] (.cube file) -> [z][x][y] (ABACUS)
-void xyz2zxy(const double* const xyz, const int nx, const int ny, const int nz, double* const zxy);
-/// change the index order: [z][x][y] (ABACUS) -> [x][y][z] (.cube file)
-void zxy2xyz(const double* const zxy, const int nx, const int ny, const int nz, double* const xyz);
+        const int precision,
+        const int ndata_line = 6);
 
     /**
      * @brief The trilinear interpolation method

source/module_io/read_cube.cpp

@@ -2,19 +2,6 @@
 #include <limits>
 #include "module_hamilt_pw/hamilt_pwdft/parallel_grid.h"
 // #include "module_base/global_variable.h" // GlobalV reference removed
-void ModuleIO::xyz2zxy(const double* const xyz, const int nx, const int ny, const int nz, double* const zxy)
-{
-    for (int ix = 0; ix < nx; ix++)
-    {
-        for (int iy = 0; iy < ny; iy++)
-        {
-            for (int iz = 0; iz < nz; iz++)
-            {
-                zxy[(iz * nx + ix) * ny + iy] = xyz[(ix * ny + iy) * nz + iz];
-            }
-        }
-    }
-}
 
 bool ModuleIO::read_grid(
     const Parallel_Grid& pgrid,
@@ -44,34 +31,40 @@ bool ModuleIO::read_grid(
     const int& ny = pgrid.ny;
     const int& nz = pgrid.nz;
     const int& nxyz = nx * ny * nz;
-    std::vector<double> data_zxy_full(nxyz, 0.0);    // [iz][ix][iy]
+    std::vector<double> data_xyz_full(nxyz, 0.0);
     if (my_rank == 0)
     {
-        const CubeInfo& cube_info = ModuleIO::read_cube(fn);
-
-        const int& nx_read = cube_info.nvoxel[0];
-        const int& ny_read = cube_info.nvoxel[1];
-        const int& nz_read = cube_info.nvoxel[2];
+        std::vector<std::string> comment;
+        int natom = 0;
+        std::vector<double> cel_pos;
+        std::vector<int> nvoxel;
+        std::vector<std::vector<double>> axis_vecs;
+        std::vector<int> atom_type;
+        std::vector<double> atom_charge;
+        std::vector<std::vector<double>> atom_pos;
+        std::vector<double> data_read;
+        bool valid = ModuleIO::read_cube(fn, comment, natom, cel_pos, nvoxel, axis_vecs, atom_type, atom_charge, atom_pos, data_read);
+
+        const int& nx_read = nvoxel[0];
+        const int& ny_read = nvoxel[1];
+        const int& nz_read = nvoxel[2];
 
         // if mismatch, trilinear interpolate
         if (nx == nx_read && ny == ny_read && nz == nz_read)
         {
-            ModuleIO::xyz2zxy(cube_info.data.data(), nx, ny, nz, data_zxy_full.data());
+            std::memcpy(data_xyz_full.data(), data_read.data(), nxyz * sizeof(double));
         }
         else
         {
-            std::vector<double> data_xyz_full(nxyz, 0.0);
-            trilinear_interpolate(cube_info.data.data(), nx_read, ny_read, nz_read, nx, ny, nz, data_xyz_full.data());
-            ModuleIO::xyz2zxy(data_xyz_full.data(), nx, ny, nz, data_zxy_full.data());
+            trilinear_interpolate(data_read.data(), nx_read, ny_read, nz_read, nx, ny, nz, data_xyz_full.data());
         }
     }
 
     // distribute
 #ifdef __MPI 
-    const int nxy = nx * ny;
-    for (int iz = 0;iz < nz;++iz) { pgrid.zpiece_to_all(data_zxy_full.data() + iz * nxy, iz, data); }
+    pgrid.bcast(data_xyz_full.data(), data);
 #else
-    std::memcpy(data, data_zxy_full.data(), nxyz * sizeof(double));
+    std::memcpy(data, data_xyz_full.data(), nxyz * sizeof(double));
 #endif
     return true;
 }
@@ -147,43 +140,51 @@ void ModuleIO::trilinear_interpolate(
     delete[] read_rho;
 }
 
-ModuleIO::CubeInfo ModuleIO::read_cube(const std::string& file)
+bool ModuleIO::read_cube(const std::string& file,
+    std::vector<std::string>& comment,
+    int& natom,
+    std::vector<double>& cell_pos,
+    std::vector<int>& nvoxel,
+    std::vector<std::vector<double>>& axis_vecs,
+    std::vector<int>& atom_type,
+    std::vector<double>& atom_charge,
+    std::vector<std::vector<double>>& atom_pos,
+    std::vector<double>& data)
 {
     std::ifstream ifs(file);
 
-    if (!ifs) { return ModuleIO::CubeInfo({}, 0, {}, {}, {}, {}, {}, {}, {}, false); }
+    if (!ifs) { return false; }
 
-    std::vector<std::string> comment(2);
+    comment.resize(2);
     for (auto& c : comment) { std::getline(ifs, c); }
 
-    int natom;
     ifs >> natom;
-    std::vector<double> cell_pos(3);
+    cell_pos.resize(3);
     for (auto& cp : cell_pos) { ifs >> cp; }
 
-    std::vector<int> nvoxel(3);
-    std::vector<std::vector<double>> axis_vecs(3);
+    nvoxel.resize(3);
+    axis_vecs.resize(0);
     for (int i = 0;i < 3;++i)
     {
         std::vector<double> vec(3);
         ifs >> nvoxel[i] >> vec[0] >> vec[1] >> vec[2];
         axis_vecs.push_back(vec);
     }
 
-    std::vector<int> itype(natom);
-    std::vector<double> charge(natom);
-    std::vector<std::vector<double>> atom_pos(natom);
+    atom_type.resize(natom);
+    atom_charge.resize(natom);
+    atom_pos.resize(0);
     for (int i = 0;i < natom;++i)
     {
         std::vector<double> apos(3);
-        ifs >> itype[i] >> charge[i] >> apos[0] >> apos[1] >> apos[2];
+        ifs >> atom_type[i] >> atom_charge[i] >> apos[0] >> apos[1] >> apos[2];
         atom_pos.push_back(apos);
     }
 
     const int nxyz = nvoxel[0] * nvoxel[1] * nvoxel[2];
-    std::vector<double> data(nxyz);
+    data.resize(nxyz);
     for (int i = 0;i < nxyz;++i) { ifs >> data[i]; }
 
     ifs.close();
-    return ModuleIO::CubeInfo(comment, natom, cell_pos, nvoxel, axis_vecs, itype, charge, atom_pos, std::move(data), true);
+    return true;
 }

source/module_io/test_serial/rho_io_test.cpp

@@ -123,11 +123,27 @@ TEST_F(RhoIOTest, TrilinearInterpolate)
             }
         }
     }
+
+    // The old implementation is inconsistent: ifdef MPI, [x][y][z]; else, [z][x][y].
+    // Now we use [x][y][z] for both MPI and non-MPI, so here we need to chage the index order.
+    auto permute_xyz2zxy = [&](const double* const xyz, double* const zxy) -> void
+        {
+            for (int ix = 0; ix < nx; ix++)
+            {
+                for (int iy = 0; iy < ny; iy++)
+                {
+                    for (int iz = 0; iz < nz; iz++)
+                    {
+                        zxy[(iz * nx + ix) * ny + iy] = xyz[(ix * ny + iy) * nz + iz];
+                    }
+                }
+            }
+        };
     const int nxyz = nx * ny * nz;
     std::vector<double> data_xyz(nxyz);
     std::vector<double> data(nxyz); // z > x > y
     ModuleIO::trilinear_interpolate(data_read.data(), nx_read, ny_read, nz_read, nx, ny, nz, data_xyz.data());
-    ModuleIO::xyz2zxy(data_xyz.data(), nx, ny, nz, data.data());
+    permute_xyz2zxy(data_xyz.data(), data.data());
     EXPECT_DOUBLE_EQ(data[0], 0.0010824725010374092);
     EXPECT_DOUBLE_EQ(data[10], 0.058649850374240906);
     EXPECT_DOUBLE_EQ(data[100], 0.018931708073604996);