diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index e4c996ff1..5efa5ab68 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -134,6 +134,7 @@ set (QuokkaSourcesNoEOS "${CMAKE_CURRENT_SOURCE_DIR}/main.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/io/DiagFilter.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/io/DiagFramePlane.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/io/DiagPDF.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/io/projection.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/cooling/GrackleLikeCooling.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/cooling/GrackleDataReader.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/cooling/TabulatedCooling.cpp" diff --git a/src/QuokkaSimulation.hpp b/src/QuokkaSimulation.hpp index f6fff3738..9ffee1bb0 100644 --- a/src/QuokkaSimulation.hpp +++ b/src/QuokkaSimulation.hpp @@ -190,7 +190,7 @@ template class QuokkaSimulation : public AMRSimulation std::unordered_map> override; + [[nodiscard]] auto ComputeProjections(const amrex::Direction dir) const -> std::unordered_map> override; // compute statistics auto ComputeStatistics() -> std::map override; @@ -600,7 +600,7 @@ template void QuokkaSimulation::ComputeDerivedVa } template -auto QuokkaSimulation::ComputeProjections(int /*dir*/) const -> std::unordered_map> +auto QuokkaSimulation::ComputeProjections(const amrex::Direction /*dir*/) const -> std::unordered_map> { // compute projections and return as unordered_map -- user should implement return std::unordered_map>{}; diff --git a/src/io/projection.cpp b/src/io/projection.cpp new file mode 100644 index 000000000..1f35cfd02 --- /dev/null +++ b/src/io/projection.cpp @@ -0,0 +1,581 @@ +//============================================================================== +// 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 projection.cpp +/// \brief AMReX I/O for 2D projections + +#include "AMReX_Array.H" +#include "AMReX_BLassert.H" +#include "AMReX_DistributionMapping.H" +#include "AMReX_FPC.H" +#include "AMReX_Geometry.H" +#include "AMReX_GpuDevice.H" +#include "AMReX_IntVect.H" +#include "AMReX_Orientation.H" +#include "AMReX_PlotFileUtil.H" +#include "AMReX_REAL.H" +#include "AMReX_SPACE.H" + +#include "projection.hpp" + +namespace quokka::diagnostics +{ + +namespace detail +{ +auto direction_to_string(const amrex::Direction dir) -> std::string +{ + if (dir == amrex::Direction::x) { + return std::string("x"); + } +#if AMREX_SPACEDIM >= 2 + if (dir == amrex::Direction::y) { + return std::string("y"); + } +#endif +#if AMREX_SPACEDIM == 3 + if (dir == amrex::Direction::z) { + return std::string("z"); + } +#endif + + amrex::Error("invalid direction in quokka::diagnostics::direction_to_string!"); + return std::string(""); +} + +void printLowerDimIntVect(std::ostream &a_File, const amrex::IntVect &a_IntVect, int skipDim) +{ + int doneDim = 0; + a_File << '('; + for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) { + if (idim != skipDim) { + a_File << a_IntVect[idim]; + doneDim++; + if (doneDim < AMREX_SPACEDIM - 1) { + a_File << ","; + } + } + } + a_File << ')'; +} + +void printLowerDimBox(std::ostream &a_File, const amrex::Box &a_box, int skipDim) +{ + a_File << '('; + printLowerDimIntVect(a_File, a_box.smallEnd(), skipDim); + a_File << ' '; + printLowerDimIntVect(a_File, a_box.bigEnd(), skipDim); + a_File << ' '; + printLowerDimIntVect(a_File, a_box.type(), skipDim); + a_File << ')'; +} + +void Write2DMultiLevelPlotfile(const std::string &a_pltfile, int a_nlevels, const amrex::Vector &a_slice, + const amrex::Vector &a_varnames, const amrex::Vector &a_geoms, const amrex::Real &a_time, + const amrex::Vector &a_steps, const amrex::Vector &a_rref) +{ + // Write a 2D AMReX Plotfile with the contents of the vector of MultiFabs `a_slice'. + + const std::string levelPrefix = "Level_"; + const std::string mfPrefix = "Cell"; + const std::string versionName = "HyperCLaw-V1.1"; + + bool const callBarrier(false); + amrex::PreBuildDirectorHierarchy(a_pltfile, levelPrefix, a_nlevels, callBarrier); + amrex::ParallelDescriptor::Barrier(); + + if (amrex::ParallelDescriptor::MyProc() == amrex::ParallelDescriptor::NProcs() - 1) { + + // Write 2D pltfile header + amrex::Vector boxArrays(a_nlevels); + for (int level(0); level < boxArrays.size(); ++level) { + boxArrays[level] = a_slice[level]->boxArray(); + } + + amrex::VisMF::IO_Buffer io_buffer(amrex::VisMF::IO_Buffer_Size); + std::string const HeaderFileName(a_pltfile + "/Header"); + std::ofstream HeaderFile; + HeaderFile.rdbuf()->pubsetbuf(io_buffer.dataPtr(), io_buffer.size()); + HeaderFile.open(HeaderFileName.c_str(), std::ofstream::out | std::ofstream::trunc | std::ofstream::binary); + if (!HeaderFile.good()) { + amrex::FileOpenFailed(HeaderFileName); + } + Write2DPlotfileHeader(HeaderFile, a_nlevels, boxArrays, a_varnames, a_geoms, a_time, a_steps, a_rref, versionName, levelPrefix, mfPrefix); + HeaderFile.flush(); + HeaderFile.close(); + } + + // Write a 2D version of the MF at each level + for (int level = 0; level < a_nlevels; ++level) { + VisMF2D(*a_slice[level], amrex::MultiFabFileFullPrefix(level, a_pltfile, levelPrefix, mfPrefix)); + } +} + +void Write2DPlotfileHeader(std::ostream &HeaderFile, int nlevels, const amrex::Vector &bArray, const amrex::Vector &varnames, + const amrex::Vector &geom, const amrex::Real &time, const amrex::Vector &level_steps, + const amrex::Vector &ref_ratio, const std::string &versionName, const std::string &levelPrefix, + const std::string &mfPrefix) +{ + int const finest_level(nlevels - 1); + HeaderFile.precision(17); + + int const lowerSpaceDim = AMREX_SPACEDIM - 1; + + HeaderFile << versionName << '\n'; + HeaderFile << varnames.size() << '\n'; + for (const auto &varname : varnames) { + HeaderFile << varname << "\n"; + } + HeaderFile << lowerSpaceDim << '\n'; + HeaderFile << time << '\n'; + HeaderFile << finest_level << '\n'; + for (int idim = 0; idim < lowerSpaceDim; ++idim) { + HeaderFile << geom[0].ProbLo(idim) << ' '; + } + HeaderFile << '\n'; + for (int idim = 0; idim < lowerSpaceDim; ++idim) { + HeaderFile << geom[0].ProbHi(idim) << ' '; + } + HeaderFile << '\n'; + for (int i = 0; i < finest_level; ++i) { + HeaderFile << ref_ratio[i][0] << ' '; + } + HeaderFile << '\n'; + for (int i = 0; i <= finest_level; ++i) { + printLowerDimBox(HeaderFile, geom[i].Domain(), 2); + HeaderFile << ' '; + } + HeaderFile << '\n'; + for (int i = 0; i <= finest_level; ++i) { + HeaderFile << level_steps[i] << ' '; + } + HeaderFile << '\n'; + for (int i = 0; i <= finest_level; ++i) { + for (int idim = 0; idim < lowerSpaceDim; ++idim) { + HeaderFile << geom[i].CellSizeArray()[idim] << ' '; + } + HeaderFile << '\n'; + } + HeaderFile << static_cast(geom[0].Coord()) << '\n'; + HeaderFile << "0\n"; + + for (int level = 0; level <= finest_level; ++level) { + HeaderFile << level << ' ' << bArray[level].size() << ' ' << time << '\n'; + HeaderFile << level_steps[level] << '\n'; + + const amrex::IntVect &domain_lo = geom[level].Domain().smallEnd(); + for (int i = 0; i < bArray[level].size(); ++i) { + // Need to shift because the RealBox ctor we call takes the + // physical location of index (0,0,0). This does not affect + // the usual cases where the domain index starts with 0. + const amrex::Box &b = amrex::shift(bArray[level][i], -domain_lo); + amrex::RealBox const loc = amrex::RealBox(b, geom[level].CellSize(), geom[level].ProbLo()); + for (int idim = 0; idim < lowerSpaceDim; ++idim) { + HeaderFile << loc.lo(idim) << ' ' << loc.hi(idim) << '\n'; + } + } + HeaderFile << amrex::MultiFabHeaderPath(level, levelPrefix, mfPrefix) << '\n'; + } +} + +void VisMF2D(const amrex::MultiFab &a_mf, const std::string &a_mf_name) +{ + auto whichRD = amrex::FArrayBox::getDataDescriptor(); + bool const doConvert(*whichRD != amrex::FPC::NativeRealDescriptor()); + + amrex::Long bytesWritten(0); + + std::string const filePrefix(a_mf_name + "_D_"); + + bool const calcMinMax = false; + amrex::VisMF::Header::Version const currentVersion = amrex::VisMF::Header::Version_v1; + amrex::VisMF::How const how = amrex::VisMF::How::NFiles; + amrex::VisMF::Header hdr(a_mf, how, currentVersion, calcMinMax); + + int const nOutFiles = std::max(1, std::min(amrex::ParallelDescriptor::NProcs(), 256)); + bool const groupSets = false; + bool const setBuf = true; + + amrex::NFilesIter nfi(nOutFiles, filePrefix, groupSets, setBuf); + + // Check if mf has sparse data + bool useSparseFPP = false; + const amrex::Vector &pmap = a_mf.DistributionMap().ProcessorMap(); + std::set procsWithData; + amrex::Vector procsWithDataVector; + for (const int i : pmap) { + procsWithData.insert(i); + } + if (static_cast(procsWithData.size()) < nOutFiles) { + useSparseFPP = true; + for (const int it : procsWithData) { + procsWithDataVector.push_back(it); + } + } + + if (useSparseFPP) { + nfi.SetSparseFPP(procsWithDataVector); + } else { + nfi.SetDynamic(); + } + for (; nfi.ReadyToWrite(); ++nfi) { + int const whichRDBytes(whichRD->numBytes()); + int nFABs(0); + amrex::Long writeDataItems(0); + amrex::Long writeDataSize(0); + for (amrex::MFIter mfi(a_mf); mfi.isValid(); ++mfi) { + const amrex::FArrayBox &fab = a_mf[mfi]; + std::stringstream hss; + write_2D_header(hss, fab, fab.nComp()); + bytesWritten += static_cast(hss.tellp()); + bytesWritten += fab.box().numPts() * a_mf.nComp() * whichRDBytes; + ++nFABs; + } + std::unique_ptr> allFabData{}; + bool canCombineFABs = false; + if ((nFABs > 1 || doConvert) && amrex::VisMF::GetUseSingleWrite()) { + allFabData = std::make_unique>(bytesWritten); + } // ---- else { no need to make a copy for one fab } + canCombineFABs = allFabData != nullptr; + + if (canCombineFABs) { + amrex::Long writePosition = 0; + for (amrex::MFIter mfi(a_mf); mfi.isValid(); ++mfi) { + int hLength(0); + const amrex::FArrayBox &fab = a_mf[mfi]; + writeDataItems = fab.box().numPts() * a_mf.nComp(); + writeDataSize = writeDataItems * whichRDBytes; + char *afPtr = allFabData->data() + writePosition; // NOLINT(cppcoreguidelines-pro-bounds-pointer-arithmetic) + std::stringstream hss; + write_2D_header(hss, fab, fab.nComp()); + hLength = static_cast(hss.tellp()); + auto tstr = hss.str(); + std::memcpy(afPtr, tstr.c_str(), hLength); // ---- the fab header + amrex::Real const *fabdata = fab.dataPtr(); +#ifdef AMREX_USE_GPU + std::unique_ptr hostfab; + if (fab.arena()->isManaged() || fab.arena()->isDevice()) { + hostfab = std::make_unique(fab.box(), fab.nComp(), amrex::The_Pinned_Arena()); + amrex::Gpu::dtoh_memcpy_async(hostfab->dataPtr(), fab.dataPtr(), fab.size() * sizeof(amrex::Real)); + amrex::Gpu::streamSynchronize(); + fabdata = hostfab->dataPtr(); + } +#endif + memcpy(afPtr + hLength, fabdata, writeDataSize); // NOLINT(cppcoreguidelines-pro-bounds-pointer-arithmetic) + writePosition += hLength + writeDataSize; + } + nfi.Stream().write(allFabData->data(), bytesWritten); + nfi.Stream().flush(); + // delete[] allFabData; + } else { + for (amrex::MFIter mfi(a_mf); mfi.isValid(); ++mfi) { + int hLength = 0; + const amrex::FArrayBox &fab = a_mf[mfi]; + writeDataItems = fab.box().numPts() * a_mf.nComp(); + writeDataSize = writeDataItems * whichRDBytes; + std::stringstream hss; + write_2D_header(hss, fab, fab.nComp()); + hLength = static_cast(hss.tellp()); + auto tstr = hss.str(); + nfi.Stream().write(tstr.c_str(), hLength); // ---- the fab header + nfi.Stream().flush(); + amrex::Real const *fabdata = fab.dataPtr(); +#ifdef AMREX_USE_GPU + std::unique_ptr hostfab; + if (fab.arena()->isManaged() || fab.arena()->isDevice()) { + hostfab = std::make_unique(fab.box(), fab.nComp(), amrex::The_Pinned_Arena()); + amrex::Gpu::dtoh_memcpy_async(hostfab->dataPtr(), fab.dataPtr(), fab.size() * sizeof(amrex::Real)); + amrex::Gpu::streamSynchronize(); + fabdata = hostfab->dataPtr(); + } +#endif + nfi.Stream().write(reinterpret_cast(fabdata), // NOLINT(cppcoreguidelines-pro-type-reinterpret-cast) + writeDataSize); + nfi.Stream().flush(); + } + } + } + + int coordinatorProc(amrex::ParallelDescriptor::IOProcessorNumber()); + if (nfi.GetDynamic()) { + coordinatorProc = nfi.CoordinatorProc(); + } + hdr.CalculateMinMax(a_mf, coordinatorProc); + + Find2FOffsets(a_mf, filePrefix, hdr, currentVersion, nfi, nOutFiles, amrex::ParallelDescriptor::Communicator()); + + Write2DMFHeader(a_mf_name, hdr, coordinatorProc, amrex::ParallelDescriptor::Communicator()); +} + +void Write2DMFHeader(const std::string &a_mf_name, amrex::VisMF::Header &hdr, int coordinatorProc, MPI_Comm comm) +{ + const int myProc(amrex::ParallelDescriptor::MyProc(comm)); + if (myProc == coordinatorProc) { + + std::string MFHdrFileName(a_mf_name); + + MFHdrFileName += "_H"; + + amrex::VisMF::IO_Buffer io_buffer(amrex::VisMF::IO_Buffer_Size); + + std::ofstream MFHdrFile; + + MFHdrFile.rdbuf()->pubsetbuf(io_buffer.dataPtr(), io_buffer.size()); + MFHdrFile.open(MFHdrFileName.c_str(), std::ios::out | std::ios::trunc); + + if (!MFHdrFile.good()) { + amrex::FileOpenFailed(MFHdrFileName); + } + + MFHdrFile.setf(std::ios::floatfield, std::ios::scientific); + MFHdrFile << hdr.m_vers << '\n'; + MFHdrFile << static_cast(hdr.m_how) << '\n'; + MFHdrFile << hdr.m_ncomp << '\n'; + if (hdr.m_ngrow == hdr.m_ngrow[0]) { + MFHdrFile << hdr.m_ngrow[0] << '\n'; + } else { + MFHdrFile << hdr.m_ngrow << '\n'; + } + + MFHdrFile << "(" << hdr.m_ba.size() << " 0" << '\n'; + for (int i = 0; i < hdr.m_ba.size(); ++i) { + printLowerDimBox(MFHdrFile, hdr.m_ba[i], 2); + MFHdrFile << '\n'; + } + MFHdrFile << ") \n"; + + MFHdrFile << hdr.m_fod << '\n'; + + MFHdrFile << hdr.m_min.size() << "," << hdr.m_min[0].size() << '\n'; + MFHdrFile.precision(16); + for (auto &hdr_i : hdr.m_min) { + for (const double j : hdr_i) { + MFHdrFile << j << ","; + } + MFHdrFile << "\n"; + } + + MFHdrFile << "\n"; + + MFHdrFile << hdr.m_max.size() << "," << hdr.m_max[0].size() << '\n'; + for (auto &hdr_i : hdr.m_max) { + for (const double j : hdr_i) { + MFHdrFile << j << ","; + } + MFHdrFile << "\n"; + } + + MFHdrFile.flush(); + MFHdrFile.close(); + } +} + +void Find2FOffsets(const amrex::FabArray &mf, const std::string &filePrefix, amrex::VisMF::Header &hdr, + amrex::VisMF::Header::Version /*whichVersion*/, amrex::NFilesIter &nfi, int nOutFiles, MPI_Comm comm) +{ + bool const groupSets = false; + + const int myProc(amrex::ParallelDescriptor::MyProc(comm)); + const int nProcs(amrex::ParallelDescriptor::NProcs(comm)); + int coordinatorProc(amrex::ParallelDescriptor::IOProcessorNumber(comm)); + if (nfi.GetDynamic()) { + coordinatorProc = nfi.CoordinatorProc(); + } + + auto whichRD = amrex::FArrayBox::getDataDescriptor(); + int const whichRDBytes(whichRD->numBytes()); + int const nComps(mf.nComp()); + + if (myProc == coordinatorProc) { // ---- calculate offsets + const amrex::BoxArray &mfBA = mf.boxArray(); + const amrex::DistributionMapping &mfDM = mf.DistributionMap(); + amrex::Vector fabHeaderBytes(mfBA.size(), 0); + int const nFiles(amrex::NFilesIter::ActualNFiles(nOutFiles)); + int whichFileNumber(-1); + std::string whichFileName; + amrex::Vector currentOffset(nProcs, 0L); + + // ---- find the length of the fab header instead of asking the file system + for (int i(0); i < mfBA.size(); ++i) { + std::stringstream hss; + amrex::FArrayBox const tempFab(mf.fabbox(i), nComps, false); // ---- no alloc + write_2D_header(hss, tempFab, tempFab.nComp()); + fabHeaderBytes[i] = static_cast(hss.tellp()); + } + + std::map> rankBoxOrder; // ---- [rank, boxarray index array] + for (int i(0); i < mfBA.size(); ++i) { + rankBoxOrder[mfDM[i]].push_back(i); + } + + amrex::Vector fileNumbers; + if (nfi.GetDynamic()) { + fileNumbers = nfi.FileNumbersWritten(); + } else if (nfi.GetSparseFPP()) { // if sparse, write to (file number = rank) + fileNumbers.resize(nProcs); + for (int i(0); i < nProcs; ++i) { + fileNumbers[i] = i; + } + } else { + fileNumbers.resize(nProcs); + for (int i(0); i < nProcs; ++i) { + fileNumbers[i] = amrex::NFilesIter::FileNumber(nFiles, i, groupSets); + } + } + const amrex::Vector> &fileNumbersWriteOrder = nfi.FileNumbersWriteOrder(); + + for (const auto &fn : fileNumbersWriteOrder) { + for (const int rank : fn) { + auto rboIter = rankBoxOrder.find(rank); + + if (rboIter != rankBoxOrder.end()) { + amrex::Vector const &index = rboIter->second; + whichFileNumber = fileNumbers[rank]; + whichFileName = amrex::VisMF::BaseName(amrex::NFilesIter::FileName(whichFileNumber, filePrefix)); + + for (const int idx : index) { + hdr.m_fod[idx].m_name = whichFileName; + hdr.m_fod[idx].m_head = currentOffset[whichFileNumber]; + currentOffset[whichFileNumber] += mf.fabbox(idx).numPts() * nComps * whichRDBytes + fabHeaderBytes[idx]; + } + } + } + } + } +} + +void write_2D_header(std::ostream &os, const amrex::FArrayBox &f, int nvar) +{ + os << "FAB " << amrex::FPC::NativeRealDescriptor(); + amrex::StreamRetry sr(os, "FABio_write_header", 4); + while (sr.TryOutput()) { + printLowerDimBox(os, f.box(), 2); + os << ' ' << nvar << '\n'; + } +} + +auto transform_box_to_2D(amrex::Direction const &dir, amrex::Box const &box) -> amrex::Box +{ + // transform box dimensions (Nx, Ny, Nz) -> (Nx', Ny', 1) where *one* of Nx, Ny, Nz == 1. + // NOTE: smallBox is assumed to be {0, 0, 0}. + amrex::IntVect dim = box.bigEnd(); + amrex::IntVect bigEnd; + + if (dir == amrex::Direction::x) { // y-z plane + bigEnd = amrex::IntVect(amrex::Dim3{dim[1], dim[2], 0}); +#if AMREX_SPACEDIM >= 2 + } else if (dir == amrex::Direction::y) { // x-z plane + bigEnd = amrex::IntVect(amrex::Dim3{dim[0], dim[2], 0}); +#endif +#if AMREX_SPACEDIM == 3 + } else if (dir == amrex::Direction::z) { // x-y plane + bigEnd = amrex::IntVect(amrex::Dim3{dim[0], dim[1], 0}); +#endif + } else { + amrex::Abort("detail::transform_box_to_2D: invalid direction!"); + } + + return amrex::Box(amrex::IntVect(amrex::Dim3{0, 0, 0}), bigEnd); +} + +auto transform_realbox_to_2D(amrex::Direction const &dir, amrex::RealBox const &box) -> amrex::RealBox +{ + // transform box dimensions (Nx, Ny, Nz) -> (Nx', Ny', 1) where *one* of Nx, Ny, Nz == 1. + // NOTE: smallBox is assumed to be {0, 0, 0}. + amrex::Real const *hi = box.hi(); + amrex::Real const *lo = box.lo(); + std::array new_hi{}; + std::array new_lo{}; + + if (dir == amrex::Direction::x) { // y-z plane + new_lo = {AMREX_D_DECL(lo[1], lo[2], lo[0])}; + new_hi = {AMREX_D_DECL(hi[1], hi[2], hi[0])}; +#if AMREX_SPACEDIM >= 2 + } else if (dir == amrex::Direction::y) { // x-z plane + new_lo = {AMREX_D_DECL(lo[0], lo[2], lo[1])}; + new_hi = {AMREX_D_DECL(hi[0], hi[2], hi[1])}; +#endif +#if AMREX_SPACEDIM == 3 + } else if (dir == amrex::Direction::z) { // x-y plane + new_lo = {AMREX_D_DECL(lo[0], lo[1], lo[2])}; + new_hi = {AMREX_D_DECL(hi[0], hi[1], hi[2])}; +#endif + } else { + amrex::Abort("detail::transform_box_to_2D: invalid direction!"); + } + + return amrex::RealBox(new_lo, new_hi); +} + +} // namespace detail + +void WriteProjection(const amrex::Direction dir, std::unordered_map> const &proj, amrex::Real time, int istep) +{ + // write projections to plotfile + auto const &firstFab = proj.begin()->second; + amrex::Vector varnames; + + // NOTE: Write2DMultiLevelPlotfile assumes the slice lies in the x-y plane + // (i.e. normal to the z axis) and the Geometry object corresponds to this. + // For a z-projection, this works as expected. For an {x,y}-projection, + // it is necessary to transform the geometry so that the data is stored in + // the x-y plane. + amrex::Geometry geom3d{}; + geom3d.Setup(); // read from ParmParse + const amrex::Box box2d = detail::transform_box_to_2D(dir, firstFab.box()); + const amrex::RealBox domain2d = detail::transform_realbox_to_2D(dir, geom3d.ProbDomain()); + const amrex::Geometry geom2d(box2d, &domain2d); + // amrex::Print() << box2d << "\n"; + // amrex::Print() << domain2d << "\n"; + + // construct output multifab on rank 0 + const amrex::BoxArray ba(box2d); + const amrex::DistributionMapping dm(amrex::Vector{0}); + const int ncomp = static_cast(proj.size()); + amrex::MultiFab mf_all(ba, dm, ncomp, 0); + + // copy all projections into a single Multifab with x-y geometry + auto iter = proj.begin(); + for (int icomp = 0; icomp < ncomp; ++icomp) { + const std::string &varname = iter->first; + const amrex::BaseFab &baseFab = iter->second; + varnames.push_back(varname); + // amrex::Print() << "varname: " << varname << " icomp: " << icomp << "\n"; + + // copy mf_comp into mf_all + auto output_arr = mf_all.arrays(); + auto const &input_arr = baseFab.const_array(); + + if (dir == amrex::Direction::x) { + amrex::ParallelFor(mf_all, + [=] AMREX_GPU_DEVICE(int bx, int i, int j, int k) noexcept { output_arr[bx](i, j, k, icomp) = input_arr(0, i, j); }); + } +#if AMREX_SPACEDIM >= 2 + else if (dir == amrex::Direction::y) { + amrex::ParallelFor(mf_all, + [=] AMREX_GPU_DEVICE(int bx, int i, int j, int k) noexcept { output_arr[bx](i, j, k, icomp) = input_arr(i, 0, j); }); + } +#endif +#if AMREX_SPACEDIM == 3 + else if (dir == amrex::Direction::z) { + amrex::ParallelFor(mf_all, + [=] AMREX_GPU_DEVICE(int bx, int i, int j, int k) noexcept { output_arr[bx](i, j, k, icomp) = input_arr(i, j, 0); }); + } +#endif + + amrex::Gpu::streamSynchronize(); + ++iter; + } + + // write mf_all to disk + const std::string basename = "proj_" + detail::direction_to_string(dir) + "_plt"; + const std::string filename = amrex::Concatenate(basename, istep, 5); + const amrex::Vector mfs = {&mf_all}; + amrex::Print() << "Writing projection " << filename << "\n"; + + detail::Write2DMultiLevelPlotfile(filename, 1, mfs, varnames, {geom2d}, time, {istep}, {}); +} + +} // namespace quokka::diagnostics diff --git a/src/io/projection.hpp b/src/io/projection.hpp new file mode 100644 index 000000000..fe9534837 --- /dev/null +++ b/src/io/projection.hpp @@ -0,0 +1,113 @@ +#ifndef PROJECTION_HPP_ // NOLINT +#define PROJECTION_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 projection.hpp +/// \brief AMReX I/O for 2D projections + +#include + +// AMReX headers +#include "AMReX_MultiFab.H" +#include "AMReX_MultiFabUtil.H" +#include "AMReX_Orientation.H" +#include "AMReX_VisMF.H" +#include + +namespace quokka::diagnostics +{ + +namespace detail +{ + +auto direction_to_string(const amrex::Direction dir) -> std::string; +auto transform_box_to_2D(amrex::Direction const &dir, amrex::Box const &box) -> amrex::Box; +auto transform_realbox_to_2D(amrex::Direction const &dir, amrex::RealBox const &box) -> amrex::RealBox; + +void printLowerDimIntVect(std::ostream &a_File, const amrex::IntVect &a_IntVect, int skipDim); +void printLowerDimBox(std::ostream &a_File, const amrex::Box &a_box, int skipDim); + +void Write2DMultiLevelPlotfile(const std::string &a_pltfile, int a_nlevels, const amrex::Vector &a_slice, + const amrex::Vector &a_varnames, const amrex::Vector &a_geoms, const amrex::Real &a_time, + const amrex::Vector &a_steps, const amrex::Vector &a_rref); + +void Write2DPlotfileHeader(std::ostream &HeaderFile, int nlevels, const amrex::Vector &bArray, const amrex::Vector &varnames, + const amrex::Vector &geom, const amrex::Real &time, const amrex::Vector &level_steps, + const amrex::Vector &ref_ratio, const std::string &versionName, const std::string &levelPrefix, + const std::string &mfPrefix); + +void VisMF2D(const amrex::MultiFab &a_mf, const std::string &a_mf_name); + +void Write2DMFHeader(const std::string &a_mf_name, amrex::VisMF::Header &hdr, int coordinatorProc, MPI_Comm comm); + +void Find2FOffsets(const amrex::FabArray &mf, const std::string &filePrefix, amrex::VisMF::Header &hdr, + amrex::VisMF::Header::Version /*whichVersion*/, amrex::NFilesIter &nfi, int nOutFiles, MPI_Comm comm); + +void write_2D_header(std::ostream &os, const amrex::FArrayBox &f, int nvar); + +} // namespace detail + +template +auto ComputePlaneProjection(amrex::Vector const &state_new, const int finest_level, amrex::Vector const &geom, + amrex::Vector const &ref_ratio, const amrex::Direction dir, F const &user_f) -> amrex::BaseFab +{ + // compute plane-parallel projection of user_f(i, j, k, state) along the given axis. + BL_PROFILE("quokka::DiagProjection::computePlaneProjection()"); + + // allocate temporary multifabs + amrex::Vector q; + q.resize(finest_level + 1); + + for (int lev = 0; lev <= finest_level; ++lev) { + q[lev].define(state_new[lev].boxArray(), state_new[lev].DistributionMap(), 1, 0); + } + + // evaluate user_f on all levels + for (int lev = 0; lev <= finest_level; ++lev) { + auto const &state = state_new[lev].const_arrays(); + auto const &result = q[lev].arrays(); + amrex::ParallelFor(q[lev], [=] AMREX_GPU_DEVICE(int bx, int i, int j, int k) { result[bx](i, j, k) = user_f(i, j, k, state[bx]); }); + } + amrex::Gpu::streamSynchronize(); + + // average down + for (int lev = finest_level; lev < 0; --lev) { + amrex::average_down(q[lev], q[lev - 1], geom[lev], geom[lev - 1], 0, 1, ref_ratio[lev - 1]); + } + + auto const &domain_box = geom[0].Domain(); + auto const &dx = geom[0].CellSizeArray(); + auto const &arr = q[0].const_arrays(); + amrex::BaseFab proj = + amrex::ReduceToPlane(int(dir), domain_box, q[0], [=] AMREX_GPU_DEVICE(int box_no, int i, int j, int k) -> amrex::Real { + return dx[int(dir)] * arr[box_no](i, j, k); // data at (i,j,k) of Box box_no + }); + amrex::Gpu::streamSynchronize(); + + // copy to host pinned memory to work around AMReX bug + amrex::BaseFab proj_host(proj.box(), 1, amrex::The_Pinned_Arena()); + proj_host.copy(proj); + amrex::Gpu::streamSynchronize(); + + if constexpr (std::is_same::value) { + amrex::ParallelReduce::Sum(proj_host.dataPtr(), static_cast(proj_host.size()), amrex::ParallelDescriptor::ioProcessor, + amrex::ParallelDescriptor::Communicator()); + } else if constexpr (std::is_same::value) { + amrex::ParallelReduce::Min(proj_host.dataPtr(), static_cast(proj_host.size()), amrex::ParallelDescriptor::ioProcessor, + amrex::ParallelDescriptor::Communicator()); + } else { + amrex::Abort("invalid reduce op!"); + } + + // return BaseFab in host memory + return proj_host; +} + +void WriteProjection(const amrex::Direction dir, std::unordered_map> const &proj, amrex::Real time, int istep); + +} // namespace quokka::diagnostics + +#endif // PROJECTION_HPP_ \ No newline at end of file diff --git a/src/linear_advection/AdvectionSimulation.hpp b/src/linear_advection/AdvectionSimulation.hpp index 13f4325a5..6c9f24312 100644 --- a/src/linear_advection/AdvectionSimulation.hpp +++ b/src/linear_advection/AdvectionSimulation.hpp @@ -78,7 +78,7 @@ template class AdvectionSimulation : public AMRSimulation

std::unordered_map> override; + [[nodiscard]] auto ComputeProjections(const amrex::Direction dir) const -> std::unordered_map> override; // compute statistics auto ComputeStatistics() -> std::map override; @@ -173,7 +173,7 @@ template void AdvectionSimulation::ComputeDerive } template -auto AdvectionSimulation::ComputeProjections(int /*dir*/) const -> std::unordered_map> +auto AdvectionSimulation::ComputeProjections(const amrex::Direction /*dir*/) const -> std::unordered_map> { // user should implement return std::unordered_map>{}; diff --git a/src/problems/NSCBC/channel.cpp b/src/problems/NSCBC/channel.cpp index 21a36834d..955160a56 100644 --- a/src/problems/NSCBC/channel.cpp +++ b/src/problems/NSCBC/channel.cpp @@ -295,7 +295,7 @@ auto problem_main() -> int // Compute test success condition int status = 0; - const double error_tol = 3.0e-5; + const double error_tol = 3.5e-5; if (epsilon > error_tol) { status = 1; } diff --git a/src/problems/ShockCloud/cloud.cpp b/src/problems/ShockCloud/cloud.cpp index aec244207..cb5ca3a70 100644 --- a/src/problems/ShockCloud/cloud.cpp +++ b/src/problems/ShockCloud/cloud.cpp @@ -29,6 +29,7 @@ #include "hydro/NSCBC_inflow.hpp" #include "hydro/NSCBC_outflow.hpp" #include "hydro/hydro_system.hpp" +#include "io/projection.hpp" #include "physics_info.hpp" #include "radiation/radiation_system.hpp" @@ -600,35 +601,33 @@ template <> auto QuokkaSimulation::ComputeStatistics() -> std::map auto QuokkaSimulation::ComputeProjections(const int dir) const -> std::unordered_map> +template <> +auto QuokkaSimulation::ComputeProjections(const amrex::Direction dir) const -> std::unordered_map> { std::unordered_map> proj; // compute (total) density projection - proj["nH"] = computePlaneProjection( - [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::Array4 const &state) noexcept { + proj["nH"] = quokka::diagnostics::ComputePlaneProjection( + state_new_cc_, finestLevel(), geom, ref_ratio, dir, [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::Array4 const &state) noexcept { Real const rho = state(i, j, k, HydroSystem::density_index); return (quokka::TabulatedCooling::cloudy_H_mass_fraction * rho) / m_H; - }, - dir); + }); // compute cloud partial density projection - proj["nH_cloud"] = computePlaneProjection( - [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::Array4 const &state) noexcept { + proj["nH_cloud"] = quokka::diagnostics::ComputePlaneProjection( + state_new_cc_, finestLevel(), geom, ref_ratio, dir, [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::Array4 const &state) noexcept { // partial cloud density Real const rho_cloud = state(i, j, k, HydroSystem::scalar0_index + 1); return (quokka::TabulatedCooling::cloudy_H_mass_fraction * rho_cloud) / m_H; - }, - dir); + }); // compute non-cloud partial density projection - proj["nH_wind"] = computePlaneProjection( - [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::Array4 const &state) noexcept { + proj["nH_wind"] = quokka::diagnostics::ComputePlaneProjection( + state_new_cc_, finestLevel(), geom, ref_ratio, dir, [=] AMREX_GPU_DEVICE(int i, int j, int k, amrex::Array4 const &state) noexcept { // partial wind density Real const rho_wind = state(i, j, k, HydroSystem::scalar0_index + 2); return (quokka::TabulatedCooling::cloudy_H_mass_fraction * rho_wind) / m_H; - }, - dir); + }); return proj; } diff --git a/src/simulation.hpp b/src/simulation.hpp index 08d64debf..3019bcbe9 100644 --- a/src/simulation.hpp +++ b/src/simulation.hpp @@ -51,6 +51,7 @@ namespace filesystem = experimental::filesystem; #include "AMReX_IntVect.H" #include "AMReX_Interpolater.H" #include "AMReX_MultiFabUtil.H" +#include "AMReX_Orientation.H" #include "AMReX_ParallelDescriptor.H" #include "AMReX_REAL.H" #include "AMReX_SPACE.H" @@ -85,6 +86,7 @@ namespace filesystem = experimental::filesystem; #include "fundamental_constants.H" #include "grid.hpp" #include "io/DiagBase.H" +#include "io/projection.hpp" #include "physics_info.hpp" #ifdef QUOKKA_USE_OPENPMD @@ -234,7 +236,7 @@ template class AMRSimulation : public amrex::AmrCore virtual void ComputeDerivedVar(int lev, std::string const &dname, amrex::MultiFab &mf, int ncomp) const = 0; // compute projected vars - [[nodiscard]] virtual auto ComputeProjections(int dir) const -> std::unordered_map> = 0; + [[nodiscard]] virtual auto ComputeProjections(const amrex::Direction dir) const -> std::unordered_map> = 0; // compute statistics virtual auto ComputeStatistics() -> std::map = 0; @@ -304,8 +306,6 @@ template class AMRSimulation : public amrex::AmrCore int bcomp, int orig_comp); // template specialized by problem generator - template auto computePlaneProjection(F const &user_f, int dir) const -> amrex::BaseFab; - // compute volume integrals template auto computeVolumeIntegral(F const &user_f) -> amrex::Real; @@ -2441,61 +2441,6 @@ template void AMRSimulation::ReadMetadataFile(st } } -template -template -auto AMRSimulation::computePlaneProjection(F const &user_f, const int dir) const -> amrex::BaseFab -{ - // compute plane-parallel projection of user_f(i, j, k, state) along the given axis. - BL_PROFILE("AMRSimulation::computePlaneProjection()"); - - // allocate temporary multifabs - amrex::Vector q; - q.resize(finest_level + 1); - for (int lev = 0; lev <= finest_level; ++lev) { - q[lev].define(boxArray(lev), DistributionMap(lev), 1, 0); - } - - // evaluate user_f on all levels - for (int lev = 0; lev <= finest_level; ++lev) { - auto const &state = state_new_cc_[lev].const_arrays(); - auto const &result = q[lev].arrays(); - amrex::ParallelFor(q[lev], [=] AMREX_GPU_DEVICE(int bx, int i, int j, int k) { result[bx](i, j, k) = user_f(i, j, k, state[bx]); }); - } - amrex::Gpu::streamSynchronize(); - - // average down - for (int lev = finest_level; lev < 0; --lev) { - amrex::average_down(q[lev], q[lev - 1], geom[lev], geom[lev - 1], 0, 1, ref_ratio[lev - 1]); - } - - auto const &domain_box = geom[0].Domain(); - auto const &dx = geom[0].CellSizeArray(); - auto const &arr = q[0].const_arrays(); - amrex::BaseFab proj = - amrex::ReduceToPlane(dir, domain_box, q[0], [=] AMREX_GPU_DEVICE(int box_no, int i, int j, int k) -> amrex::Real { - return dx[dir] * arr[box_no](i, j, k); // data at (i,j,k) of Box box_no - }); - amrex::Gpu::streamSynchronize(); - - // copy to host pinned memory to work around AMReX bug - amrex::BaseFab proj_host(proj.box(), 1, amrex::The_Pinned_Arena()); - proj_host.copy(proj); - amrex::Gpu::streamSynchronize(); - - if constexpr (std::is_same::value) { - amrex::ParallelReduce::Sum(proj_host.dataPtr(), static_cast(proj_host.size()), amrex::ParallelDescriptor::ioProcessor, - amrex::ParallelDescriptor::Communicator()); - } else if constexpr (std::is_same::value) { - amrex::ParallelReduce::Min(proj_host.dataPtr(), static_cast(proj_host.size()), amrex::ParallelDescriptor::ioProcessor, - amrex::ParallelDescriptor::Communicator()); - } else { - amrex::Abort("invalid reduce op!"); - } - - // return BaseFab in host memory - return proj_host; -} - template void AMRSimulation::WriteProjectionPlotfile() const { std::vector dirs{}; @@ -2504,50 +2449,27 @@ template void AMRSimulation::WriteProjectionPlot auto dir_from_string = [=](const std::string &dir_str) { if (dir_str == "x") { - return 0; + return amrex::Direction::x; } +#if AMREX_SPACEDIM >= 2 if (dir_str == "y") { - return 1; + return amrex::Direction::y; } +#endif +#if AMREX_SPACEDIM == 3 if (dir_str == "z") { - return 2; + return amrex::Direction::z; } - return -1; +#endif + amrex::Error("invalid direction for projection!"); + return amrex::Direction::x; }; for (auto &dir_str : dirs) { // compute projections along axis 'dir' - int dir = dir_from_string(dir_str); + amrex::Direction dir = dir_from_string(dir_str); std::unordered_map> proj = ComputeProjections(dir); - - auto const &firstFab = proj.begin()->second; - const amrex::BoxArray ba(firstFab.box()); - const amrex::DistributionMapping dm(amrex::Vector{0}); - amrex::MultiFab mf_all(ba, dm, static_cast(proj.size()), 0); - amrex::Vector varnames; - - // write 2D plotfiles - auto iter = proj.begin(); - for (int icomp = 0; icomp < static_cast(proj.size()); ++icomp) { - const std::string &varname = iter->first; - const amrex::BaseFab &baseFab = iter->second; - - const amrex::BoxArray ba(baseFab.box()); - const amrex::DistributionMapping dm(amrex::Vector{0}); - amrex::MultiFab mf(ba, dm, 1, 0, amrex::MFInfo().SetAlloc(false)); - if (amrex::ParallelDescriptor::IOProcessor()) { - mf.setFab(0, amrex::FArrayBox(baseFab.array())); - } - amrex::MultiFab::Copy(mf_all, mf, 0, icomp, 1, 0); - varnames.push_back(varname); - ++iter; - } - - const std::string basename = "proj_" + dir_str + "_plt"; - const std::string filename = amrex::Concatenate(basename, istep[0], 5); - amrex::Print() << "Writing projection " << filename << "\n"; - const amrex::Geometry mygeom(firstFab.box()); - amrex::WriteSingleLevelPlotfile(filename, mf_all, varnames, mygeom, tNew_[0], istep[0]); + quokka::diagnostics::WriteProjection(dir, proj, tNew_[0], istep[0]); } }