Skip to content

Commit

Permalink
add axis-aligned projection outputs (#367)
Browse files Browse the repository at this point in the history 
* 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 <[email protected]>
  • Loading branch information
@BenWibking @github-actions @psharda
3 people authored Sep 9, 2023
1 parent bce9c28 commit b169a30
Show file tree
Hide file tree
Showing 3 changed files with 163 additions and 8 deletions.
9 changes: 9 additions & 0 deletions src/AdvectionSimulation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,8 @@ template <typename problem_t> class AdvectionSimulation : public AMRSimulation<p

// compute derived variables
void ComputeDerivedVar(int lev, std::string const &dname, amrex::MultiFab &mf, int ncomp) const override;
// compute projected vars
[[nodiscard]] auto ComputeProjections(int dir) const -> std::unordered_map<std::string, amrex::BaseFab<amrex::Real>> override;

// compute statistics
auto ComputeStatistics() -> std::map<std::string, amrex::Real> override;
Expand Down Expand Up @@ -155,6 +157,13 @@ template <typename problem_t> void AdvectionSimulation<problem_t>::ComputeDerive
// user should implement
}

template <typename problem_t>
auto AdvectionSimulation<problem_t>::ComputeProjections(int /*dir*/) const -> std::unordered_map<std::string, amrex::BaseFab<amrex::Real>>
{
// user should implement
return std::unordered_map<std::string, amrex::BaseFab<amrex::Real>>{};
}

template <typename problem_t>
auto AdvectionSimulation<problem_t>::ComputeStatistics() -> std::map<std::string, amrex::Real>
{
Expand Down
15 changes: 13 additions & 2 deletions src/RadhydroSimulation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#include <limits>
#include <string>
#include <tuple>
#include <unordered_map>
#include <utility>

#include "AMReX.H"
Expand Down Expand Up @@ -191,8 +192,11 @@ template <typename problem_t> class RadhydroSimulation : public AMRSimulation<pr
// compute derived variables
void ComputeDerivedVar(int lev, std::string const &dname, amrex::MultiFab &mf, int ncomp) const override;

// compute statistics
auto ComputeStatistics() -> std::map<std::string, amrex::Real> override;
// compute projected vars
[[nodiscard]] auto ComputeProjections(int dir) const -> std::unordered_map<std::string, amrex::BaseFab<amrex::Real>> override;

// compute statistics
auto ComputeStatistics() -> std::map<std::string, amrex::Real> override;

// fix-up states
void FixupState(int level) override;
Expand Down Expand Up @@ -497,6 +501,13 @@ void RadhydroSimulation<problem_t>::ComputeDerivedVar(int lev, std::string const
// compute derived variables and save in 'mf' -- user should implement
}

template <typename problem_t>
auto RadhydroSimulation<problem_t>::ComputeProjections(int /*dir*/) const -> std::unordered_map<std::string, amrex::BaseFab<amrex::Real>>
{
// compute projections and return as unordered_map -- user should implement
return std::unordered_map<std::string, amrex::BaseFab<amrex::Real>>{};
}

template <typename problem_t>
auto RadhydroSimulation<problem_t>::ComputeStatistics() -> std::map<std::string, amrex::Real>
{
Expand Down
147 changes: 141 additions & 6 deletions src/simulation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,7 @@ template <typename problem_t> 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
Expand Down Expand Up @@ -159,8 +160,11 @@ template <typename problem_t> 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<std::string, amrex::Real> = 0;
// compute projected vars
[[nodiscard]] virtual auto ComputeProjections(int dir) const -> std::unordered_map<std::string, amrex::BaseFab<amrex::Real>> = 0;

// compute statistics
virtual auto ComputeStatistics() -> std::map<std::string, amrex::Real> = 0;

// fix-up any unphysical states created by AMR operations
// (e.g., caused by the flux register or from interpolation)
Expand Down Expand Up @@ -218,6 +222,8 @@ template <typename problem_t> 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 <typename ReduceOp, typename F> auto computePlaneProjection(F const &user_f, int dir) const -> amrex::BaseFab<amrex::Real>;

// I/O functions
[[nodiscard]] auto PlotFileName(int lev) const -> std::string;
[[nodiscard]] auto CustomPlotFileName(const char *base, int lev) const -> std::string;
Expand All @@ -228,6 +234,7 @@ template <typename problem_t> 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();
Expand Down Expand Up @@ -439,8 +446,11 @@ template <typename problem_t> void AMRSimulation<problem_t>::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_);
Expand Down Expand Up @@ -533,7 +543,11 @@ template <typename problem_t> void AMRSimulation<problem_t>::setInitialCondition
WritePlotFile();
}

