add support for writing openPMD/ADIOS2 plotfiles (#466)

### Description
This adds support for writing ADIOS2 format plotfiles using openPMD-api
(https://github.com/openPMD/openPMD-api). This is the file format used
by WarpX for its large-scale runs, and it is the fastest option for
large-scale I/O (~Terabytes or larger).

This adds support for cell-centered variables only. (However, it does
include cell-centered averages of the face-centered variables, and all
of the other derived variables.)

openPMD currently does not support ghost cells, so we exclude them when
writing the plotfile in openPMD/ADIOS2 format.

There is no support for checkpoints or restarts with this format (nor is
there in WarpX, or any other code that I'm aware of).

### Related issues
N/A

### Checklist
_Before this pull request can be reviewed, all of these tasks should be
completed. Denote completed tasks with an `x` inside the square brackets
`[ ]` in the Markdown source below:_
- [x] I have added a description (see above).
- [x] I have added a link to any related issues see (see above).
- [x] I have read the [Contributing
Guide](https://github.com/quokka-astro/quokka/blob/development/CONTRIBUTING.md).
- [x] I have added tests for any new physics that this PR adds to the
code.
- [x] I have tested this PR on my local computer and all tests pass.
- [x] I have manually triggered the GPU tests with the magic comment
`/azp run`.
- [x] I have requested a reviewer for this PR.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

.github/workflows/clang-tidy.yml

@@ -23,7 +23,7 @@ jobs:
         config_file: src/.clang-tidy
         build_dir: build
         apt_packages: libopenmpi-dev,libhdf5-mpi-dev,python3-dev,python3-numpy,python3-matplotlib
-        cmake_command: cmake . -B build -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DQUOKKA_PYTHON=ON
+        cmake_command: cmake . -B build -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DQUOKKA_PYTHON=ON -DQUOKKA_OPENPMD=ON -DopenPMD_USE_ADIOS2=OFF
         split_workflow: true
 
     # Uploads an artefact containing clang_fixes.json

.gitmodules

@@ -14,3 +14,6 @@
 [submodule "extern/yaml-cpp"]
 	path = extern/yaml-cpp
 	url = https://github.com/jbeder/yaml-cpp.git
+[submodule "extern/openPMD-api"]
+	path = extern/openPMD-api
+	url = https://github.com/openPMD/openPMD-api.git

CMakeLists.txt

@@ -31,6 +31,7 @@ option(QUOKKA_PYTHON "Compile with Python support (on/off)" ON)
 option(DISABLE_FMAD "Disable fused multiply-add instructions on GPU (on/off)" ON)
 option(ENABLE_ASAN "Enable AddressSanitizer and UndefinedBehaviorSanitizer" OFF)
 option(WARNINGS_AS_ERRORS "Treat compiler warnings as errors" OFF)
+option(QUOKKA_OPENPMD "Enable OpenPMD output (on/off)" OFF)
 
 if(AMReX_GPU_BACKEND MATCHES "CUDA")
   enable_language(CUDA)

README.md

@@ -34,6 +34,7 @@ Quokka also features advanced Adaptive Quokka Refinement:tm: technology:
 * ROCm 5.2.0+ (optional, for AMD GPUs)
 * Ninja (optional, for faster builds)
 * Python 3.7+ (optional)
+* ADIOS2 2.9+ with GPU-aware support (optional, for writing terabyte-sized or larger outputs)
 
 ## Quickstart
 

src/CMakeLists.txt

@@ -11,6 +11,25 @@ if(Python_FOUND)
     link_libraries(${Python_LIBRARIES})
 endif()
 
+if(QUOKKA_OPENPMD)
+  message(STATUS "Building Quokka with OpenPMD support")
+  add_compile_definitions(QUOKKA_USE_OPENPMD)
+  set(openPMD_USE_ADIOS2 ON CACHE BOOL "") # ADIOS2 is required
+  set(openPMD_USE_HDF5 OFF CACHE BOOL "")
+  set(openPMD_USE_PYTHON OFF CACHE BOOL "")
+  set(openPMD_BUILD_TESTING OFF CACHE BOOL "")
+  set(openPMD_BUILD_EXAMPLES OFF CACHE BOOL "")
+  set(openPMD_BUILD_CLI_TOOLS OFF CACHE BOOL "")
+  set(openPMD_INSTALL OFF CACHE BOOL "")
+  add_subdirectory(${QuokkaCode_SOURCE_DIR}/extern/openPMD-api ${QuokkaCode_BINARY_DIR}/openPMD-api)
+  include_directories(${OpenPMD_INCLUDE_DIRS_RET})
+  link_libraries(openPMD::openPMD)
+  set(openPMDSources "${CMAKE_CURRENT_SOURCE_DIR}/openPMD.cpp")
+  message(STATUS "WARNING: OpenPMD plotfiles are ENABLED. Face-centered variables will only be available as cell-centered averages in plotfiles!")
+else()
+  set(openPMDSources "")
+endif()
+
 # HDF5
 find_package(HDF5 REQUIRED)
 
@@ -97,7 +116,7 @@ include(CTest)
 
 #create an object library for files that should be compiled with all test problems
 set (QuokkaObjSources "${CMAKE_CURRENT_SOURCE_DIR}/main.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/CloudyCooling.cpp" 
-                      "${CMAKE_CURRENT_SOURCE_DIR}/GrackleDataReader.cpp"
+                      "${CMAKE_CURRENT_SOURCE_DIR}/GrackleDataReader.cpp" "${openPMDSources}"
                       "${gamma_law_sources}")
 #we don't use it anymore because it gives issues on Cray systems
 #add_library(QuokkaObj OBJECT ${QuokkaObjSources} ${gamma_law_sources})

src/PrimordialChem/CMakeLists.txt

@@ -5,7 +5,7 @@ setup_target_for_microphysics_compilation(${microphysics_network_name} "${CMAKE_
 #this is critical to ensure the correct Microphysics files are linked to primordial chem target
 include_directories(BEFORE ${primordial_chem_dirs} "${CMAKE_CURRENT_BINARY_DIR}/" "includes/extern_parameters.H" "includes/network_properties.H")
 
-add_executable(test_primordial_chem test_primordial_chem.cpp ../main.cpp ../GrackleDataReader.cpp ../CloudyCooling.cpp ../Chemistry.cpp ${primordial_chem_sources})
+add_executable(test_primordial_chem test_primordial_chem.cpp ../main.cpp "${openPMDSources}" ../GrackleDataReader.cpp ../CloudyCooling.cpp ../Chemistry.cpp ${primordial_chem_sources})
 target_compile_definitions(test_primordial_chem PUBLIC PRIMORDIAL_CHEM) #this will add #define PRIMORDIAL_CHEM
 
 #de-link QuokkaObj from this target to avoid multiple definitions of same variables

src/openPMD.cpp

@@ -0,0 +1,144 @@
+//==============================================================================
+// TwoMomentRad - a radiation transport library for patch-based AMR codes
+// Copyright 2020 Benjamin Wibking.
+// Released under the MIT license. See LICENSE file included in the GitHub repo.
+//==============================================================================
+//! \file openPMD.cpp
+///  \brief openPMD I/O for snapshots
+
+#include <cstdint>
+#include <string>
+
+// AMReX headers
+#include "AMReX.H"
+#include "AMReX_Geometry.H"
+#include "AMReX_IntVect.H"
+#include "AMReX_MultiFab.H"
+
+// openPMD headers
+#include "openPMD/openPMD.hpp"
+
+namespace quokka::OpenPMDOutput
+{
+
+namespace detail
+{
+/** \brief
+ * Convert an IntVect to a std::vector<std::uint64_t>
+ * and reverse the order of the elements
+ * (used for compatibility with the openPMD API)
+ */
+auto getReversedVec(const amrex::IntVect &v) -> std::vector<std::uint64_t>
+{
+	// Convert the IntVect v to and std::vector u
+	std::vector<std::uint64_t> u = {AMREX_D_DECL(static_cast<std::uint64_t>(v[0]), static_cast<std::uint64_t>(v[1]), static_cast<std::uint64_t>(v[2]))};
+	// Reverse the order of elements, if v corresponds to the indices of a Fortran-order array (like an AMReX FArrayBox)
+	// but u is intended to be used with a C-order API (like openPMD)
+	std::reverse(u.begin(), u.end());
+	return u;
+}
+
+/** \brief
+ * Convert Real* pointer to a std::vector<double>,
+ * and reverse the order of the elements
+ * (used for compatibility with the openPMD API)
+ */
+auto getReversedVec(const amrex::Real *v) -> std::vector<double>
+{
+	// Convert Real* v to and std::vector u
+	std::vector<double> u = {AMREX_D_DECL(static_cast<double>(v[0]), static_cast<double>(v[1]), static_cast<double>(v[2]))}; // NOLINT
+	// Reverse the order of elements, if v corresponds to the indices of a Fortran-order array (like an AMReX FArrayBox)
+	// but u is intended to be used with a C-order API (like openPMD)
+	std::reverse(u.begin(), u.end());
+	return u;
+}
+
+void SetupMeshComponent(openPMD::Mesh &mesh, amrex::Geometry &full_geom)
+{
+	amrex::Box const &global_box = full_geom.Domain();
+	auto global_size = getReversedVec(global_box.size());
+	std::vector<double> const grid_spacing = getReversedVec(full_geom.CellSize());
+	std::vector<double> const global_offset = getReversedVec(full_geom.ProbLo());
+
+	// Prepare the type of dataset that will be written
+	mesh.setDataOrder(openPMD::Mesh::DataOrder::C);
+	mesh.setGridSpacing(grid_spacing);
+	mesh.setGridGlobalOffset(global_offset);
+	mesh.setAttribute("fieldSmoothing", "none");
+
+	auto mesh_comp = mesh[openPMD::MeshRecordComponent::SCALAR];
+	auto const dataset = openPMD::Dataset(openPMD::determineDatatype<amrex::Real>(), global_size);
+	mesh_comp.resetDataset(dataset);
+	std::vector<amrex::Real> const relativePosition{0.5, 0.5, 0.5}; // cell-centered only (for now)
+	mesh_comp.setPosition(relativePosition);
+}
+
+auto GetMeshComponentName(int meshLevel, std::string const &field_name) -> std::string
+{
+	std::string new_field_name = field_name;
+	if (meshLevel > 0) {
+		new_field_name += std::string("_lvl").append(std::to_string(meshLevel));
+	}
+	return new_field_name;
+}
+} // namespace detail
+//----------------------------------------------------------------------------------------
+//! \fn void OpenPMDOutput:::WriteOutputFile(Mesh *pm)
+//  \brief  Write cell-centered MultiFab using openPMD
+void WriteFile(const std::vector<std::string> &varnames, int const output_levels, amrex::Vector<const amrex::MultiFab *> &mf,
+	       amrex::Vector<amrex::Geometry> &geom, const std::string &output_basename, amrex::Real const time, int const file_number)
+{
+	// open file
+	std::string const filename = output_basename + ".bp";
+	auto series = openPMD::Series(filename, openPMD::Access::CREATE, amrex::ParallelDescriptor::Communicator());
+	series.setSoftware("Quokka", "1.0");
+
+	auto series_iteration = series.iterations[file_number];
+	series_iteration.open();
+	series_iteration.setTime(time);
+	auto meshes = series_iteration.meshes;
+
+	// loop over levels up to output_levels
+	for (int lev = 0; lev < output_levels; lev++) {
+		amrex::Geometry full_geom = geom[lev];
+		amrex::Box const &global_box = full_geom.Domain();
+		int const ncomp = mf[lev]->nComp();
+
+		for (int icomp = 0; icomp < ncomp; icomp++) {
+			const std::string field_name = detail::GetMeshComponentName(lev, varnames[icomp]);
+			if (!meshes.contains(field_name)) {
+				auto mesh = meshes[field_name];
+				detail::SetupMeshComponent(mesh, full_geom);
+			}
+		}
+
+		// pass data pointers for each box to ADIOS
+		for (int icomp = 0; icomp < ncomp; icomp++) {
+			std::string const field_name = detail::GetMeshComponentName(lev, varnames[icomp]);
+			openPMD::MeshRecordComponent mesh = series_iteration.meshes[field_name][openPMD::MeshRecordComponent::SCALAR];
+
+			// Loop through the multifab, and store each box as a chunk in the openPMD file.
+			for (amrex::MFIter mfi(*mf[lev]); mfi.isValid(); ++mfi) {
+				amrex::FArrayBox const &fab = (*mf[lev])[mfi];
+				amrex::Box const &local_box = fab.box();
+
+				// Determine the offset and size of this chunk
+				amrex::IntVect const box_offset = local_box.smallEnd() - global_box.smallEnd();
+				auto chunk_offset = detail::getReversedVec(box_offset); // this overflows if ghost zones are used (!)
+				auto chunk_size = detail::getReversedVec(local_box.size());
+
+				// pass device pointer directly to ADIOS
+				amrex::Real const *local_data = fab.dataPtr(icomp);
+				mesh.storeChunkRaw(local_data, chunk_offset, chunk_size);
+			}
+		}
+
+		// flush this level to disk
+		series.flush();
+	}
+
+	// close file
+	series.close();
+}
+
+} // namespace quokka::OpenPMDOutput

src/openPMD.hpp

@@ -0,0 +1,40 @@
+#ifndef OPENPMD_HPP_ // NOLINT
+#define OPENPMD_HPP_
+//==============================================================================
+// TwoMomentRad - a radiation transport library for patch-based AMR codes
+// Copyright 2020 Benjamin Wibking.
+// Released under the MIT license. See LICENSE file included in the GitHub repo.
+//==============================================================================
+//! \file openPMD.hpp
+///  \brief openPMD I/O for snapshots
+
+#include <string>
+
+// AMReX headers
+#include "AMReX_Geometry.H"
+#include "AMReX_IntVect.H"
+#include "AMReX_MultiFab.H"
+#include <AMReX.H>
+
+// openPMD headers
+#include "openPMD/openPMD.hpp"
+
+namespace quokka::OpenPMDOutput
+{
+
+namespace detail
+{
+
+auto getReversedVec(const amrex::IntVect &v) -> std::vector<std::uint64_t>;
+auto getReversedVec(const amrex::Real *v) -> std::vector<double>;
+void SetupMeshComponent(openPMD::Mesh &mesh, int /*meshLevel*/, const std::string &comp_name, amrex::Geometry &full_geom);
+auto GetMeshComponentName(int meshLevel, std::string const &field_name) -> std::string;
+
+} // namespace detail
+
+void WriteFile(const std::vector<std::string> &varnames, int output_levels, amrex::Vector<const amrex::MultiFab *> &mf, amrex::Vector<amrex::Geometry> &geom,
+	       const std::string &output_basename, amrex::Real time, int file_number);
+
+} // namespace quokka::OpenPMDOutput
+
+#endif // OPENPMD_HPP_