feat(task356, task3): implement MPI-IO parallel cube file write

source/source_io/module_output/cube_io.h

@@ -3,6 +3,11 @@
 #include "source_cell/unitcell.h"
 
 #include <string>
+
+#ifdef __MPI
+#include <mpi.h>
+#endif
+
 class Parallel_Grid;
 
 namespace ModuleIO
@@ -101,6 +106,26 @@ void trilinear_interpolate(const double* const data_in,
                            const int& ny,
                            const int& nz,
                            double* data_out);
+
+/// MPI-IO parallel cube file write. All ranks must have the full data array.
+/// Each rank writes its z-slice range via collective MPI-IO.
+#ifdef __MPI
+void write_cube_mpi(const std::string& file,
+                    const std::vector<std::string>& comment,
+                    const int& natom,
+                    const std::vector<double>& origin,
+                    const int& nx, const int& ny, const int& nz,
+                    const std::vector<double>& dx,
+                    const std::vector<double>& dy,
+                    const std::vector<double>& dz,
+                    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 int precision,
+                    const MPI_Comm& comm);
+#endif
+
 } // namespace ModuleIO
 
 #endif

source/source_io/module_output/read_cube.cpp

@@ -198,11 +198,44 @@ bool ModuleIO::read_cube(const std::string& file,
 
     const int nxyz = nx * ny * nz;
     data.resize(nxyz);
-    for (int i = 0;i < nxyz;++i) 
-    { 
-	    ifs >> data[i]; 
+
+    // Check for MPI binary format marker (CIPM in little-endian)
+    static constexpr uint32_t CUBE_MPI_MARKER = 0x4D504943; // "CIPM"
+
+    ifs >> std::ws;
+    const std::streampos data_start = ifs.tellg();
+
+    // Peek at the first 4 bytes to detect MPI binary format
+    char peek_buf[4] = {};
+    ifs.read(peek_buf, 4);
+    uint32_t magic = 0;
+    std::memcpy(&magic, peek_buf, 4);
+
+    if (magic == CUBE_MPI_MARKER)
+    {
+        // MPI-parallel binary format: 4-byte marker + nxyz doubles in binary
+        ifs.close();
+
+        std::ifstream ifs_bin(file, std::ios::binary);
+        if (ifs_bin)
+        {
+            ifs_bin.seekg(data_start + static_cast<std::streampos>(4));
+            ifs_bin.read(reinterpret_cast<char*>(data.data()),
+                         nxyz * static_cast<std::streamsize>(sizeof(double)));
+        }
+        ifs_bin.close();
+    }
+    else
+    {
+        // Text format: rewind and parse as floating-point numbers
+        ifs.clear();
+        ifs.seekg(data_start);
+        for (int i = 0; i < nxyz; ++i)
+        {
+            ifs >> data[i];
+        }
+        ifs.close();
     }
 
-    ifs.close();
     return true;
 }

source/source_io/module_output/write_cube.cpp

@@ -258,3 +258,133 @@ void ModuleIO::write_cube(const std::string& file,
     }
     ofs.close();
 }
