Skip to content

Commit

Permalink
refactored ApplyRadiationFourForce and CoolingFunctionCalculateFourForce
Browse files Browse the repository at this point in the history
  • Loading branch information
mari2895 committed Jan 8, 2024
1 parent 5ec5d55 commit f57529a
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 84 deletions.
10 changes: 5 additions & 5 deletions src/phoebus_driver.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -720,17 +720,17 @@ TaskListStatus PhoebusDriver::RadiationPostStep() {
pdo_gain_reducer);
int reg_dep_id = 0;
}
auto calculate_four_force =
tl.AddTask(finish_gain_reducer, radiation::CoolingFunctionCalculateFourForce,
sc0.get(), dt);
auto apply_four_force = tl.AddTask(
calculate_four_force, radiation::ApplyRadiationFourForce, sc0.get(), dt);
}
TaskRegion &async_region = tc.AddRegion(num_independent_task_lists);
for (int ib = 0; ib < num_independent_task_lists; ib++) {
auto pmb = blocks[ib].get();
auto &tl = async_region[ib];
auto &sc0 = pmb->meshblock_data.Get(stage_name[integrator->nstages]);
auto calculate_four_force =
tl.AddTask(finish_gain_reducer, radiation::CoolingFunctionCalculateFourForce,
sc0.get(), dt);
auto apply_four_force = tl.AddTask(
calculate_four_force, radiation::ApplyRadiationFourForce, sc0.get(), dt);
}
} else if (rad_method == "monte_carlo") {
return MonteCarloStep();
Expand Down
114 changes: 58 additions & 56 deletions src/radiation/cooling_function.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -116,38 +116,30 @@ TaskStatus CheckDoGain(MeshData<Real> *rc, bool *do_gain_global) {
return TaskStatus::complete;
}

TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const double dt) {
TaskStatus CoolingFunctionCalculateFourForce(MeshData<Real> *rc, const double dt) {
namespace p = fluid_prim;
namespace c = fluid_cons;
namespace iv = internal_variables;
auto *pmb = rc->GetParentPointer();
using parthenon::MakePackDescriptor;
Mesh *pmesh = rc->GetMeshPointer();
auto &resolved_pkgs = pmesh->resolved_packages;
const int ndim = pmesh->ndim;

std::vector<std::string> vars({c::density::name(), p::density::name(),
p::velocity::name(), p::temperature::name(),
p::ye::name(), c::energy::name(), iv::Gcov::name(),
iv::GcovHeat::name(), iv::GcovCool::name(),
iv::Gye::name(), iv::tau::name(), p::energy::name()});
static auto desc =
MakePackDescriptor<c::density, p::density, p::velocity, p::temperature, p::ye,
c::energy, iv::Gcov, iv::GcovHeat, iv::GcovCool, iv::Gye,
iv::tau, p::energy>(resolved_pkgs.get());
const int nblocks = v.GetNBlocks();

PackIndexMap imap;
auto v = rc->PackVariables(vars, imap);
const int crho = imap[c::density::name()].first;
const int prho = imap[p::density::name()].first;
const int pvlo = imap[p::velocity::name()].first;
const int pvhi = imap[p::velocity::name()].second;
const int ptemp = imap[p::temperature::name()].first;
const int pye = imap[p::ye::name()].first;
const int penergy = imap[p::energy::name()].first;
const int Gcov_lo = imap[iv::Gcov::name()].first;
const int Gcov_hi = imap[iv::Gcov::name()].second;
const int Gye = imap[iv::Gye::name()].first;
const int ptau = imap[iv::tau::name()].first;
const int GcovHeat = imap[iv::GcovHeat::name()].first;
const int GcovCool = imap[iv::GcovCool::name()].first;
auto v = desc.GetPack(rc);

IndexRange ib = rc->GetBoundsI(IndexDomain::interior);
IndexRange jb = rc->GetBoundsJ(IndexDomain::interior);
IndexRange kb = rc->GetBoundsK(IndexDomain::interior);

const int nblocks = v.GetNBlocks();

auto &unit_conv =
pmb->packages.Get("phoebus")->Param<phoebus::UnitConversions>("unit_conv");
auto rad = pmb->packages.Get("radiation").get();
Expand Down Expand Up @@ -176,17 +168,17 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const doub
const Real time_conversion_factor = unit_conv.GetTimeCGSToCode();

parthenon::par_for(
DEFAULT_LOOP_PATTERN, "CoolingFunctionCalculateFourForce", DevExecSpace(), kb.s,
kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) {
DEFAULT_LOOP_PATTERN, "CoolingFunctionCalculateFourForce", DevExecSpace(), 0,
nblocks - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
// Initialize five-force to zero
for (int mu = Gcov_lo; mu <= Gcov_lo + 3; mu++) {
v(mu, k, j, i) = 0.;
for (int mu = iv::Gcov(0); mu <= iv::Gcov(0) + 3; mu++) {
v(b, mu, k, j, i) = 0.;
}
v(Gye, k, j, i) = 0.;
v(b, iv::Gye(), k, j, i) = 0.;
});

// Light Bulb with Liebendorfer model
auto &coords = pmb->coords;
const bool do_liebendorfer = rad->Param<bool>("do_liebendorfer");
const bool do_lightbulb = rad->Param<bool>("do_lightbulb");
if (do_lightbulb) {
Expand All @@ -200,16 +192,19 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const doub
const bool do_gain = pdo_gain_reducer->val;

parthenon::par_for(
DEFAULT_LOOP_PATTERN, "CoolingFunctionCalculateFourForce", DevExecSpace(), kb.s,
kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int k, const int j, const int i) {
const Real r = std::abs(coords.Xc<1>(k, j, i)); // TODO(MG) coord transform game
const Real rho = v(prho, k, j, i) * density_conversion_factor; // Density in CGS
const Real cdensity = v(crho, k, j, i); // conserved density
DEFAULT_LOOP_PATTERN, "CoolingFunctionCalculateFourForce", DevExecSpace(), 0,
nblocks - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
auto &coords = v.GetCoordinates(b) const Real r =
std::abs(coords.Xc<1>(k, j, i));
const Real rho =
v(b, p::density(), k, j, i) * density_conversion_factor; // Density in CGS
const Real cdensity = v(b, c::density(), k, j, i); // conserved density
Real Gcov[4][4];
geom.SpacetimeMetric(CellLocation::Cent, k, j, i, Gcov);
Real Ucon[4];
Real vel[3] = {v(pvlo, k, j, i), v(pvlo + 1, k, j, i), v(pvlo + 2, k, j, i)};
Real vel[3] = {v(b, p::velocity(0), k, j, i), v(b, p::velocity(0) + 1, k, j, i),
v(b, p::velocity(0) + 2, k, j, i)};
GetFourVelocity(vel, geom, CellLocation::Cent, k, j, i, Ucon);
Geometry::Tetrads Tetrads(Ucon, Gcov);
Real Jye = 0.0;
Expand All @@ -227,7 +222,7 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const doub
constexpr Real MeVToCGS = 1.16040892301e10;
constexpr Real Tnorm = 2.0 * MeVToCGS;

Real Ye = v(pye, k, j, i);
Real Ye = v(b, p::ye(), k, j, i);

if (do_liebendorfer) {
constexpr Real Ye_beta = 0.27;
Expand Down Expand Up @@ -257,19 +252,20 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const doub
}
Real heat;
Real cool;
const Real tau = v(ptau, k, j, i);
const Real tau = v(b, iv::tau(), k, j, i);
const Real hfac = LightBulb::HeatAndCool::HFAC * lum;
const Real cfac = LightBulb::HeatAndCool::CFAC;
Real Xa, Xh, Xn, Xp, Abar, Zbar;
Real lambda[2];
lambda[0] = Ye;
eos_sc.MassFractionsFromDensityTemperature(
rho, v(ptemp, k, j, i) * temperature_conversion_factor, Xa, Xh, Xn, Xp,
Abar, Zbar, lambda);
rho, v(b, p::temperature(), k, j, i) * temperature_conversion_factor, Xa,
Xh, Xn, Xp, Abar, Zbar, lambda);
heat = do_gain * (Xn + Xp) * hfac * std::exp(-tau) *
pow((rnorm / (r * length_conversion_factor)), 2);
cool = (Xn + Xp) * cfac * std::exp(-tau) *
pow((v(ptemp, k, j, i) * temperature_conversion_factor / Tnorm), 6);
cool =
(Xn + Xp) * cfac * std::exp(-tau) *
pow((v(b, p::temp(), k, j, i) * temperature_conversion_factor / Tnorm), 6);

Real CGSToCodeFact =
energy_conversion_factor / mass_conversion_factor / time_conversion_factor;
Expand All @@ -283,11 +279,13 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const doub
Tetrads.TetradToCoordCov(Gcov_tetrad, Gcov_coord);
for (int mu = Gcov_lo; mu <= Gcov_lo + 3; mu++) {
// detg included above
Kokkos::atomic_add(&(v(mu, k, j, i)), -Gcov_coord[mu - Gcov_lo]);
Kokkos::atomic_add(&(v(b, mu, k, j, i)), -Gcov_coord[mu - iv::Gcov(0)]);
}
v(GcovHeat, k, j, i) = v(prho, k, j, i) * density_conversion_factor * heat;
v(GcovCool, k, j, i) = v(prho, k, j, i) * density_conversion_factor * cool;
Kokkos::atomic_add(&(v(Gye, k, j, i)), Jye);
v(GcovHeat, k, j, i) =
v(b, p::density(), k, j, i) * density_conversion_factor * heat;
v(GcovCool, k, j, i) =
v(b, p::density(), k, j, i) * density_conversion_factor * cool;
Kokkos::atomic_add(&(v(b, iv::Gye(), k, j, i)), Jye);
});
#else
PARTHENON_THROW("Lighbulb only supported with HDF5");
Expand All @@ -299,32 +297,36 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const doub
auto s = species[sidx];

parthenon::par_for(
DEFAULT_LOOP_PATTERN, "CoolingFunctionCalculateFourForce", DevExecSpace(),
kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int k, const int j, const int i) {
DEFAULT_LOOP_PATTERN, "CoolingFunctionCalculateFourForce", DevExecSpace(), 0,
nblocks - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
Real Gcov[4][4];
geom.SpacetimeMetric(CellLocation::Cent, k, j, i, Gcov);
Real Ucon[4];
Real vel[3] = {v(pvlo, k, j, i), v(pvlo + 1, k, j, i),
v(pvlo + 2, k, j, i)};
Real vel[3] = {v(b, p::velocity(0), k, j, i),
v(b, p::velocity(0) + 1, k, j, i),
v(b, p::velocity(0) + 2, k, j, i)};
GetFourVelocity(vel, geom, CellLocation::Cent, k, j, i, Ucon);
Geometry::Tetrads Tetrads(Ucon, Gcov);

const Real Ye = v(pye, k, j, i);
const Real Ye = v(b, p::ye(), k, j, i);

double J = d_opacity.Emissivity(v(prho, k, j, i), v(ptemp, k, j, i), Ye, s);
double Jye = mp_code * d_opacity.NumberEmissivity(v(prho, k, j, i),
v(ptemp, k, j, i), Ye, s);
double J = d_opacity.Emissivity(v(b, p::density(), k, j, i),
v(b, p::temperature(), k, j, i), Ye, s);
double Jye = mp_code * d_opacity.NumberEmissivity(
v(b, p::density(), k, j, i),
v(b, p::temperature(), k, j, i), Ye, s);

Real Gcov_tetrad[4] = {-J, 0., 0., 0.};
Real Gcov_coord[4];
Tetrads.TetradToCoordCov(Gcov_tetrad, Gcov_coord);
Real detG = geom.DetG(CellLocation::Cent, k, j, i);

for (int mu = Gcov_lo; mu <= Gcov_lo + 3; mu++) {
Kokkos::atomic_add(&(v(mu, k, j, i)), -detG * Gcov_coord[mu - Gcov_lo]);
for (int mu = iv::Gcov(0); mu <= iv::Gcov(0) + 3; mu++) {
Kokkos::atomic_add(&(v(b, mu, k, j, i)),
-detG * Gcov_coord[mu - Gcov_lo]);
}
Kokkos::atomic_add(&(v(Gye, k, j, i)), -LeptonSign(s) * detG * Jye);
Kokkos::atomic_add(&(v(b, iv::Gye, k, j, i)), -LeptonSign(s) * detG * Jye);
});
}
}
Expand Down
41 changes: 20 additions & 21 deletions src/radiation/radiation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -538,37 +538,36 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
return physics;
}

