Skip to content

Commit

Permalink
feat: move to RK2. add tracers to advection pgen
Browse files Browse the repository at this point in the history
  • Loading branch information
@AstroBarker
AstroBarker committed Jan 21, 2024
1 parent e6aa218 commit a0eaca3
Show file tree
Hide file tree
Showing 3 changed files with 115 additions and 32 deletions.
57 changes: 57 additions & 0 deletions src/pgen/advection.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
"Problem \"advection\" requires \"Minkowski\" geometry!");

auto &rc = pmb->meshblock_data.Get();
auto tracer_pkg = pmb->packages.Get("tracers");
bool do_tracers = tracer_pkg->Param<bool>("active");

PackIndexMap imap;
auto v =
Expand Down Expand Up @@ -102,6 +104,61 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
});

fluid::PrimitiveToConserved(rc.get());

/* tracer init section */
if (do_tracers) {
auto &sc = pmb->swarm_data.Get();
auto &swarm = pmb->swarm_data.Get()->Get("tracers");
auto rng_pool = tracer_pkg->Param<RNGPool>("rng_pool");

const Real &x_min = pmb->coords.Xf<1>(ib.s);
const Real &y_min = pmb->coords.Xf<2>(jb.s);
const Real &z_min = pmb->coords.Xf<3>(kb.s);
const Real &x_max = pmb->coords.Xf<1>(ib.e + 1);
const Real &y_max = pmb->coords.Xf<2>(jb.e + 1);
const Real &z_max = pmb->coords.Xf<3>(kb.e + 1);

// as for num_tracers on each block... will get too many on multiple blocks
// TODO: distribute amongst blocks.
const auto num_tracers_total = tracer_pkg->Param<int>("num_tracers");
const int number_block = num_tracers_total;

ParArrayND<int> new_indices;
swarm->AddEmptyParticles(number_block, new_indices);

auto &x = swarm->Get<Real>("x").Get();
auto &y = swarm->Get<Real>("y").Get();
auto &z = swarm->Get<Real>("z").Get();
auto &id = swarm->Get<int>("id").Get();

auto swarm_d = swarm->GetDeviceContext();

const int gid = pmb->gid;
const int max_active_index = swarm->GetMaxActiveIndex();
std::printf("ma act %d %d %d %d\n", max_active_index, num_tracers_total, number_block,
nblocks);
pmb->par_for(
"ProblemGenerator::Torus::DistributeTracers", 0, max_active_index,
KOKKOS_LAMBDA(const int n) {
if (swarm_d.IsActive(n)) {
auto rng_gen = rng_pool.get_state();

// sample in ye ball
Real r2 = 1.0 + rin * rin; // init > rin^2
while (r2 > rin * rin) {
x(n) = x_min + rng_gen.drand() * (x_max - x_min);
y(n) = y_min + rng_gen.drand() * (y_max - y_min);
z(n) = z_min + rng_gen.drand() * (z_max - z_min);
r2 = x(n) * x(n) + y(n) * y(n) + z(n) * z(n);
}
id(n) = num_tracers_total * gid + n;

bool on_current_mesh_block = true;
swarm_d.GetNeighborBlockIndex(n, x(n), y(n), z(n), on_current_mesh_block);
rng_pool.free_state(rng_gen);
}
});
}
}

} // namespace advection
46 changes: 16 additions & 30 deletions src/tracers/tracers.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
#include "geometry/geometry.hpp"
#include "geometry/geometry_utils.hpp"
#include "phoebus_utils/cell_locations.hpp"
#include "phoebus_utils/phoebus_interpolation.hpp"
#include "phoebus_utils/relativity_utils.hpp"
#include "phoebus_utils/variables.hpp"

Expand Down Expand Up @@ -79,7 +78,6 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
} // Initialize

