Skip to content

Commit

Permalink
Merge pull request #92 from seahorce-scidac/output_loc
Browse files Browse the repository at this point in the history 
optionally output cell-centered location
  • Loading branch information
@asalmgren
asalmgren authored Dec 15, 2023
2 parents cad1533 + 8f3461d commit a3dbcf0
Show file tree
Hide file tree
Showing 3 changed files with 67 additions and 51 deletions.
71 changes: 43 additions & 28 deletions Source/IO/Plotfile.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,8 @@ ROMSX::setPlotVariables (const std::string& pp_plot_var_names, Vector<std::strin
tmp_plot_names.push_back(cons_names[i]);
}
}
// check for velocity since it's not in cons_names
// if we are asked for any velocity component, we will need them all
// Check for velocity since it's not in cons_names
// If we are asked for any velocity component, we will need them all
if (containerHasElement(plot_var_names, "x_velocity") ||
containerHasElement(plot_var_names, "y_velocity") ||
containerHasElement(plot_var_names, "z_velocity")) {
Expand All @@ -52,6 +52,15 @@ ROMSX::setPlotVariables (const std::string& pp_plot_var_names, Vector<std::strin
tmp_plot_names.push_back("z_velocity");
}

// If we are asked for any location component, we will provide them all
if (containerHasElement(plot_var_names, "x_cc") ||
containerHasElement(plot_var_names, "y_cc") ||
containerHasElement(plot_var_names, "z_cc")) {
tmp_plot_names.push_back("x_cc");
tmp_plot_names.push_back("y_cc");
tmp_plot_names.push_back("z_cc");
}

for (int i = 0; i < derived_names.size(); ++i) {
if ( containerHasElement(plot_var_names, derived_names[i]) ) {
#ifdef ROMSX_USE_PARTICLES
Expand Down Expand Up @@ -89,7 +98,6 @@ ROMSX::PlotFileVarNames ( Vector<std::string> plot_var_names ) const
void
ROMSX::WritePlotFile (int which, Vector<std::string> plot_var_names)
{

const Vector<std::string> varnames = PlotFileVarNames(plot_var_names);
const int ncomp_mf = varnames.size();
const auto ngrow_vars = IntVect(NGROW-1,NGROW-1,0);
Expand Down Expand Up @@ -158,50 +166,57 @@ ROMSX::WritePlotFile (int which, Vector<std::string> plot_var_names)
mf_comp += AMREX_SPACEDIM;
}

// Finally, check for any derived quantities and compute them, inserting
// them into our output multifab
auto calculate_derived = [&](std::string der_name,
decltype(derived::romsx_dernull)& der_function)
{
if (containerHasElement(plot_var_names, der_name)) {
MultiFab dmf(mf[lev], make_alias, mf_comp, 1);

for (MFIter mfi(dmf, TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const Box& bx = mfi.tilebox();
auto& dfab = dmf[mfi];
auto& sfab = (*cons_new[lev])[mfi];
der_function(bx, dfab, 0, 1, sfab, Geom(lev), t_new[0], nullptr, lev);
}
// Fill cell-centered location
Real dx = Geom()[lev].CellSizeArray()[0];
Real dy = Geom()[lev].CellSizeArray()[1];

mf_comp++;
}
};
// Next, check for location names -- if we write one we write all
if (containerHasElement(plot_var_names, "x_cc") ||
containerHasElement(plot_var_names, "y_cc") ||
containerHasElement(plot_var_names, "z_cc"))
{
MultiFab dmf(mf[lev], make_alias, mf_comp, AMREX_SPACEDIM);
for (MFIter mfi(dmf, TilingIfNotGPU()); mfi.isValid(); ++mfi) {
const Box& bx = mfi.tilebox();
const Array4<Real> loc_arr = dmf.array(mfi);
const Array4<Real const> zp_arr = vec_z_phys_nd[lev]->const_array(mfi);

ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) {
loc_arr(i,j,k,0) = (i+0.5) * dx;
loc_arr(i,j,k,1) = (j+0.5) * dy;
loc_arr(i,j,k,2) = Real(0.125) * (zp_arr(i,j ,k ) + zp_arr(i+1,j ,k ) +
zp_arr(i,j+1,k ) + zp_arr(i+1,j+1,k ) +
zp_arr(i,j ,k+1) + zp_arr(i+1,j ,k+1) +
zp_arr(i,j+1,k+1) + zp_arr(i+1,j+1,k+1) );
});
} // mfi
mf_comp += AMREX_SPACEDIM;
} // if containerHasElement

#ifdef ROMSX_USE_PARTICLES
if (containerHasElement(plot_var_names, "tracer_particle_count"))
{
MultiFab temp_dat(mf[lev].boxArray(), mf[lev].DistributionMap(), 1, 0);
temp_dat.setVal(0);
particleData.tracer_particles->Redistribute();
particleData.tracer_particles->Increment(temp_dat, lev);
MultiFab::Copy(mf[lev], temp_dat, 0, mf_comp, 1, 0);
mf_comp += 1;
}
#endif
}

// Fill terrain distortion MF
for (int lev(0); lev <= finest_level; ++lev) {
MultiFab::Copy(mf_nd[lev],*vec_z_phys_nd[lev],0,2,1,0);
Real dz = Geom()[lev].CellSizeArray()[2];
for (MFIter mfi(mf_nd[lev], TilingIfNotGPU()); mfi.isValid(); ++mfi) {
for (MFIter mfi(mf_nd[lev], TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const Box& bx = mfi.tilebox();
Array4< Real> mf_arr = mf_nd[lev].array(mfi);
Array4<Real> mf_arr = mf_nd[lev].array(mfi);
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) {
mf_arr(i,j,k,2) -= k * dz;
});
}
}
} // mfi

} // lev

