Skip to content

Commit

Permalink
Merge branch '2d_spherical_dist' of github.com:zhichen3/Castro into 2…
Browse files Browse the repository at this point in the history
…d_spherical_dist
  • Loading branch information
zhichen3 committed Nov 8, 2024
2 parents b38a018 + 00c76bd commit c01d12d
Show file tree
Hide file tree
Showing 17 changed files with 94 additions and 299 deletions.
22 changes: 7 additions & 15 deletions Source/driver/Castro.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3639,21 +3639,10 @@ Castro::apply_tagging_restrictions(TagBoxArray& tags, [[maybe_unused]] Real time
{
const Real* problo = geomdata.ProbLo();
const Real* probhi = geomdata.ProbHi();
const Real* dx = geomdata.CellSize();

Real loc[3] = {0.0};

loc[0] = problo[0] + (static_cast<Real>(i) + 0.5_rt) * dx[0];
#if AMREX_SPACEDIM >= 2
loc[1] = problo[1] + (static_cast<Real>(j) + 0.5_rt) * dx[1];
#endif
#if AMREX_SPACEDIM == 3
loc[2] = problo[2] + (static_cast<Real>(k) + 0.5_rt) * dx[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]));
GpuArray<Real, 3> loc;
position(i, j, k, geomdata, loc);
Real r = distance(geomdata, loc);

Real max_dist_lo = 0.0;
Real max_dist_hi = 0.0;
Expand Down Expand Up @@ -4357,9 +4346,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
25 changes: 23 additions & 2 deletions Source/driver/Castro_util.H
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,6 @@ bool mom_flux_has_p (const int mom_dir, const int flux_dir, const int coord)
}



AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void position(int i, int j, int k,
GeometryData const& geomdata, GpuArray<Real, 3>& loc,
Expand Down Expand Up @@ -147,7 +146,7 @@ void position(int i, int j, int k,
}

for (int dir = 0; dir < AMREX_SPACEDIM; ++dir) {
loc[dir] = offset[dir] + idx[dir] * dx[dir];
loc[dir] = offset[dir] + idx[dir] * dx[dir] - problem::center[dir];
}

for (int dir = AMREX_SPACEDIM; dir < 3; ++dir) {
Expand All @@ -156,6 +155,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 loc[0];
}


namespace geometry_util
{

Expand Down
131 changes: 27 additions & 104 deletions Source/driver/Derive.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -577,24 +577,13 @@ extern "C"
auto const dat = datfab.array();
auto const der = derfab.array();

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];
#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;
#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;
#endif
GpuArray<Real, 3> loc;
position(i, j, k, geomdata, loc);

if (domain_is_plane_parallel) {
#if AMREX_SPACEDIM == 2
Expand All @@ -607,15 +596,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 @@ -632,24 +621,13 @@ extern "C"
auto const dat = datfab.array();
auto const der = derfab.array();

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];
#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;
#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;
#endif
GpuArray<Real, 3> loc;
position(i, j, k, geomdata, loc);

if (domain_is_plane_parallel) {
#if AMREX_SPACEDIM == 2
Expand All @@ -662,11 +640,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 +654,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 Expand Up @@ -712,36 +690,16 @@ extern "C"
{

int idir = 0;
auto dx = geom.CellSizeArray();
auto problo = geom.ProbLoArray();
auto geomdata = geom.data();

auto const dat = datfab.array();
auto const L = derfab.array();

amrex::ParallelFor(bx,
[=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
Real loc[3];

//loc calculated like sum_utils.cpp
//This might be equivalent and more modular: position(i, j, k, geom, loc);
loc[0] = problo[0] + (0.5_rt + i) * dx[0];

#if AMREX_SPACEDIM >= 2
loc[1] = problo[1] + (0.5_rt + j) * dx[1];
#else
loc[1] = 0.0_rt;
#endif

#if AMREX_SPACEDIM == 3
loc[2] = problo[2] + (0.5_rt + k) * dx[2];
#else
loc[2] = 0.0_rt;
#endif

for (int dir = 0; dir < AMREX_SPACEDIM; ++dir) {
loc[dir] -= problem::center[dir];
}
GpuArray<Real, 3> loc;
position(i, j, k, geomdata, loc);

// Explicitly computing only the required cross-product as in inertial_to_rotational_velocity
if (idir == 0) { // cross_product(loc, mom): ang_mom(1)->x)
Expand All @@ -764,33 +722,16 @@ extern "C"
{

int idir = 1;
auto dx = geom.CellSizeArray();
auto problo = geom.ProbLoArray();
auto geomdata = geom.data();

auto const dat = datfab.array();
auto const L = derfab.array();

amrex::ParallelFor(bx,
[=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
Real loc[3];

loc[0] = problo[0] + (0.5_rt + i) * dx[0];

#if AMREX_SPACEDIM >= 2
loc[1] = problo[1] + (0.5_rt + j) * dx[1];
#else
loc[1] = 0.0_rt;
#endif

#if AMREX_SPACEDIM == 3
loc[2] = problo[2] + (0.5_rt + k) * dx[2];
#else
loc[2] = 0.0_rt;
#endif
for (int dir = 0; dir < AMREX_SPACEDIM; ++dir) {
loc[dir] -= problem::center[dir];
}
GpuArray<Real, 3> loc;
position(i, j, k, geomdata, loc);

if (idir == 0) { // cross_product(loc, mom): ang_mom(1)->x)
L(i,j,k,0) = loc[1] * dat(i,j,k,3) - loc[2] * dat(i,j,k,2);
Expand All @@ -812,34 +753,16 @@ extern "C"
{

int idir = 2;
auto dx = geom.CellSizeArray();
auto problo = geom.ProbLoArray();
auto geomdata = geom.data();

auto const dat = datfab.array();
auto const L = derfab.array();

amrex::ParallelFor(bx,
[=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
Real loc[3];

loc[0] = problo[0] + (0.5_rt + i) * dx[0];

#if AMREX_SPACEDIM >= 2
loc[1] = problo[1] + (0.5_rt + j) * dx[1];
#else
loc[1] = 0.0_rt;
#endif

#if AMREX_SPACEDIM == 3
loc[2] = problo[2] + (0.5_rt + k) * dx[2];
#else
loc[2] = 0.0_rt;
#endif

for (int dir = 0; dir < AMREX_SPACEDIM; ++dir) {
loc[dir] -= problem::center[dir];
}
GpuArray<Real, 3> loc;
position(i, j, k, geomdata, loc);

if (idir == 0) { // cross_product(loc, mom): ang_mom(1)->x)
L(i,j,k,0) = loc[1] * dat(i,j,k,3) - loc[2] * dat(i,j,k,2);
Expand Down
4 changes: 0 additions & 4 deletions Source/driver/sum_utils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -384,10 +384,6 @@ Castro::gwstrain (Real time,
GpuArray<Real, 3> r;
position(i, j, k, geomdata, r);

for (int n = 0; n < 3; ++n) {
r[n] -= problem::center[n];
}

Real rhoInv;
if (rho(i,j,k) * maskFactor > 0.0_rt) {
rhoInv = 1.0_rt / rho(i,j,k);
Expand Down
8 changes: 1 addition & 7 deletions Source/gravity/Castro_gravity.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -375,10 +375,7 @@ void Castro::construct_old_gravity_source(MultiFab& source, MultiFab& state_in,

#ifdef HYBRID_MOMENTUM
GpuArray<Real, 3> loc;
for (int n = 0; n < 3; ++n) {
position(i, j, k, geomdata, loc);
loc[n] -= problem::center[n];
}
position(i, j, k, geomdata, loc);

GpuArray<Real, 3> hybrid_src;

Expand Down Expand Up @@ -574,9 +571,6 @@ void Castro::construct_new_gravity_source(MultiFab& source, MultiFab& state_old,
#ifdef HYBRID_MOMENTUM
GpuArray<Real, 3> loc;
position(i, j, k, geomdata, loc);
for (int n = 0; n < 3; ++n) {
loc[n] -= problem::center[n];
}

GpuArray<Real, 3> hybrid_src;

Expand Down
Loading

0 comments on commit c01d12d

Please sign in to comment.