Skip to content

Commit

Permalink
add a distance function (#2990)
Browse files Browse the repository at this point in the history 
Adds a separate function to calculate distance, mainly because the old way gets the 2d spherical case wrong.
  • Loading branch information
@zhichen3
zhichen3 authored Nov 13, 2024
1 parent b38a018 commit 161c875
Show file tree
Hide file tree
Showing 6 changed files with 69 additions and 45 deletions.
17 changes: 9 additions & 8 deletions Source/driver/Castro.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Real, 3> loc = {0.0};

loc[0] = problo[0] + (static_cast<Real>(i) + 0.5_rt) * dx[0];
loc[0] = problo[0] + (static_cast<Real>(i) + 0.5_rt) * dx[0] - problem::center[0];
#if AMREX_SPACEDIM >= 2
loc[1] = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1];
loc[1] = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1] - problem::center[1];
#endif
#if AMREX_SPACEDIM == 3
loc[2] = problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[2];
loc[2] = problo[2] + (static_cast<Real>(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;
Expand Down Expand Up @@ -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;
}

}
Expand Down
22 changes: 22 additions & 0 deletions Source/driver/Castro_util.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Real, 3>& 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
{

Expand Down
49 changes: 22 additions & 27 deletions Source/driver/Derive.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Real>(i) + 0.5_rt) * dx[0] - problem::center[0];
GpuArray<Real, 3> loc = {0.0};
loc[0] = problo[0] + (static_cast<Real>(i) + 0.5_rt) * dx[0] - problem::center[0];
#if AMREX_SPACEDIM >= 2
Real y = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1] - problem::center[1];
#else
Real y = 0.0_rt;
loc[1] = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1] - problem::center[1];
#endif
#if AMREX_SPACEDIM == 3
Real z = problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[2] - problem::center[2];
#else
Real z = 0.0_rt;
loc[2] = problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[2] - problem::center[2];
#endif

if (domain_is_plane_parallel) {
Expand All @@ -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 );
}

});
Expand All @@ -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<Real>(i) + 0.5_rt) * dx[0] - problem::center[0];
GpuArray<Real, 3> loc = {0.0};
loc[0] = problo[0] + (static_cast<Real>(i) + 0.5_rt) * dx[0] - problem::center[0];
#if AMREX_SPACEDIM >= 2
Real y = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1] - problem::center[1];
#else
Real y = 0.0_rt;
loc[1] = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1] - problem::center[1];
#endif
#if AMREX_SPACEDIM == 3
Real z = problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[2] - problem::center[2];
#else
Real z = 0.0_rt;
loc[2] = problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[2] - problem::center[2];
#endif

if (domain_is_plane_parallel) {
Expand All @@ -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
Expand All @@ -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));
}
Expand Down
13 changes: 8 additions & 5 deletions Source/gravity/Gravity.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1334,6 +1334,7 @@ Gravity::interpolate_monopole_grav(int level, RealVector& radial_grav, MultiFab&
const Real dr = dx[0] / static_cast<Real>(gravity::drdxfac);

const auto problo = geom.ProbLoArray();
const auto geomdata = geom.data();

#ifdef _OPENMP
#pragma omp parallel
Expand Down Expand Up @@ -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<int>(r / dr);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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<Real>(i) + 0.5_rt) * dx[0] - problem::center[0];
GpuArray<Real, 3> loc;
loc[0] = problo[0] + (static_cast<Real>(i) + 0.5_rt) * dx[0] - problem::center[0];
Real lo_i = problo[0] + static_cast<Real>(i) * dx[0] - problem::center[0];

Real yc = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1] - problem::center[1];
loc[1]= problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1] - problem::center[1];
Real lo_j = problo[1] + static_cast<Real>(j) * dx[1] - problem::center[1];

Real zc = problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[2] - problem::center[2];
loc[2]= problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[2] - problem::center[2];
Real lo_k = problo[2] + static_cast<Real>(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<int>(r * drinv);

// We may be coming in here with a masked out zone (in a zone on a coarse
Expand Down
10 changes: 6 additions & 4 deletions Source/reactions/Castro_react.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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<Real, 3> rr = {0.0_rt};

rr[0] = problo[0] + dx[0] * (static_cast<Real>(i) + 0.5_rt) - problem::center[0];
#if AMREX_SPACEDIM >= 2
Expand All @@ -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);
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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<Real, 3> rr = {0.0_rt};

rr[0] = problo[0] + dx[0] * (static_cast<Real>(i) + 0.5_rt) - problem::center[0];
#if AMREX_SPACEDIM >= 2
Expand All @@ -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);
Expand Down
3 changes: 2 additions & 1 deletion Source/sources/Castro_sponge.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 161c875

Please sign in to comment.