+#ifdef __MPI
+
+// Binary marker to identify MPI-parallel binary cube files
+// "CIPM" = Cube I/O Parallel Marker (little-endian)
+static constexpr uint32_t CUBE_MPI_MARKER = 0x4D504943;
+
+void ModuleIO::write_cube_mpi(const std::string& file,
+                               const std::vector<std::string>& comment,
+                               const int& natom,
+                               const std::vector<double>& origin,
+                               const int& nx,
+                               const int& ny,
+                               const int& nz,
+                               const std::vector<double>& dx,
+                               const std::vector<double>& dy,
+                               const std::vector<double>& dz,
+                               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 int precision,
+                               const MPI_Comm& comm)
+{
+    assert(comment.size() >= 2);
+    assert(origin.size() >= 3);
+    assert(data.size() >= static_cast<size_t>(nx) * ny * nz);
+
+    int nprocs = 1;
+    int my_rank = 0;
+    MPI_Comm_size(comm, &nprocs);
+    MPI_Comm_rank(comm, &my_rank);
+
+    const int nxy = nx * ny;
+
+    // Compute z-slice distribution across ranks
+    int nz_local = nz / nprocs;
+    int remainder = nz % nprocs;
+    int z_start = 0;
+    for (int r = 0; r < my_rank; ++r)
+    {
+        z_start += nz_local + (r < remainder ? 1 : 0);
+    }
+    if (my_rank < remainder)
+        nz_local += 1;
+
+    const int my_count = nz_local * nxy;
+
+    // Build the text header string (rank 0 only) — identical to standard cube header
+    std::string header_str;
+    size_t header_bytes = 0;
+
+    if (my_rank == 0)
+    {
+        std::ostringstream hdr;
+        hdr << std::fixed;
+
+        for (int i = 0; i < 2; ++i)
+            hdr << comment[i] << "\n";
+
+        hdr << std::setprecision(1);
+        hdr << natom << " " << origin[0] << " " << origin[1] << " " << origin[2] << " \n";
+
+        hdr << std::setprecision(6);
+        hdr << nx << " " << dx[0] << " " << dx[1] << " " << dx[2] << "\n";
+        hdr << ny << " " << dy[0] << " " << dy[1] << " " << dy[2] << "\n";
+        hdr << nz << " " << dz[0] << " " << dz[1] << " " << dz[2] << "\n";
+
+        for (int i = 0; i < natom; ++i)
+        {
+            hdr << " " << atom_type[i] << " " << atom_charge[i] << " "
+                << atom_pos[i][0] << " " << atom_pos[i][1] << " " << atom_pos[i][2] << "\n";
+        }
+
+        header_str = hdr.str();
+        header_bytes = header_str.size();
+    }
+
+    MPI_Bcast(&header_bytes, 1, MPI_UNSIGNED_LONG, 0, comm);
+
+    // Open file with MPI-IO
+    MPI_File fh;
+    MPI_File_open(comm, file.c_str(),
+                  MPI_MODE_WRONLY | MPI_MODE_CREATE,
+                  MPI_INFO_NULL, &fh);
+
+    // Rank 0 writes the text header + binary marker
+    if (my_rank == 0)
+    {
+        MPI_File_write_at(fh, 0, header_str.data(), static_cast<int>(header_bytes),
+                          MPI_CHAR, MPI_STATUS_IGNORE);
+        // Write binary marker so readers can detect MPI-parallel binary format
+        MPI_File_write_at(fh, static_cast<MPI_Offset>(header_bytes),
+                          &CUBE_MPI_MARKER, 1, MPI_UINT32_T, MPI_STATUS_IGNORE);
+    }
+
+    // Barrier to ensure header is written before data
+    MPI_Barrier(comm);
+
+    // Each rank writes its z-slice data using MPI file views for correct
+    // z-fastest ordering. File layout: for each ixy, all nz values
+    // consecutively. Each rank writes its nz_local values within each nz block.
+    MPI_Offset data_offset = static_cast<MPI_Offset>(header_bytes) + sizeof(uint32_t);
+
+    // Create strided file view: for each ixy row, this rank writes nz_local
+    // values with a stride of nz (skipping data written by other ranks).
+    MPI_Datatype filetype;
+    MPI_Type_vector(nxy, nz_local, nz, MPI_DOUBLE, &filetype);
+    MPI_Type_commit(&filetype);
+
+    // Set file view: the view starts at this rank's z offset within each block
+    MPI_File_set_view(fh, data_offset + z_start * sizeof(double),
+                      MPI_DOUBLE, filetype, "native", MPI_INFO_NULL);
+
+    // Pack local buffer in row-major order matching the file view
+    std::vector<double> buf(my_count);
+    for (int ixy = 0; ixy < nxy; ++ixy)
+    {
+        for (int iz = 0; iz < nz_local; ++iz)
+        {
+            buf[ixy * nz_local + iz] = data[ixy * nz + (z_start + iz)];
+        }
+    }
+
+    MPI_File_write_all(fh, buf.data(), my_count, MPI_DOUBLE, MPI_STATUS_IGNORE);
+    MPI_Type_free(&filetype);
+
+    MPI_File_close(&fh);
+}
+
+#endif

source/source_io/test/CMakeLists.txt

@@ -285,6 +285,17 @@ AddTest(
 )
 endif()
 
+AddTest(
+  TARGET MODULE_IO_write_cube_test
+  LIBS parameter ${math_libs} base device planewave
+  SOURCES write_cube_test.cpp ../module_output/write_cube.cpp ../module_output/read_cube.cpp ../../source_pw/module_pwdft/parallel_grid.cpp ../../source_basis/module_pw/test/test_tool.cpp
+)
+
+add_test(NAME MODULE_IO_write_cube_test_parallel
+      COMMAND mpirun -np 4 ./MODULE_IO_write_cube_test
+      WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
+)
+
 AddTest(
   TARGET MODULE_IO_write_elf_logic_test
   SOURCES write_elf_logic_test.cpp