WIP: Add OpenPMD support

.github/workflows/ci-macos.yml

@@ -28,10 +28,14 @@ jobs:
         cache-dependency-path: '**/requirements.txt'
     - run: pip install -r requirements.txt
 
-    - name: Install dependencies
+    - name: Install dependencies (Homebrew)
       run: |
         brew install openmpi hdf5-mpi adios2 || true
 
+    - name: Install OpenPMD
+      run: |
+        openPMD_USE_MPI=ON python3 -m pip install openpmd-api --no-binary openpmd-api
+
     - name: Configure
       run: cmake -B build -DCMAKE_BUILD_TYPE=Release
 

CMakeLists.txt

@@ -36,6 +36,7 @@ option(PARTHENON_DISABLE_MPI "MPI is enabled by default if found, set this to Tr
 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 \
@@ -200,6 +201,26 @@ 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"
+    # we need newer than the latest 0.15.2 release to support writing attriutes from a subset of ranks
+    GIT_TAG        "1c7d7ff") # develop as of 2024-09-02
+  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

@@ -191,6 +191,7 @@ add_library(parthenon
   outputs/history.cpp
   outputs/io_wrapper.cpp
   outputs/io_wrapper.hpp
+  outputs/output_attr.hpp
   outputs/output_utils.cpp
   outputs/output_utils.hpp
   outputs/outputs.cpp
@@ -203,10 +204,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
@@ -326,6 +331,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_attr.hpp

@@ -0,0 +1,47 @@
+//========================================================================================
+// Parthenon performance portable AMR framework
+// Copyright(C) 2023-2024 The Parthenon collaboration
+// Licensed under the 3-clause BSD License, see LICENSE file for details
+//========================================================================================
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
+//
+// This program was produced under U.S. Government contract 89233218CNA000001 for Los
+// Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
+// for the U.S. Department of Energy/National Nuclear Security Administration. All rights
+// in the program are reserved by Triad National Security, LLC, and the U.S. Department
+// of Energy/National Nuclear Security Administration. The Government is granted for
+// itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+// license in this material to reproduce, prepare derivative works, distribute copies to
+// the public, perform publicly and display publicly, and to permit others to do so.
+//========================================================================================
+
+#ifndef OUTPUTS_OUTPUT_ATTR_HPP_
+#define OUTPUTS_OUTPUT_ATTR_HPP_
+
+#include <vector>
+
+// JMM: This could probably be done with template magic but I think
+// using a macro is honestly the simplest and cleanest solution here.
+// Template solution would be to define a variatic class to conain the
+// list of types and then a hierarchy of structs/functions to turn
+// that into function calls. Preprocessor seems easier, given we're
+// not manipulating this list in any way.
+// The following types are the ones we allow to be stored as attributes in outputs
+// (specifically within Params).
+#define PARTHENON_ATTR_VALID_VEC_TYPES(T)                                                \
+  T, std::vector<T>, ParArray1D<T>, ParArray2D<T>, ParArray3D<T>, HostArray1D<T>,        \
+      HostArray2D<T>, HostArray3D<T>, Kokkos::View<T *>, Kokkos::View<T **>,             \
+      ParArrayND<T>, ParArrayHost<T>
+// JMM: This is the list of template specializations we
+// "pre-instantiate" We only pre-instantiate device memory, not host
+// memory. The reason is that when building with the Kokkos serial
+// backend, DevMemSpace and HostMemSpace are the same and so this
+// resolves to the same type in the macro, which causes problems.
+#define PARTHENON_ATTR_FOREACH_VECTOR_TYPE(T)                                            \
+  PARTHENON_ATTR_APPLY(T);                                                               \
+  PARTHENON_ATTR_APPLY(Kokkos::View<T *, LayoutWrapper, DevMemSpace>);                   \
+  PARTHENON_ATTR_APPLY(Kokkos::View<T **, LayoutWrapper, DevMemSpace>);                  \
+  PARTHENON_ATTR_APPLY(Kokkos::View<T ***, LayoutWrapper, DevMemSpace>);                 \
+  PARTHENON_ATTR_APPLY(device_view_t<T>)
+
+#endif // OUTPUTS_OUTPUT_ATTR_HPP_

src/outputs/output_utils.cpp

@@ -18,6 +18,7 @@
 #include <algorithm>
 #include <functional>
 #include <map>
+#include <memory>
 #include <set>
 #include <string>
 #include <type_traits>
@@ -33,6 +34,7 @@
 #include "mesh/meshblock.hpp"
 #include "outputs/output_utils.hpp"
 #include "parameter_input.hpp"
+#include "utils/mpi_types.hpp"
 
 namespace parthenon {
 namespace OutputUtils {
@@ -253,6 +255,47 @@ std::vector<int> ComputeDerefinementCount(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<std::size_t>
+FlattendedLocalToGlobal(Mesh *pm, const std::vector<std::size_t> &data_local);
+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.
@@ -329,6 +372,35 @@ 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;
+}
+
 void CheckParameterInputConsistent(ParameterInput *pin) {
 #ifdef MPI_PARALLEL
   CheckMPISizeT();

src/outputs/output_utils.hpp

@@ -213,8 +213,8 @@ struct SwarmInfo {
   std::size_t count_on_rank = 0;                        // per-meshblock
   std::size_t global_offset;                            // global
   std::size_t global_count;                             // global
-  std::vector<std::size_t> counts;                      // per-meshblock
-  std::vector<std::size_t> offsets;                     // global
+  std::vector<std::size_t> counts;                      // on local meshblocks
+  std::vector<std::size_t> offsets; // global offset for local meshblocks
   // std::vector<ParArray1D<bool>> masks; // used for reading swarms without defrag
   std::vector<std::size_t> max_indices;   // JMM: If we defrag, unneeded?
   void AddOffsets(const SP_Swarm &swarm); // sets above metadata
@@ -236,7 +236,7 @@ struct SwarmInfo {
   // Copies swarmvar to host in prep for output
   template <typename T>
   std::vector<T> FillHostBuffer(const std::string vname,
-                                ParticleVariableVector<T> &swmvarvec) {
+                                const ParticleVariableVector<T> &swmvarvec) const {
     const auto &vinfo = var_info.at(vname);
     std::vector<T> host_data(count_on_rank * vinfo.nvar);
     std::size_t ivec = 0;
@@ -245,6 +245,7 @@ struct SwarmInfo {
         for (int n4 = 0; n4 < vinfo.GetN(4); ++n4) {
           for (int n3 = 0; n3 < vinfo.GetN(3); ++n3) {
             for (int n2 = 0; n2 < vinfo.GetN(2); ++n2) {
+              // TODO(pgrete) understand what's doing on with the blocks here...
               std::size_t block_idx = 0;
               for (auto &swmvar : swmvarvec) {
                 // Copied extra times. JMM: If we defrag, unneeded?
@@ -344,13 +345,29 @@ std::vector<int64_t> ComputeLocs(Mesh *pm);
 std::vector<int> ComputeIDsAndFlags(Mesh *pm);
 std::vector<int> ComputeDerefinementCount(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);
+
 void CheckParameterInputConsistent(ParameterInput *pin);
 } // namespace OutputUtils
 } // namespace parthenon

src/outputs/outputs.cpp

@@ -214,8 +214,13 @@ Outputs::Outputs(Mesh *pm, ParameterInput *pin, SimTime *tm) {
       // set output variable and optional data format string used in formatted writes
       if ((op.file_type != "hst") && (op.file_type != "rst") &&
           (op.file_type != "ascent") && (op.file_type != "histogram")) {
-        op.variables = pin->GetOrAddVector<std::string>(pib->block_name, "variables",
-                                                        std::vector<std::string>());
+        // differentiating here whether a block exists or not to not add an empty
+        // parameter to the input file (which might interfere with restarts)
+        if (pin->DoesParameterExist(pib->block_name, "variables")) {
+          op.variables = pin->GetVector<std::string>(pib->block_name, "variables");
+        } else {
+          op.variables = std::vector<std::string>();
+        }
         // JMM: If the requested var isn't present for a given swarm,
         // it is simply not output.
         op.swarms.clear(); // Not sure this is needed
@@ -263,6 +268,20 @@ 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
+        const auto backend_config =
+            pin->GetOrAddString(op.block_name, "backend_config", "default");
+
+        pnew_type = new OpenPMDOutput(op, backend_config);
+#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

@@ -23,6 +23,7 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "Kokkos_ScatterView.hpp"
@@ -171,6 +172,22 @@ 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, std::string backend_config)
+      : OutputType(oparams), backend_config_(std::move(backend_config)) {}
+  void WriteOutputFile(Mesh *pm, ParameterInput *pin, SimTime *tm,
+                       const SignalHandler::OutputSignal signal) override;
+
+ private:
+  //  path to file containing config passed to backend
+  std::string backend_config_;
+};
+
 #ifdef ENABLE_HDF5
 //----------------------------------------------------------------------------------------
 //! \class PHDF5Output