std::string plotfilename;
if (which == 1)
Expand Down
1 change: 0 additions & 1 deletion Source/ROMSX.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -817,7 +817,6 @@ private:

amrex::Vector<std::string> plot_var_names_1;
amrex::Vector<std::string> plot_var_names_2;
const amrex::Vector<std::string> velocity_names {"x_velocity", "y_velocity", "z_velocity"};
const amrex::Vector<std::string> cons_names {"temp", "salt", "scalar"};

// Note that the order of variable names here must match the order in Derive.cpp
Expand Down
46 changes: 24 additions & 22 deletions Source/Utils/DepthStretchTransform.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ ROMSX::stretch_transform (int lev)

for ( MFIter mfi(*cons_new[lev], TilingIfNotGPU()); mfi.isValid(); ++mfi )
{

Array4<Real> const& z_w = (mf_z_w)->array(mfi);
Array4<Real> const& z_r = (mf_z_r)->array(mfi);
Array4<Real> const& s_r = (mf_s_r)->array(mfi);
Expand Down Expand Up @@ -56,10 +55,6 @@ ROMSX::stretch_transform (int lev)
z_w(i,j,k) = h(i,j,0);
});

if (verbose > 2) {
PrintToFile("Zt_avg1_dst").SetPrecision(18) << FArrayBox(Zt_avg1) << std::endl;
}

// ROMS Transform 2
Gpu::streamSynchronize();
amrex::ParallelFor(gbx3D,
Expand Down Expand Up @@ -195,7 +190,7 @@ ROMSX::stretch_transform (int lev)
Hz(i,j,k)=z_w(i,j,k)-z_w(i,j,k-1);
}
});
}
} // mfi

Real time = t_new[lev];
FillPatch(lev, time, *vec_z_w[lev], GetVecOfPtrs(vec_z_w));
Expand All @@ -217,23 +212,30 @@ ROMSX::stretch_transform (int lev)
// WARNING: z_w(i,j,k) refers to the face on the HIGH side of cell (i,j,k)
// WARNING: z_phys_nd(i,j,k) refers to the node on the LOW side of cell (i,j,k)
//
amrex::ParallelFor(Box(z_phys_nd),
[=] AMREX_GPU_DEVICE (int i, int j, int k)
{
// For now assume all boundaries are constant height --
// we will enforce periodicity below
int kk = (k == 0) ? hi.z : k-1;
if ( i >= lo.x && i <= hi.x-1 && j >= lo.y && j <= hi.y-1 )
if (solverChoice.flat_bathymetry) {
Real dz = geom[lev].CellSize(2);
ParallelFor(Box(z_phys_nd), [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
z_phys_nd(i,j,k)=0.25*( z_w(i,j ,kk) + z_w(i+1,j ,kk) +
z_w(i,j+1,kk) + z_w(i+1,j+1,kk) );
} else {
int ii = std::min(std::max(i, lo.x), hi.x);
int jj = std::min(std::max(j, lo.y), hi.y);
z_phys_nd(i,j,k) = z_w(ii,jj,kk);
}
if (k == 0) z_phys_nd(i,j,k) *= -1.;
});
z_phys_nd(i,j,k) = k * dz;
});
} else {
ParallelFor(Box(z_phys_nd), [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
// For now assume all boundaries are constant height --
// we will enforce periodicity below
int kk = (k == 0) ? hi.z : k-1;
if ( i >= lo.x && i <= hi.x-1 && j >= lo.y && j <= hi.y-1 )
{
z_phys_nd(i,j,k)=0.25*( z_w(i,j ,kk) + z_w(i+1,j ,kk) +
z_w(i,j+1,kk) + z_w(i+1,j+1,kk) );
} else {
int ii = std::min(std::max(i, lo.x), hi.x);
int jj = std::min(std::max(j, lo.y), hi.y);
z_phys_nd(i,j,k) = z_w(ii,jj,kk);
}
if (k == 0) z_phys_nd(i,j,k) *= -1.;
});
} // if not flat
} // mf

// Note that we do *not* want to do a multilevel fill here -- we have
Expand Down

0 comments on commit a3dbcf0

Please sign in to comment.