From 161c875ae969b327d0aef545145278f8a61ba1e0 Mon Sep 17 00:00:00 2001 From: Zhi Chen <62574124+zhichen3@users.noreply.github.com> Date: Wed, 13 Nov 2024 11:56:18 -0500 Subject: [PATCH] add a distance function (#2990) Adds a separate function to calculate distance, mainly because the old way gets the 2d spherical case wrong. --- Source/driver/Castro.cpp | 17 ++++++----- Source/driver/Castro_util.H | 22 ++++++++++++++ Source/driver/Derive.cpp | 49 ++++++++++++++----------------- Source/gravity/Gravity.cpp | 13 ++++---- Source/reactions/Castro_react.cpp | 10 ++++--- Source/sources/Castro_sponge.cpp | 3 +- 6 files changed, 69 insertions(+), 45 deletions(-) diff --git a/Source/driver/Castro.cpp b/Source/driver/Castro.cpp index 6149e7f988..2a47ac55dd 100644 --- a/Source/driver/Castro.cpp +++ b/Source/driver/Castro.cpp @@ -3641,19 +3641,17 @@ Castro::apply_tagging_restrictions(TagBoxArray& tags, [[maybe_unused]] Real time const Real* probhi = geomdata.ProbHi(); const Real* dx = geomdata.CellSize(); - Real loc[3] = {0.0}; + GpuArray loc = {0.0}; - loc[0] = problo[0] + (static_cast(i) + 0.5_rt) * dx[0]; + loc[0] = problo[0] + (static_cast(i) + 0.5_rt) * dx[0] - problem::center[0]; #if AMREX_SPACEDIM >= 2 - loc[1] = problo[1] + (static_cast(j) + 0.5_rt) * dx[1]; + loc[1] = problo[1] + (static_cast(j) + 0.5_rt) * dx[1] - problem::center[1]; #endif #if AMREX_SPACEDIM == 3 - loc[2] = problo[2] + (static_cast(k) + 0.5_rt) * dx[2]; + loc[2] = problo[2] + (static_cast(k) + 0.5_rt) * dx[2] - problem::center[2]; #endif - Real r = std::sqrt((loc[0] - problem::center[0]) * (loc[0] - problem::center[0]) + - (loc[1] - problem::center[1]) * (loc[1] - problem::center[1]) + - (loc[2] - problem::center[2]) * (loc[2] - problem::center[2])); + Real r = distance(geomdata, loc); Real max_dist_lo = 0.0; Real max_dist_hi = 0.0; @@ -4357,9 +4355,12 @@ Castro::define_new_center(const MultiFab& S, Real time) // Now broadcast to everyone else. ParallelDescriptor::Bcast(&problem::center[0], AMREX_SPACEDIM, owner); - // Make sure if R-Z that center stays exactly on axis + // Make sure if R-Z and SPHERICAL that center stays exactly on axis if ( Geom().IsRZ() ) { problem::center[0] = 0; + } else if ( Geom().IsSPHERICAL() ) { + problem::center[0] = 0; + problem::center[1] = 0; } } diff --git a/Source/driver/Castro_util.H b/Source/driver/Castro_util.H index 016215aa32..c6162b0190 100644 --- a/Source/driver/Castro_util.H +++ b/Source/driver/Castro_util.H @@ -156,6 +156,28 @@ void position(int i, int j, int k, } + +AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE +Real distance(GeometryData const& geomdata, GpuArray& loc) +{ + // Returns the distance from the center provided with loc, the position. + // loc is the form of [x, y, z,] in Cartesian, [r, z, phi] in cylindrical + // and [r, theta, phi] in spherical + + const auto coord = geomdata.Coord(); + + if (coord == CoordSys::cartesian) { + return std::sqrt(loc[0]*loc[0] + loc[1]*loc[1] + loc[2]*loc[2]); + } + + if (coord == CoordSys::RZ) { + return std::sqrt(loc[0]*loc[0] + loc[1]*loc[1]); + } + + return std::abs(loc[0]); +} + + namespace geometry_util { diff --git a/Source/driver/Derive.cpp b/Source/driver/Derive.cpp index 63d31449d1..78a1dfcc35 100644 --- a/Source/driver/Derive.cpp +++ b/Source/driver/Derive.cpp @@ -579,21 +579,18 @@ extern "C" auto dx = geom.CellSizeArray(); auto problo = geom.ProbLoArray(); + auto geomdata = geom.data(); amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept { - - Real x = problo[0] + (static_cast(i) + 0.5_rt) * dx[0] - problem::center[0]; + GpuArray loc = {0.0}; + loc[0] = problo[0] + (static_cast(i) + 0.5_rt) * dx[0] - problem::center[0]; #if AMREX_SPACEDIM >= 2 - Real y = problo[1] + (static_cast(j) + 0.5_rt) * dx[1] - problem::center[1]; -#else - Real y = 0.0_rt; + loc[1] = problo[1] + (static_cast(j) + 0.5_rt) * dx[1] - problem::center[1]; #endif #if AMREX_SPACEDIM == 3 - Real z = problo[2] + (static_cast(k) + 0.5_rt) * dx[2] - problem::center[2]; -#else - Real z = 0.0_rt; + loc[2] = problo[2] + (static_cast(k) + 0.5_rt) * dx[2] - problem::center[2]; #endif if (domain_is_plane_parallel) { @@ -607,15 +604,15 @@ extern "C" // where e_r and e_phi are the cylindrical unit vectors // we need the distance in the x-y plane from the origin - Real r = std::sqrt(x*x + y*y); - der(i,j,k,0) = (dat(i,j,k,1)*x + dat(i,j,k,2)*y) / (dat(i,j,k,0)*r); + Real r = std::sqrt(loc[0]*loc[0] + loc[1]*loc[1]); + der(i,j,k,0) = (dat(i,j,k,1)*loc[0] + dat(i,j,k,2)*loc[1]) / (dat(i,j,k,0)*r); #endif } else { - Real r = std::sqrt(x*x + y*y + z*z); + Real r = distance(geomdata, loc); - der(i,j,k,0) = (dat(i,j,k,1)*x + - dat(i,j,k,2)*y + - dat(i,j,k,3)*z) / ( dat(i,j,k,0)*r ); + der(i,j,k,0) = (dat(i,j,k,1)*loc[0] + + dat(i,j,k,2)*loc[1] + + dat(i,j,k,3)*loc[2]) / ( dat(i,j,k,0)*r ); } }); @@ -634,21 +631,19 @@ extern "C" auto dx = geom.CellSizeArray(); auto problo = geom.ProbLoArray(); + auto geomdata = geom.data(); amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept { - Real x = problo[0] + (static_cast(i) + 0.5_rt) * dx[0] - problem::center[0]; + GpuArray loc = {0.0}; + loc[0] = problo[0] + (static_cast(i) + 0.5_rt) * dx[0] - problem::center[0]; #if AMREX_SPACEDIM >= 2 - Real y = problo[1] + (static_cast(j) + 0.5_rt) * dx[1] - problem::center[1]; -#else - Real y = 0.0_rt; + loc[1] = problo[1] + (static_cast(j) + 0.5_rt) * dx[1] - problem::center[1]; #endif #if AMREX_SPACEDIM == 3 - Real z = problo[2] + (static_cast(k) + 0.5_rt) * dx[2] - problem::center[2]; -#else - Real z = 0.0_rt; + loc[2] = problo[2] + (static_cast(k) + 0.5_rt) * dx[2] - problem::center[2]; #endif if (domain_is_plane_parallel) { @@ -662,11 +657,11 @@ extern "C" // where e_r and e_phi are the cylindrical unit vectors // we need the distance in the x-y plane from the origin - Real r = std::sqrt(x*x + y*y); - der(i,j,k,0) = (-dat(i,j,k,1)*y + dat(i,j,k,2)*x) / (dat(i,j,k,0)*r); + Real r = std::sqrt(loc[0]*loc[0] + loc[1]*loc[1]); + der(i,j,k,0) = (-dat(i,j,k,1)*loc[1] + dat(i,j,k,2)*loc[0]) / (dat(i,j,k,0)*r); #endif } else { - Real r = std::sqrt(x*x + y*y + z*z); + Real r = distance(geomdata, loc); // we really mean just the velocity component that is // perpendicular to radial, and in general 3-d (e.g. a @@ -676,9 +671,9 @@ extern "C" dat(i,j,k,2)*dat(i,j,k,2) + dat(i,j,k,3)*dat(i,j,k,3))/(dat(i,j,k,0)*dat(i,j,k,0)); - Real vr = (dat(i,j,k,1)*x + - dat(i,j,k,2)*y + - dat(i,j,k,3)*z) / ( dat(i,j,k,0)*r ); + Real vr = (dat(i,j,k,1)*loc[0] + + dat(i,j,k,2)*loc[1] + + dat(i,j,k,3)*loc[2]) / ( dat(i,j,k,0)*r ); der(i,j,k,0) = std::sqrt(amrex::max(vtot2 - vr*vr, 0.0_rt)); } diff --git a/Source/gravity/Gravity.cpp b/Source/gravity/Gravity.cpp index 82003be86f..cfd40f65f0 100644 --- a/Source/gravity/Gravity.cpp +++ b/Source/gravity/Gravity.cpp @@ -1334,6 +1334,7 @@ Gravity::interpolate_monopole_grav(int level, RealVector& radial_grav, MultiFab& const Real dr = dx[0] / static_cast(gravity::drdxfac); const auto problo = geom.ProbLoArray(); + const auto geomdata = geom.data(); #ifdef _OPENMP #pragma omp parallel @@ -1367,7 +1368,7 @@ Gravity::interpolate_monopole_grav(int level, RealVector& radial_grav, MultiFab& loc[2] = 0.0_rt; #endif - Real r = std::sqrt(loc[0] * loc[0] + loc[1] * loc[1] + loc[2] * loc[2]); + Real r = distance(geomdata, loc); int index = static_cast(r / dr); @@ -1457,6 +1458,7 @@ Gravity::compute_radial_mass(const Box& bx, Real drinv = 1.0_rt / dr; const int coord_type = geom.Coord(); + const auto geomdata = geom.data(); AMREX_ALWAYS_ASSERT(coord_type >= 0 && coord_type <= 2); @@ -1497,16 +1499,17 @@ Gravity::compute_radial_mass(const Box& bx, amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept { - Real xc = problo[0] + (static_cast(i) + 0.5_rt) * dx[0] - problem::center[0]; + GpuArray loc; + loc[0] = problo[0] + (static_cast(i) + 0.5_rt) * dx[0] - problem::center[0]; Real lo_i = problo[0] + static_cast(i) * dx[0] - problem::center[0]; - Real yc = problo[1] + (static_cast(j) + 0.5_rt) * dx[1] - problem::center[1]; + loc[1]= problo[1] + (static_cast(j) + 0.5_rt) * dx[1] - problem::center[1]; Real lo_j = problo[1] + static_cast(j) * dx[1] - problem::center[1]; - Real zc = problo[2] + (static_cast(k) + 0.5_rt) * dx[2] - problem::center[2]; + loc[2]= problo[2] + (static_cast(k) + 0.5_rt) * dx[2] - problem::center[2]; Real lo_k = problo[2] + static_cast(k) * dx[2] - problem::center[2]; - Real r = std::sqrt(xc * xc + yc * yc + zc * zc); + Real r = distance(geomdata, loc); int index = static_cast(r * drinv); // We may be coming in here with a masked out zone (in a zone on a coarse diff --git a/Source/reactions/Castro_react.cpp b/Source/reactions/Castro_react.cpp index b7a422f1dd..9927bc4f2a 100644 --- a/Source/reactions/Castro_react.cpp +++ b/Source/reactions/Castro_react.cpp @@ -209,6 +209,7 @@ Castro::react_state(MultiFab& s, MultiFab& r, Real time, Real dt, const int stra const auto dx = geom.CellSizeArray(); #ifdef MODEL_PARSER const auto problo = geom.ProbLoArray(); + const auto geomdata = geom.data(); #endif #if defined(AMREX_USE_GPU) @@ -270,7 +271,7 @@ Castro::react_state(MultiFab& s, MultiFab& r, Real time, Real dt, const int stra #ifdef MODEL_PARSER if (drive_initial_convection) { - Real rr[3] = {0.0_rt}; + GpuArray rr = {0.0_rt}; rr[0] = problo[0] + dx[0] * (static_cast(i) + 0.5_rt) - problem::center[0]; #if AMREX_SPACEDIM >= 2 @@ -285,7 +286,7 @@ Castro::react_state(MultiFab& s, MultiFab& r, Real time, Real dt, const int stra if (domain_is_plane_parallel) { dist = rr[AMREX_SPACEDIM-1]; } else { - dist = std::sqrt(rr[0] * rr[0] + rr[1] * rr[1] + rr[2] * rr[2]); + dist = distance(geomdata, rr); } burn_state.T_fixed = interpolate(dist, model::itemp); @@ -565,6 +566,7 @@ Castro::react_state(Real time, Real dt) const auto dx = geom.CellSizeArray(); const auto problo = geom.ProbLoArray(); + const auto geomdata = geom.data(); #if defined(AMREX_USE_GPU) ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) @@ -618,7 +620,7 @@ Castro::react_state(Real time, Real dt) #ifdef MODEL_PARSER if (drive_initial_convection) { - Real rr[3] = {0.0_rt}; + GpuArray rr = {0.0_rt}; rr[0] = problo[0] + dx[0] * (static_cast(i) + 0.5_rt) - problem::center[0]; #if AMREX_SPACEDIM >= 2 @@ -633,7 +635,7 @@ Castro::react_state(Real time, Real dt) if (domain_is_plane_parallel) { dist = rr[AMREX_SPACEDIM-1]; } else { - dist = std::sqrt(rr[0] * rr[0] + rr[1] * rr[1] + rr[2] * rr[2]); + dist = distance(geomdata, rr); } burn_state.T_fixed = interpolate(dist, model::itemp); diff --git a/Source/sources/Castro_sponge.cpp b/Source/sources/Castro_sponge.cpp index 2109d63374..06781f4403 100644 --- a/Source/sources/Castro_sponge.cpp +++ b/Source/sources/Castro_sponge.cpp @@ -81,6 +81,7 @@ Castro::apply_sponge(const Box& bx, auto dx = geom.CellSizeArray(); auto problo = geom.ProbLoArray(); + auto geomdata = geom.data(); amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept @@ -124,7 +125,7 @@ Castro::apply_sponge(const Box& bx, Real sponge_factor = 0.0_rt; if (sponge_lower_radius >= 0.0_rt && sponge_upper_radius > sponge_lower_radius) { - Real rad = std::sqrt(r[0]*r[0] + r[1]*r[1] + r[2]*r[2]); + Real rad = distance(geomdata, r); if (rad < sponge_lower_radius) { sponge_factor = sponge_lower_factor;