if (statisticsInterval_ > 0) {
if (projectionInterval_ > 0) {
WriteProjectionPlotfile();
}

if (statisticsInterval_ > 0) {
WriteStatisticsFile();
}

Expand Down Expand Up @@ -675,6 +689,7 @@ template <typename problem_t> void AMRSimulation<problem_t>::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_;
Expand Down Expand Up @@ -739,6 +754,11 @@ template <typename problem_t> void AMRSimulation<problem_t>::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_;
Expand Down Expand Up @@ -805,7 +825,12 @@ template <typename problem_t> void AMRSimulation<problem_t>::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();
}
Expand Down Expand Up @@ -1778,6 +1803,116 @@ template <typename problem_t> void AMRSimulation<problem_t>::WritePlotFile()
#endif
}

template <typename problem_t>
template <typename ReduceOp, typename F>
auto AMRSimulation<problem_t>::computePlaneProjection(F const &user_f, const int dir) const -> amrex::BaseFab<amrex::Real>
{
// compute plane-parallel projection of user_f(i, j, k, state) along the given axis.
BL_PROFILE("AMRSimulation::computePlaneProjection()");

// allocate temporary multifabs
amrex::Vector<amrex::MultiFab> 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<amrex::Real> proj =
amrex::ReduceToPlane<ReduceOp, amrex::Real>(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<amrex::Real> proj_host(proj.box(), 1, amrex::The_Pinned_Arena());
proj_host.copy<amrex::RunOn::Device>(proj);
amrex::Gpu::streamSynchronize();

if constexpr (std::is_same<ReduceOp, amrex::ReduceOpSum>::value) {
amrex::ParallelReduce::Sum(proj_host.dataPtr(), static_cast<int>(proj_host.size()), amrex::ParallelDescriptor::ioProcessor,
amrex::ParallelDescriptor::Communicator());
} else if constexpr (std::is_same<ReduceOp, amrex::ReduceOpMin>::value) {
amrex::ParallelReduce::Min(proj_host.dataPtr(), static_cast<int>(proj_host.size()), amrex::ParallelDescriptor::ioProcessor,
amrex::ParallelDescriptor::Communicator());
} else {
amrex::Abort("invalid reduce op!");
}

// return BaseFab in host memory
return proj_host;
}

template <typename problem_t> void AMRSimulation<problem_t>::WriteProjectionPlotfile() const
{
std::vector<std::string> 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<std::string, amrex::BaseFab<amrex::Real>> proj = ComputeProjections(dir);

auto const &firstFab = proj.begin()->second;
const amrex::BoxArray ba(firstFab.box());
const amrex::DistributionMapping dm(amrex::Vector<int>{0});
amrex::MultiFab mf_all(ba, dm, static_cast<int>(proj.size()), 0);
amrex::Vector<std::string> varnames;

// write 2D plotfiles
auto iter = proj.begin();
for (int icomp = 0; icomp < static_cast<int>(proj.size()); ++icomp) {
const std::string &varname = iter->first;
const amrex::BaseFab<amrex::Real> &baseFab = iter->second;

const amrex::BoxArray ba(baseFab.box());
const amrex::DistributionMapping dm(amrex::Vector<int>{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 <typename problem_t>
void AMRSimulation<problem_t>::WriteStatisticsFile() {
// append to statistics file
Expand Down

0 comments on commit b169a30

Please sign in to comment.