Write output netCDF files in parallel

nextsimhub · Dec 4, 2023 · c64bb14 · c64bb14
1 parent 6b41b8f
commit c64bb14
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 3 deletions.
diff --git a/core/src/CMakeLists.txt b/core/src/CMakeLists.txt
@@ -26,6 +26,10 @@ set(BaseSources
     "${ModelArrayStructure}/ModelArrayDetails.cpp"
     )
 
+set(ParallelNetCDFSources
+    "${CMAKE_CURRENT_SOURCE_DIR}/ParallelNetcdfFile.cpp"
+    )
+
 list(TRANSFORM BaseSources PREPEND "${CMAKE_CURRENT_SOURCE_DIR}/")
 
 set(ModuleDir "${CMAKE_CURRENT_SOURCE_DIR}/modules")
@@ -66,8 +70,25 @@ set(NextsimSources
     "${NextsimSources}"
     "${BaseSources}"
     "${ModuleSources}"
+    "${ParallelNetCDFSources}"
     PARENT_SCOPE)
 
+if(ENABLE_MPI)
+    set(NextsimSources
+        "${NextsimSources}"
+        "${BaseSources}"
+        "${ModuleSources}"
+        "${ParallelNetCDFSources}"
+        PARENT_SCOPE)
+else()
+    set(NextsimSources
+        "${NextsimSources}"
+        "${BaseSources}"
+        "${ModuleSources}"
+        PARENT_SCOPE)
+endif()
+
+
 set(NextsimIncludeDirs
     "${NextsimIncludeDirs}"
     "${ModuleDir}"

diff --git a/core/src/RectGridIO.cpp b/core/src/RectGridIO.cpp
@@ -23,6 +23,9 @@
 #include <ncVar.h>
 
 #include <algorithm>
+#ifdef USE_MPI
+#include <include/ParallelNetcdfFile.hpp>
+#endif
 #include <list>
 #include <map>
 #include <string>
@@ -177,8 +180,7 @@ void RectGridIO::dumpModelState(const ModelState& state, const ModelMetadata& me
     const std::string& filePath, bool isRestart) const
 {
 #ifdef USE_MPI
-    auto filePathRank = filePath + "_" + std::to_string(metadata.mpiMyRank);
-    netCDF::NcFile ncFile(filePathRank, netCDF::NcFile::replace);
+    netCDF::NcFilePar ncFile(filePath, netCDF::NcFile::replace, metadata.mpiComm);
 #else
     netCDF::NcFile ncFile(filePath, netCDF::NcFile::replace);
 #endif
@@ -198,22 +200,47 @@ void RectGridIO::dumpModelState(const ModelState& state, const ModelMetadata& me
 
     // Create the dimension data, since it has to be in the same group as the
     // data or the parent group
+#ifdef USE_MPI
+    netCDF::NcDim xDim = dataGroup.addDim(dimensionNames[0], metadata.globalExtentX);
+    netCDF::NcDim yDim = dataGroup.addDim(dimensionNames[1], metadata.globalExtentY);
+#else
     netCDF::NcDim xDim = dataGroup.addDim(dimensionNames[0], nx);
     netCDF::NcDim yDim = dataGroup.addDim(dimensionNames[1], ny);
-    std::vector<netCDF::NcDim> dims2 = { yDim, xDim };
+#endif
     netCDF::NcDim zDim = dataGroup.addDim(dimensionNames[2], nz);
+    std::vector<netCDF::NcDim> dims2 = { yDim, xDim };
     std::vector<netCDF::NcDim> dims3 = { zDim, yDim, xDim };
+#ifdef USE_MPI
+    // Set the origins and extensions for reading 2D data based
+    // on MPI decomposition
+    std::vector<size_t> start(3);
+    std::vector<size_t> size(3);
+    start[0] = metadata.localCornerX;
+    start[1] = metadata.localCornerY;
+    start[2] = 0;
+    size[0] = metadata.localExtentX;
+    size[1] = metadata.localExtentY;
+    size[2] = nz;
+#endif
 
     for (const auto entry : state.data) {
         const std::string& name = entry.first;
         if (entry.second.getType() == ModelArray::Type::H && entry.second.trueSize() > 0) {
             netCDF::NcVar var(dataGroup.addVar(name, netCDF::ncDouble, dims2));
             var.putAtt(mdiName, netCDF::ncDouble, MissingData::value);
+#ifdef USE_MPI
+            var.putVar(start, size, entry.second.getData());
+#else
             var.putVar(entry.second.getData());
+#endif
         } else if (entry.second.getType() == ModelArray::Type::Z && entry.second.trueSize() > 0) {
             netCDF::NcVar var(dataGroup.addVar(name, netCDF::ncDouble, dims3));
             var.putAtt(mdiName, netCDF::ncDouble, MissingData::value);
+#ifdef USE_MPI
+            var.putVar(start, size, entry.second.getData());
+#else
             var.putVar(entry.second.getData());
+#endif
         }
     }