WIP: Add OpenPMD support

CMakeLists.txt

@@ -31,11 +31,12 @@ include(CTest)
 # Compile time constants
 
 # Compile Options
-option(PARTHENON_SINGLE_PRECISION "Run in single precision" OFF)
+option(PARTHENON_SINGLE_PRENSION "Run in single precision" OFF)
 option(PARTHENON_DISABLE_MPI "MPI is enabled by default if found, set this to True to disable MPI" OFF)
 option(PARTHENON_ENABLE_HOST_COMM_BUFFERS "CUDA/HIP Only: Allocate communication buffers on host (may be slower)" OFF)
 option(PARTHENON_DISABLE_HDF5 "HDF5 is enabled by default if found, set this to True to disable HDF5" OFF)
 option(PARTHENON_DISABLE_HDF5_COMPRESSION "HDF5 compression is enabled by default, set this to True to disable compression in HDF5 output/restart files" OFF)
+option(PARTHENON_ENABLE_OPENPMD "OpenPMD is enabled by default if found, set this to True to disable OpenPMD" ON)
 option(PARTHENON_DISABLE_SPARSE "Sparse capability is enabled by default, set this to True to compile-time disable all sparse capability" OFF)
 option(PARTHENON_ENABLE_ASCENT "Enable Ascent for in situ visualization and analysis" OFF)
 option(PARTHENON_LINT_DEFAULT "Linting is turned off by default, use the \"lint\" target or set \
@@ -187,6 +188,25 @@ if (NOT PARTHENON_DISABLE_HDF5)
   install(TARGETS HDF5_C EXPORT parthenonTargets)
 endif()
 
+if (PARTHENON_ENABLE_OPENPMD)
+#TODO(pgrete) add logic for serial/parallel
+#TODO(pgrete) add logic for internal/external build
+  include(FetchContent)
+  set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
+  set(openPMD_BUILD_CLI_TOOLS OFF)
+  set(openPMD_BUILD_EXAMPLES OFF)
+  set(openPMD_BUILD_TESTING OFF)
+  set(openPMD_BUILD_SHARED_LIBS OFF)  # precedence over BUILD_SHARED_LIBS if needed
+  set(openPMD_INSTALL OFF)            # or instead use:
+  # set(openPMD_INSTALL ${BUILD_SHARED_LIBS})  # only install if used as a shared library
+  set(openPMD_USE_PYTHON OFF)
+  FetchContent_Declare(openPMD
+    GIT_REPOSITORY "https://github.com/openPMD/openPMD-api.git"
+    GIT_TAG        "0.15.2")
+  FetchContent_MakeAvailable(openPMD)
+  install(TARGETS openPMD EXPORT parthenonTargets)
+endif()
+
 # Kokkos recommendatation resulting in not using default GNU extensions
 set(CMAKE_CXX_EXTENSIONS OFF)
 

src/CMakeLists.txt

@@ -193,10 +193,14 @@ add_library(parthenon
   outputs/parthenon_hdf5_types.hpp
   outputs/parthenon_xdmf.cpp
   outputs/parthenon_hdf5.hpp
+  outputs/parthenon_opmd.cpp
+  outputs/parthenon_opmd.hpp
   outputs/parthenon_xdmf.hpp
   outputs/restart.hpp
   outputs/restart_hdf5.cpp
   outputs/restart_hdf5.hpp
+  outputs/restart_opmd.cpp
+  outputs/restart_opmd.hpp
   outputs/vtk.cpp
 
   parthenon/driver.hpp
@@ -312,6 +316,10 @@ if (ENABLE_HDF5)
   target_link_libraries(parthenon PUBLIC HDF5_C)
 endif()
 
+if (PARTHENON_ENABLE_OPENPMD)
+  target_link_libraries(parthenon PUBLIC openPMD::openPMD)
+endif()
+
 # For Cuda with NVCC (<11.2) and C++17 Kokkos currently does not work/compile with
 # relaxed-constexpr, see https://github.com/kokkos/kokkos/issues/3496
 # However, Parthenon heavily relies on it and there is no harm in compiling Kokkos

src/config.hpp.in

@@ -45,6 +45,8 @@
 // defne ENABLE_HDF5 or not at all
 #cmakedefine ENABLE_HDF5
 
+#cmakedefine PARTHENON_ENABLE_OPENPMD
+
 // define PARTHENON_DISABLE_HDF5_COMPRESSION or not at all
 #cmakedefine PARTHENON_DISABLE_HDF5_COMPRESSION
 

src/interface/params.hpp

@@ -118,6 +118,12 @@ class Params {
     return it->second;
   }
 
+  const Mutability &GetMutability(const std::string &key) const {
+    auto const it = myMutable_.find(key);
+    PARTHENON_REQUIRE_THROWS(it != myMutable_.end(), "Key " + key + " doesn't exist");
+    return it->second;
+  }
+
   std::vector<std::string> GetKeys() const {
     std::vector<std::string> keys;
     for (auto &x : myParams_) {

src/outputs/output_utils.cpp

@@ -15,6 +15,7 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 
+#include <cstdint>
 #include <map>
 #include <set>
 #include <string>
@@ -29,6 +30,7 @@
 #include "mesh/mesh.hpp"
 #include "mesh/meshblock.hpp"
 #include "outputs/output_utils.hpp"
+#include "utils/mpi_types.hpp"
 
 namespace parthenon {
 namespace OutputUtils {
@@ -241,6 +243,45 @@ std::vector<int> ComputeIDsAndFlags(Mesh *pm) {
       });
 }
 
+template <typename T>
+std::vector<T> FlattendedLocalToGlobal(Mesh *pm, const std::vector<T> &data_local) {
+  const int n_blocks_global = pm->nbtotal;
+  const int n_blocks_local = static_cast<int>(pm->block_list.size());
+
+  const int n_elem = data_local.size() / n_blocks_local;
+  PARTHENON_REQUIRE_THROWS(data_local.size() % n_blocks_local == 0,
+                           "Results from flattened input vector does not evenly divide "
+                           "into number of local blocks.");
+  std::vector<T> data_global(n_elem * n_blocks_global);
+
+  std::vector<int> counts(Globals::nranks);
+  std::vector<int> offsets(Globals::nranks);
+
+  const auto &nblist = pm->GetNbList();
+  counts[0] = n_elem * nblist[0];
+  offsets[0] = 0;
+  for (int r = 1; r < Globals::nranks; r++) {
+    counts[r] = n_elem * nblist[r];
+    offsets[r] = offsets[r - 1] + counts[r - 1];
+  }
+
+#ifdef MPI_PARALLEL
+  PARTHENON_MPI_CHECK(MPI_Allgatherv(data_local.data(), counts[Globals::my_rank],
+                                     MPITypeMap<T>::type(), data_global.data(),
+                                     counts.data(), offsets.data(), MPITypeMap<T>::type(),
+                                     MPI_COMM_WORLD));
+#else
+  return data_local;
+#endif
+  return data_global;
+}
+
+// explicit template instantiation
+template std::vector<int64_t>
+FlattendedLocalToGlobal(Mesh *pm, const std::vector<int64_t> &data_local);
+template std::vector<int> FlattendedLocalToGlobal(Mesh *pm,
+                                                  const std::vector<int> &data_local);
+
 // TODO(JMM): I could make this use the other loop
 // functionality/high-order functions.  but it was more code than this
 // for, I think, little benefit.
@@ -313,5 +354,34 @@ std::size_t MPISum(std::size_t val) {
   return val;
 }
 
+VariableVector<Real> GetVarsToWrite(const std::shared_ptr<MeshBlock> pmb,
+                                    const bool restart,
+                                    const std::vector<std::string> &variables) {
+  const auto &var_vec = pmb->meshblock_data.Get()->GetVariableVector();
+  auto vars_to_write = GetAnyVariables(var_vec, variables);
+  if (restart) {
+    // get all vars with flag Independent OR restart
+    auto restart_vars = GetAnyVariables(
+        var_vec, {parthenon::Metadata::Independent, parthenon::Metadata::Restart});
+    for (auto restart_var : restart_vars) {
+      vars_to_write.emplace_back(restart_var);
+    }
+  }
+  return vars_to_write;
+}
+
+std::vector<VarInfo> GetAllVarsInfo(const VariableVector<Real> &vars,
+                                    const IndexShape &cellbounds) {
+  std::vector<VarInfo> all_vars_info;
+  for (auto &v : vars) {
+    all_vars_info.emplace_back(v, cellbounds);
+  }
+
+  // sort alphabetically
+  std::sort(all_vars_info.begin(), all_vars_info.end(),
+            [](const VarInfo &a, const VarInfo &b) { return a.label < b.label; });
+  return all_vars_info;
+}
+
 } // namespace OutputUtils
 } // namespace parthenon

src/outputs/output_utils.hpp

@@ -339,13 +339,29 @@ std::vector<Real> ComputeXminBlocks(Mesh *pm);
 std::vector<int64_t> ComputeLocs(Mesh *pm);
 std::vector<int> ComputeIDsAndFlags(Mesh *pm);
 
+// Takes a vector containing flattened data of all rank local blocks and returns the
+// flattened data over all blocks.
+template <typename T>
+std::vector<T> FlattendedLocalToGlobal(Mesh *pm, const std::vector<T> &data_local);
+
 // TODO(JMM): Potentially unsafe if MPI_UNSIGNED_LONG_LONG isn't a size_t
 // however I think it's probably safe to assume we'll be on systems
 // where this is the case?
 // TODO(JMM): If we ever need non-int need to generalize
 std::size_t MPIPrefixSum(std::size_t local, std::size_t &tot_count);
 std::size_t MPISum(std::size_t local);
 
+// Return all variables to write, i.e., for restarts all indpendent variables and ones
+// with explicit Restart flag, but also variables explicitly defined to output in the
+// input file.
+VariableVector<Real> GetVarsToWrite(const std::shared_ptr<MeshBlock> pmb,
+                                    const bool restart,
+                                    const std::vector<std::string> &variables);
+
+// Returns a sorted vector of VarInfo associated with vars
+std::vector<VarInfo> GetAllVarsInfo(const VariableVector<Real> &vars,
+                                    const IndexShape &cellbounds);
+
 } // namespace OutputUtils
 } // namespace parthenon
 