TaskStatus ApplyRadiationFourForce(MeshBlockData<Real> *rc, const double dt) {
TaskStatus ApplyRadiationFourForce(MeshData<Real> *rc, const double dt) {
PARTHENON_REQUIRE(USE_VALENCIA, "Covariant GRMHD formulation not supported!");

namespace c = fluid_cons;
namespace iv = internal_variables;
using parthenon::MakePackDescriptor;
Mesh *pmesh = rc->GetMeshPointer();
auto &resolved_pkgs = pmesh->resolved_packages;
const int ndim = pmesh->ndim;

static auto desc =
MakePackDescriptor<c::density, c::energy, c::momentum, c::ye, iv::Gcov, iv::Gye>(
resolved_pkgs.get());

std::vector<std::string> vars({c::density::name(), c::energy::name(),
c::momentum::name(), c::ye::name(), iv::Gcov::name(),
iv::Gye::name()});
PackIndexMap imap;
auto v = rc->PackVariables(vars, imap);
const int crho = imap[c::density::name()].first;
const int ceng = imap[c::energy::name()].first;
const int cmom_lo = imap[c::momentum::name()].first;
const int cmom_hi = imap[c::momentum::name()].second;
const int cye = imap[c::ye::name()].first;
const int Gcov_lo = imap[iv::Gcov::name()].first;
const int Gcov_hi = imap[iv::Gcov::name()].second;
const int Gye = imap[iv::Gye::name()].first;
auto v = desc.GetPack(rc);
const int nblocks = v.GetNBlocks();

IndexRange ib = rc->GetBoundsI(IndexDomain::interior);
IndexRange jb = rc->GetBoundsJ(IndexDomain::interior);
IndexRange kb = rc->GetBoundsK(IndexDomain::interior);

parthenon::par_for(
DEFAULT_LOOP_PATTERN, "ApplyRadiationFourForce", DevExecSpace(), kb.s, kb.e, jb.s,
jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) {
v(ceng, k, j, i) -= v(Gcov_lo, k, j, i) * dt;
v(cmom_lo, k, j, i) += v(Gcov_lo + 1, k, j, i) * dt;
v(cmom_lo + 1, k, j, i) += v(Gcov_lo + 2, k, j, i) * dt;
v(cmom_lo + 2, k, j, i) += v(Gcov_lo + 3, k, j, i) * dt;
v(cye, k, j, i) += v(Gye, k, j, i) * dt;
DEFAULT_LOOP_PATTERN, "ApplyRadiationFourForce", DevExecSpace(), 0, nblocks - 1,
kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
v(ceng, k, j, i) -= v(b, iv::Gcov(0), k, j, i) * dt;
v(cmom_lo, k, j, i) += v(b, iv::Gcov(0) + 1, k, j, i) * dt;
v(cmom_lo + 1, k, j, i) += v(b, iv::Gcov(0) + 2, k, j, i) * dt;
v(cmom_lo + 2, k, j, i) += v(b, iv::Gcov(0) + 3, k, j, i) * dt;
v(cye, k, j, i) += v(b, iv::Gye(), k, j, i) * dt;
});
return TaskStatus::complete;
}
Expand Down
4 changes: 2 additions & 2 deletions src/radiation/radiation.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -91,12 +91,12 @@ typedef Kokkos::Random_XorShift64_Pool<> RNGPool;

std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin);

TaskStatus ApplyRadiationFourForce(MeshBlockData<Real> *rc, const Real dt);
TaskStatus ApplyRadiationFourForce(MeshData<Real> *rc, const Real dt);

Real EstimateTimestepBlock(MeshBlockData<Real> *rc);

// Cooling function tasks
TaskStatus CoolingFunctionCalculateFourForce(MeshBlockData<Real> *rc, const Real dt);
TaskStatus CoolingFunctionCalculateFourForce(MeshData<Real> *rc, const Real dt);
TaskStatus CheckDoGain(MeshData<Real> *rc, bool *do_gain_global);
TaskStatus LightBulbCalcTau(MeshData<Real> *rc);

Expand Down

0 comments on commit f57529a

Please sign in to comment.