Skip to content

Commit

Permalink
fix hybrid FFT when we used stretched grids
Browse files Browse the repository at this point in the history
  • Loading branch information
asalmgren committed Nov 19, 2024
1 parent a4b7411 commit 8199a6f
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 19 deletions.
16 changes: 13 additions & 3 deletions Source/Utils/ERF_TerrainMetrics.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -63,9 +63,6 @@ init_zlevels (Vector<Vector<Real>>& zlevels_stag,
stretched_dz_h[lev][k] = (zlevels_stag[lev][k+1] - zlevels_stag[lev][k]);
}
}

stretched_dz_d[lev].resize(domain.length(2));
Gpu::copy(Gpu::hostToDevice, stretched_dz_h[lev].begin(), stretched_dz_h[lev].end(), stretched_dz_d[lev].begin());
}

// Try reading in terrain_z_levels, which allows arbitrarily spaced grid
Expand Down Expand Up @@ -97,6 +94,19 @@ init_zlevels (Vector<Vector<Real>>& zlevels_stag,
int rr = ref_ratio[lev-1][2];
expand_and_interpolate_1d(zlevels_stag[lev], zlevels_stag[lev-1], rr, false);
}

for (int lev = 0; lev <= max_level; lev++) {
int nz = zlevels_stag[lev].size();
for (int k = 0; k < nz-1; k++)
{
stretched_dz_h[lev][k] = (zlevels_stag[lev][k+1] - zlevels_stag[lev][k]);
}
}
}

for (int lev = 0; lev <= max_level; lev++) {
stretched_dz_d[lev].resize(stretched_dz_h[lev].size());
Gpu::copy(Gpu::hostToDevice, stretched_dz_h[lev].begin(), stretched_dz_h[lev].end(), stretched_dz_d[lev].begin());
}
}

Expand Down
38 changes: 22 additions & 16 deletions Source/Utils/ERF_solve_with_fft.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -53,13 +53,10 @@ void ERF::solve_with_fft (int lev, MultiFab& rhs, MultiFab& phi, Array<MultiFab,
AMREX_ALWAYS_ASSERT(use_fft);

bool l_use_terrain = SolverChoice::terrain_type != TerrainType::None;
bool use_gmres = (l_use_terrain && !SolverChoice::terrain_is_flat);

auto const dom_lo = lbound(geom[lev].Domain());
auto const dom_hi = ubound(geom[lev].Domain());

MultiFab r_hse(base_state[lev], make_alias, BaseState::r0_comp, 1);

auto bclo = get_projection_bc(Orientation::low);
auto bchi = get_projection_bc(Orientation::high);

Expand Down Expand Up @@ -107,10 +104,7 @@ void ERF::solve_with_fft (int lev, MultiFab& rhs, MultiFab& phi, Array<MultiFab,
if (!m_2D_poisson) {
m_2D_poisson = std::make_unique<FFT::PoissonHybrid<MultiFab>>(Geom(0));
}
// m_2D_poisson->solve(phi, rhs[lev], stretched_dz_d[lev]);

Gpu::DeviceVector<Real> stretched_dz(dom_hi.z+1, geom[lev].CellSize(2));
m_2D_poisson->solve(phi, rhs, stretched_dz);
m_2D_poisson->solve(phi, rhs, stretched_dz_d[lev]);

} else {
amrex::Abort("FFT isn't appropriate for spatially varying terrain");
Expand Down Expand Up @@ -227,16 +221,28 @@ void ERF::solve_with_fft (int lev, MultiFab& rhs, MultiFab& phi, Array<MultiFab,
});

Box const& zbx = mfi.nodaltilebox(2);
const Real dz_inv = dxInv[2];
Array4<Real> const& fz_arr = fluxes[2].array(mfi);
ParallelFor(zbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
if (k == dom_lo.z || k == dom_hi.z+1) {
fz_arr(i,j,k) = 0.0;
} else {
fz_arr(i,j,k) = -(p_arr(i,j,k) - p_arr(i,j,k-1)) * dz_inv;
}
});
if (l_use_terrain && SolverChoice::terrain_is_flat) {
ParallelFor(zbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
if (k == dom_lo.z || k == dom_hi.z+1) {
fz_arr(i,j,k) = 0.0;
} else {
Real dz = 0.5 * (stretched_dz_d[lev][k] + stretched_dz_d[lev][k-1]);
fz_arr(i,j,k) = -(p_arr(i,j,k) - p_arr(i,j,k-1)) / dz;
}
});
} else {
const Real dz_inv = dxInv[2];
ParallelFor(zbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
if (k == dom_lo.z || k == dom_hi.z+1) {
fz_arr(i,j,k) = 0.0;
} else {
fz_arr(i,j,k) = -(p_arr(i,j,k) - p_arr(i,j,k-1)) * dz_inv;
}
});
}
} // mfi
}
#endif

0 comments on commit 8199a6f

Please sign in to comment.