src/outputs/outputs.cpp

@@ -259,6 +259,17 @@ Outputs::Outputs(Mesh *pm, ParameterInput *pin, SimTime *tm) {
         pnew_type = new VTKOutput(op);
       } else if (op.file_type == "ascent") {
         pnew_type = new AscentOutput(op);
+      } else if (op.file_type == "openpmd") {
+#ifdef PARTHENON_ENABLE_OPENPMD
+        pnew_type = new OpenPMDOutput(op);
+#else
+        msg << "### FATAL ERROR in Outputs constructor" << std::endl
+            << "Executable not configured for OpenPMD outputs, but OpenPMD file format "
+            << "is requested in output/restart block '" << op.block_name << "'. "
+            << "You can disable this block without deleting it by setting a dt < 0."
+            << std::endl;
+        PARTHENON_FAIL(msg);
+#endif // ifdef PARTHENON_ENABLE_OPENPMD
       } else if (op.file_type == "histogram") {
 #ifdef ENABLE_HDF5
         pnew_type = new HistogramOutput(op, pin);

src/outputs/outputs.hpp

@@ -32,6 +32,7 @@
 #include "interface/mesh_data.hpp"
 #include "io_wrapper.hpp"
 #include "kokkos_abstraction.hpp"
+#include "openPMD/Iteration.hpp"
 #include "parthenon_arrays.hpp"
 #include "utils/error_checking.hpp"
 
@@ -204,6 +205,18 @@ class AscentOutput : public OutputType {
   ParArray1D<Real> ghost_mask_;
 };
 
+//----------------------------------------------------------------------------------------
+//! \class OpenPMDOutput
+//  \brief derived OutputType class for OpenPMD based output
+
+class OpenPMDOutput : public OutputType {
+ public:
+  explicit OpenPMDOutput(const OutputParameters &oparams) : OutputType(oparams) {}
+  void WriteOutputFile(Mesh *pm, ParameterInput *pin, SimTime *tm,
+                       const SignalHandler::OutputSignal signal) override;
+
+};
+
 #ifdef ENABLE_HDF5
 //----------------------------------------------------------------------------------------
 //! \class PHDF5Output

src/outputs/parthenon_hdf5.cpp

@@ -357,22 +357,20 @@ void PHDF5Output::WriteOutputFileImpl(Mesh *pm, ParameterInput *pin, SimTime *tm
       const auto &pmb = pm->block_list[b_idx];
       bool is_allocated = false;
 
-      // for each variable that this local meshblock actually has
-      const auto vars = get_vars(pmb);
-      for (auto &v : vars) {
-        // For reference, if we update the logic here, there's also
-        // a similar block in parthenon_manager.cpp
-        if (v->IsAllocated() && (var_name == v->label())) {
-          auto v_h = v->data.GetHostMirrorAndCopy();
-          OutputUtils::PackOrUnpackVar(
-              vinfo, output_params.include_ghost_zones, index,
-              [&](auto index, int topo, int t, int u, int v, int k, int j, int i) {
-                tmpData[index] = static_cast<OutT>(v_h(topo, t, u, v, k, j, i));
-              });
-
-          is_allocated = true;
-          break;
-        }
+      // TODO(reviewers) Why was the loop originally there? Does the direct Get causes
+      // issue?
+      auto v = pmb->meshblock_data.Get()->GetVarPtr(var_name);
+      // For reference, if we update the logic here, there's also
+      // a similar block in parthenon_manager.cpp
+      if (v->IsAllocated() && (var_name == v->label())) {
+        auto v_h = v->data.GetHostMirrorAndCopy();
+        OutputUtils::PackOrUnpackVar(
+            vinfo, output_params.include_ghost_zones, index,
+            [&](auto index, int topo, int t, int u, int v, int k, int j, int i) {
+              tmpData[index] = static_cast<OutT>(v_h(topo, t, u, v, k, j, i));
+            });
+
+        is_allocated = true;
       }
 
       if (vinfo.is_sparse) {