From b169a30c59d57e80eb83fcad77d3fb76d4350334 Mon Sep 17 00:00:00 2001 From: Ben Wibking Date: Sat, 9 Sep 2023 07:38:50 -0400 Subject: [PATCH] add axis-aligned projection outputs (#367) * add axis-aligned projection outputs * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * add const * fix warning --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Piyush Sharda <34922596+psharda@users.noreply.github.com> --- src/AdvectionSimulation.hpp | 9 +++ src/RadhydroSimulation.hpp | 15 +++- src/simulation.hpp | 147 ++++++++++++++++++++++++++++++++++-- 3 files changed, 163 insertions(+), 8 deletions(-) diff --git a/src/AdvectionSimulation.hpp b/src/AdvectionSimulation.hpp index bf1522d3b..acbb619c9 100644 --- a/src/AdvectionSimulation.hpp +++ b/src/AdvectionSimulation.hpp @@ -81,6 +81,8 @@ template class AdvectionSimulation : public AMRSimulation

std::unordered_map> override; // compute statistics auto ComputeStatistics() -> std::map override; @@ -155,6 +157,13 @@ template void AdvectionSimulation::ComputeDerive // user should implement } +template +auto AdvectionSimulation::ComputeProjections(int /*dir*/) const -> std::unordered_map> +{ + // user should implement + return std::unordered_map>{}; +} + template auto AdvectionSimulation::ComputeStatistics() -> std::map { diff --git a/src/RadhydroSimulation.hpp b/src/RadhydroSimulation.hpp index 39a640ef9..2ed5090d3 100644 --- a/src/RadhydroSimulation.hpp +++ b/src/RadhydroSimulation.hpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include "AMReX.H" @@ -191,8 +192,11 @@ template class RadhydroSimulation : public AMRSimulation std::map override; + // compute projected vars + [[nodiscard]] auto ComputeProjections(int dir) const -> std::unordered_map> override; + + // compute statistics + auto ComputeStatistics() -> std::map override; // fix-up states void FixupState(int level) override; @@ -497,6 +501,13 @@ void RadhydroSimulation::ComputeDerivedVar(int lev, std::string const // compute derived variables and save in 'mf' -- user should implement } +template +auto RadhydroSimulation::ComputeProjections(int /*dir*/) const -> std::unordered_map> +{ + // compute projections and return as unordered_map -- user should implement + return std::unordered_map>{}; +} + template auto RadhydroSimulation::ComputeStatistics() -> std::map { diff --git a/src/simulation.hpp b/src/simulation.hpp index c2033ffef..96e6dc795 100644 --- a/src/simulation.hpp +++ b/src/simulation.hpp @@ -114,6 +114,7 @@ template class AMRSimulation : public amrex::AmrCore amrex::Long maxWalltime_ = 0; // default: no limit int ascentInterval_ = -1; // -1 == no in-situ renders with Ascent int plotfileInterval_ = -1; // -1 == no output + int projectionInterval_ = -1; // -1 == no output int statisticsInterval_ = -1; // -1 == no output amrex::Real plotTimeInterval_ = -1.0; // time interval for plt file amrex::Real checkpointTimeInterval_ = -1.0; // time interval for checkpoints @@ -159,8 +160,11 @@ template class AMRSimulation : public amrex::AmrCore // compute derived variables virtual void ComputeDerivedVar(int lev, std::string const &dname, amrex::MultiFab &mf, int ncomp) const = 0; - // compute statistics - virtual auto ComputeStatistics() -> std::map = 0; + // compute projected vars + [[nodiscard]] virtual auto ComputeProjections(int dir) const -> std::unordered_map> = 0; + + // compute statistics + virtual auto ComputeStatistics() -> std::map = 0; // fix-up any unphysical states created by AMR operations // (e.g., caused by the flux register or from interpolation) @@ -218,6 +222,8 @@ template class AMRSimulation : public amrex::AmrCore amrex::GeometryData const &geom, amrex::Real time, const amrex::BCRec *bcr, int bcomp, int orig_comp); // template specialized by problem generator + template auto computePlaneProjection(F const &user_f, int dir) const -> amrex::BaseFab; + // I/O functions [[nodiscard]] auto PlotFileName(int lev) const -> std::string; [[nodiscard]] auto CustomPlotFileName(const char *base, int lev) const -> std::string; @@ -228,6 +234,7 @@ template class AMRSimulation : public amrex::AmrCore void ReadMetadataFile(std::string const &chkfilename); void WriteStatisticsFile(); void WritePlotFile(); + void WriteProjectionPlotfile() const; void WriteCheckpointFile() const; void SetLastCheckpointSymlink(std::string const &checkpointname) const; void ReadCheckpointFile(); @@ -439,8 +446,11 @@ template void AMRSimulation::readParameters() // Default output interval pp.query("plotfile_interval", plotfileInterval_); - // Default statistics interval - pp.query("statistics_interval", statisticsInterval_); + // Default projection interval + pp.query("projection_interval", projectionInterval_); + + // Default statistics interval + pp.query("statistics_interval", statisticsInterval_); // Default Time interval pp.query("plottime_interval", plotTimeInterval_); @@ -533,7 +543,11 @@ template void AMRSimulation::setInitialCondition WritePlotFile(); } - if (statisticsInterval_ > 0) { + if (projectionInterval_ > 0) { + WriteProjectionPlotfile(); + } + + if (statisticsInterval_ > 0) { WriteStatisticsFile(); } @@ -675,6 +689,7 @@ template void AMRSimulation::evolve() #ifdef AMREX_USE_ASCENT int last_ascent_step = 0; #endif + int last_projection_step = 0; int last_statistics_step = 0; int last_plot_file_step = 0; double next_plot_file_time = plotTimeInterval_; @@ -739,6 +754,11 @@ template void AMRSimulation::evolve() WritePlotFile(); } + if (projectionInterval_ > 0 && (step + 1) % projectionInterval_ == 0) { + last_projection_step = step + 1; + WriteProjectionPlotfile(); + } + // Writing Plot files at time intervals if (plotTimeInterval_ > 0 && next_plot_file_time <= cur_time) { next_plot_file_time += plotTimeInterval_; @@ -805,7 +825,12 @@ template void AMRSimulation::evolve() WritePlotFile(); } - // write final statistics + // write final projection + if (projectionInterval_ > 0 && istep[0] > last_projection_step) { + WriteProjectionPlotfile(); + } + + // write final statistics if (statisticsInterval_ > 0 && istep[0] > last_statistics_step) { WriteStatisticsFile(); } @@ -1778,6 +1803,116 @@ template void AMRSimulation::WritePlotFile() #endif } +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{}; + const amrex::ParmParse pp; + pp.queryarr("projection.dirs", dirs); + + auto dir_from_string = [=](const std::string &dir_str) { + if (dir_str == "x") { + return 0; + } + if (dir_str == "y") { + return 1; + } + if (dir_str == "z") { + return 2; + } + return -1; + }; + + for (auto &dir_str : dirs) { + // compute projections along axis 'dir' + int 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; + 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]); + } +} + template void AMRSimulation::WriteStatisticsFile() { // append to statistics file