TaskStatus AdvectTracers(MeshBlockData<Real> *rc, const Real dt) {
using namespace LCInterp;
namespace p = fluid_prim;

auto *pmb = rc->GetParentPointer();
Expand All @@ -104,41 +102,29 @@ TaskStatus AdvectTracers(MeshBlockData<Real> *rc, const Real dt) {

auto geom = Geometry::GetCoordinateSystem(rc);

// update loop.
// update loop. RK2
const int max_active_index = swarm->GetMaxActiveIndex();
pmb->par_for(
"Advect Tracers", 0, max_active_index, KOKKOS_LAMBDA(const int n) {
if (swarm_d.IsActive(n)) {
int k, j, i;
swarm_d.Xtoijk(x(n), y(n), z(n), i, j, k);

// geom quantities
Real gcov4[4][4];
geom.SpacetimeMetric(0.0, x(n), y(n), z(n), gcov4);
Real lapse = geom.Lapse(0.0, x(n), y(n), z(n));
Real shift[3];
geom.ContravariantShift(0.0, x(n), y(n), z(n), shift);

// Get shift, W, lapse
const Real Wvel_X1 = LCInterp::Do(0, x(n), y(n), z(n), pack, pvel_lo);
const Real Wvel_X2 = LCInterp::Do(0, x(n), y(n), z(n), pack, pvel_lo + 1);
const Real Wvel_X3 = LCInterp::Do(0, x(n), y(n), z(n), pack, pvel_hi);
const Real Wvel[] = {Wvel_X1, Wvel_X2, Wvel_X3};
const Real W = phoebus::GetLorentzFactor(Wvel, gcov4);
const Real vel_X1 = Wvel_X1 / W;
const Real vel_X2 = Wvel_X2 / W;
const Real vel_X3 = Wvel_X3 / W;

const Real rhs1 = (lapse * vel_X1 - shift[0]) * dt;
const Real rhs2 = (lapse * vel_X2 - shift[1]) * dt;
Real rhs3 = 0.0;
if (ndim > 2) {
rhs3 = (lapse * vel_X3 - shift[2]) * dt;
}

x(n) += rhs1;
y(n) += rhs2;
z(n) += rhs3;
Real rhs1, rhs2, rhs3;

// predictor
tracers_rhs(pack, geom, pvel_lo, pvel_hi, ndim, dt, x(n), y(n), z(n), rhs1,
rhs2, rhs3);
const Real kx = x(n) + 0.5 * dt * rhs1;
const Real ky = y(n) + 0.5 * dt * rhs2;
const Real kz = z(n) + 0.5 * dt * rhs3;

// corrector
tracers_rhs(pack, geom, pvel_lo, pvel_hi, ndim, dt, kx, ky, kz, rhs1, rhs2,
rhs3);
x(n) += rhs1 * dt;
y(n) += rhs2 * dt;
z(n) += rhs3 * dt;

bool on_current_mesh_block = true;
swarm_d.GetNeighborBlockIndex(n, x(n), y(n), z(n), on_current_mesh_block);
Expand Down
44 changes: 42 additions & 2 deletions src/tracers/tracers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,21 +21,61 @@
#include <parthenon/driver.hpp>
#include <parthenon/package.hpp>

#include "geometry/geometry.hpp"
#include "geometry/geometry_utils.hpp"
#include "microphysics/eos_phoebus/eos_phoebus.hpp"
#include "phoebus_utils/cell_locations.hpp"
#include "phoebus_utils/phoebus_interpolation.hpp"
#include "phoebus_utils/relativity_utils.hpp"
#include "phoebus_utils/variables.hpp"

using namespace parthenon::driver::prelude;
using namespace parthenon::package::prelude;
using namespace parthenon;
using Microphysics::EOS::EOS;

#include "phoebus_utils/variables.hpp"

typedef Kokkos::Random_XorShift64_Pool<> RNGPool;

namespace tracers {

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

/**
* RHS of tracer advection equations
* alpha v^i - beta^i
* dt not included
**/
template <typename Pack, typename Geometry>
KOKKOS_INLINE_FUNCTION void tracers_rhs(Pack &pack, Geometry &geom, const int pvel_lo,
const int pvel_hi, const int ndim, const Real dt,
const Real x, const Real y, const Real z,
Real &rhs1, Real &rhs2, Real &rhs3) {

// geom quantities
Real gcov4[4][4];
geom.SpacetimeMetric(0.0, x, y, z, gcov4);
Real lapse = geom.Lapse(0.0, x, y, z);
Real shift[3];
geom.ContravariantShift(0.0, x, y, z, shift);

// Get shift, W, lapse
const Real Wvel_X1 = LCInterp::Do(0, x, y, z, pack, pvel_lo);
const Real Wvel_X2 = LCInterp::Do(0, x, y, z, pack, pvel_lo + 1);
const Real Wvel_X3 = LCInterp::Do(0, x, y, z, pack, pvel_hi);
const Real Wvel[] = {Wvel_X1, Wvel_X2, Wvel_X3};
const Real W = phoebus::GetLorentzFactor(Wvel, gcov4);
const Real vel_X1 = Wvel_X1 / W;
const Real vel_X2 = Wvel_X2 / W;
const Real vel_X3 = Wvel_X3 / W;

rhs1 = (lapse * vel_X1 - shift[0]);
rhs2 = (lapse * vel_X2 - shift[1]);
rhs3 = 0.0;
if (ndim > 2) {
rhs3 = (lapse * vel_X3 - shift[2]);
}
}

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

void FillTracers(MeshBlockData<Real> *rc);
Expand Down

0 comments on commit a0eaca3

Please sign in to comment.