/* PBHVlasov example for GRTeclyn (pbhgr project, stage 2). See VlasovParticles.hpp. */
#include "VlasovParticles.hpp"
#include "DimensionDefinitions.hpp"
#include "TensorAlgebra.hpp"
#include <AMReX_MultiFabUtil.H>
#include <cmath>
#include <cstdint>
using namespace VlasovIdx;
namespace
{
// metric quantities from the interpolated field block f[NFLD] (values + derivatives)
struct MetricAtParticle
{
amrex::Real alpha, chi;
amrex::Real beta[3], dalpha[3], dchi[3];
amrex::Real h[3][3], dh[3][3][3], dbeta[3][3]; // dbeta[i][j] = d_i beta^j, dh[i][j][k] = d_i h_jk
amrex::Real gUU[3][3]; // gamma^{ij} = chi h~^{ij}
amrex::Real dgUU[3][3][3]; // d_i gamma^{jk}
AMREX_GPU_DEVICE static MetricAtParticle from(const amrex::Real *f)
{
MetricAtParticle M;
auto val = [&](int field, int d) { return f[4 * field + d]; }; // d = 0 value, 1..3 derivative
M.alpha = val(f_lapse, 0);
M.chi = val(f_chi, 0);
static constexpr int sym[3][3] = {{0, 1, 2}, {1, 3, 4}, {2, 4, 5}};
for (int i = 0; i < 3; ++i)
{
M.beta[i] = val(f_shift + i, 0);
M.dalpha[i] = val(f_lapse, 1 + i);
M.dchi[i] = val(f_chi, 1 + i);
for (int j = 0; j < 3; ++j)
{
M.dbeta[i][j] = val(f_shift + j, 1 + i);
M.h[i][j] = val(f_h + sym[i][j], 0);
for (int k = 0; k < 3; ++k)
M.dh[k][i][j] = val(f_h + sym[i][j], 1 + k);
}
}
// inverse of the symmetric conformal metric (det h~ = 1 up to numerical error; use the full inverse)
const amrex::Real det = M.h[0][0] * (M.h[1][1] * M.h[2][2] - M.h[1][2] * M.h[2][1]) -
M.h[0][1] * (M.h[1][0] * M.h[2][2] - M.h[1][2] * M.h[2][0]) +
M.h[0][2] * (M.h[1][0] * M.h[2][1] - M.h[1][1] * M.h[2][0]);
amrex::Real hUU[3][3];
hUU[0][0] = (M.h[1][1] * M.h[2][2] - M.h[1][2] * M.h[2][1]) / det;
hUU[0][1] = (M.h[0][2] * M.h[2][1] - M.h[0][1] * M.h[2][2]) / det;
hUU[0][2] = (M.h[0][1] * M.h[1][2] - M.h[0][2] * M.h[1][1]) / det;
hUU[1][1] = (M.h[0][0] * M.h[2][2] - M.h[0][2] * M.h[2][0]) / det;
hUU[1][2] = (M.h[0][2] * M.h[1][0] - M.h[0][0] * M.h[1][2]) / det;
hUU[2][2] = (M.h[0][0] * M.h[1][1] - M.h[0][1] * M.h[1][0]) / det;
hUU[1][0] = hUU[0][1]; hUU[2][0] = hUU[0][2]; hUU[2][1] = hUU[1][2];
for (int i = 0; i < 3; ++i)
for (int j = 0; j < 3; ++j)
{
M.gUU[i][j] = M.chi * hUU[i][j];
for (int k = 0; k < 3; ++k)
{
// d_k h~^{ij} = - h~^{il} h~^{jm} d_k h~_lm
amrex::Real dhUU = 0.0;
for (int l = 0; l < 3; ++l)
for (int m = 0; m < 3; ++m)
dhUU -= hUU[i][l] * hUU[j][m] * M.dh[k][l][m];
M.dgUU[k][i][j] = M.dchi[k] * hUU[i][j] + M.chi * dhUU;
}
}
return M;
}
AMREX_GPU_DEVICE amrex::Real Gamma(const amrex::Real *u) const
{
amrex::Real s = 1.0;
for (int i = 0; i < 3; ++i)
for (int j = 0; j < 3; ++j)
s += gUU[i][j] * u[i] * u[j];
return std::sqrt(s);
}
AMREX_GPU_DEVICE void rhs(const amrex::Real *u, amrex::Real *dx, amrex::Real *du) const
{
const amrex::Real G = Gamma(u);
for (int i = 0; i < 3; ++i)
{
amrex::Real gu = 0.0;
for (int j = 0; j < 3; ++j)
gu += gUU[i][j] * u[j];
dx[i] = alpha * gu / G - beta[i];
amrex::Real t = -G * dalpha[i]; // -Gamma d_i alpha (H = alpha Gamma - beta^j u_j)
for (int j = 0; j < 3; ++j)
{
t += u[j] * dbeta[i][j];
for (int k = 0; k < 3; ++k)
t -= 0.5 * alpha / G * u[j] * u[k] * dgUU[i][j][k];
}
du[i] = t;
}
}
};
} // namespace
namespace
{
// Deterministic hash-based random numbers (independent of the MPI layout)
inline std::uint64_t splitmix64(std::uint64_t &s)
{
s += 0x9E3779B97F4A7C15ULL;
std::uint64_t z = s;
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
return z ^ (z >> 31);
}
inline double hash_uniform(std::uint64_t &s) { return (splitmix64(s) >> 11) * (1.0 / 9007199254740992.0); }
inline double hash_gauss(std::uint64_t &s)
{
double u1 = hash_uniform(s);
const double u2 = hash_uniform(s);
if (u1 < 1e-300)
u1 = 1e-300;
return std::sqrt(-2.0 * std::log(u1)) * std::cos(2.0 * M_PI * u2);
}
// Unit-variance isotropic Gaussian velocity for sub-particle `sub` of the lattice cell `key`. For 2^3 particles per
// cell the eight velocities are the corners of a randomly rotated cube scaled by a Maxwellian speed: every particle is
// marginally N(0, 1) per component, while each cell has zero net momentum and an isotropic second moment (quiet start).
// Other lattices fall back to independent Gaussians.
inline void unit_velocity(std::uint64_t key, int sub, int n_per_dir, std::uint64_t seed, amrex::Real *v)
{
std::uint64_t s = key ^ (seed * 0xD1B54A32D192ED03ULL);
splitmix64(s);
if (n_per_dir == 2)
{
double q[4], nq = 0.0;
for (int i = 0; i < 4; ++i)
{
q[i] = hash_gauss(s);
nq += q[i] * q[i];
}
nq = std::sqrt(nq);
for (int i = 0; i < 4; ++i)
q[i] /= nq;
const double g1 = hash_gauss(s), g2 = hash_gauss(s), g3 = hash_gauss(s);
const double speed = std::sqrt((g1 * g1 + g2 * g2 + g3 * g3) / 3.0); // times the corner (+-1, +-1, +-1)
const double c[3] = {(sub & 1) ? 1.0 : -1.0, (sub & 2) ? 1.0 : -1.0, (sub & 4) ? 1.0 : -1.0};
// rotation matrix of the unit quaternion (w, x, y, z) = (q0, q1, q2, q3)
const double w = q[0], x = q[1], y = q[2], z = q[3];
const double Rm[3][3] = {{1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)},
{2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)},
{2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)}};
for (int i = 0; i < 3; ++i)
v[i] = speed * (Rm[i][0] * c[0] + Rm[i][1] * c[1] + Rm[i][2] * c[2]);
}
else
{
s ^= static_cast<std::uint64_t>(sub + 1) * 0x2545F4914F6CDD1DULL;
splitmix64(s);
for (int i = 0; i < 3; ++i)
v[i] = hash_gauss(s);
}
}
// Particle lattice built level by level in each level's own index space: n_per_dir^3 particles per level-l cell
// inside the level-l grids grown by buffer_fine cells, excluding the (grown) region of level l+1. Cells are
// visited through the level-0 tiles owned by this rank (refined to level l), so creation is distributed;
// the particles are added at level 0 and redistributed to their levels.
template <class F>
void make_lattice(VlasovParticles &pc, const amrex::Geometry &geom0, int n_per_dir, int buffer_fine, F &&fill)
{
using ParticleType = VlasovParticles::ParticleType;
const int finest = pc.finestLevel();
const auto plo = geom0.ProbLoArray();
const auto dx0 = geom0.CellSizeArray();
// cumulative refinement ratio from level 0 to level l
std::vector<int> ratio(finest + 1, 1);
for (int l = 1; l <= finest; ++l)
ratio[l] = ratio[l - 1] * pc.GetParGDB()->refRatio(l - 1)[0];
// grown grids per level (level 0: the whole domain), and the region of each level = grown grids minus the
// grown grids of the next level (coarsened to this level)
std::vector<amrex::BoxArray> grown(finest + 1);
for (int l = 0; l <= finest; ++l)
{
if (l == 0)
grown[l] = amrex::BoxArray(geom0.Domain());
else
{
grown[l] = pc.ParticleBoxArray(l);
grown[l].grow(buffer_fine);
grown[l] = amrex::intersect(grown[l], amrex::refine(geom0.Domain(), ratio[l]));
grown[l].removeOverlap(); // grown boxes overlap each other: make the list disjoint
}
}
std::vector<amrex::BoxList> region(finest + 1);
for (int l = 0; l <= finest; ++l)
{
if (l == finest)
region[l] = grown[l].boxList();
else
{
amrex::BoxList next_bl = grown[l + 1].boxList();
next_bl.coarsen(pc.GetParGDB()->refRatio(l)[0]);
amrex::BoxArray next(next_bl); // plain array (coarsened lists cannot call removeOverlap)
next.removeOverlap();
for (int b = 0; b < grown[l].size(); ++b)
region[l].join(next.complementIn(grown[l][b]));
}
}
for (amrex::MFIter mfi = pc.MakeMFIter(0); mfi.isValid(); ++mfi)
{
const amrex::Box &tile0 = mfi.tilebox();
amrex::Gpu::HostVector<ParticleType> host;
for (int l = 0; l <= finest; ++l)
{
const amrex::Box tile_l = amrex::refine(tile0, ratio[l]);
const int n = n_per_dir;
const amrex::Real dxl[3] = {dx0[0] / ratio[l], dx0[1] / ratio[l], dx0[2] / ratio[l]};
const amrex::Real dV_sub = dxl[0] * dxl[1] * dxl[2] / (static_cast<amrex::Real>(n) * n * n);
for (const amrex::Box &rb : region[l])
{
const amrex::Box cells = tile_l & rb;
if (cells.isEmpty())
continue;
for (amrex::IntVect iv = cells.smallEnd(); iv <= cells.bigEnd(); cells.next(iv))
for (int a = 0; a < n; ++a)
for (int b = 0; b < n; ++b)
for (int c = 0; c < n; ++c)
{
ParticleType p;
p.id() = ParticleType::NextID();
p.cpu() = amrex::ParallelDescriptor::MyProc();
p.pos(0) = plo[0] + (iv[0] + (a + 0.5) / n) * dxl[0];
p.pos(1) = plo[1] + (iv[1] + (b + 0.5) / n) * dxl[1];
p.pos(2) = plo[2] + (iv[2] + (c + 0.5) / n) * dxl[2];
for (int q = 0; q < VlasovIdx::NAOS; ++q)
p.rdata(q) = 0.0;
// per-cell key (level + cell index) and sub-lattice index for balanced velocity sets
const std::uint64_t key = (static_cast<std::uint64_t>(l) << 58) ^
(static_cast<std::uint64_t>(iv[0]) * 0x9E3779B97F4A7C15ULL) ^
(static_cast<std::uint64_t>(iv[1]) * 0xC2B2AE3D27D4EB4FULL) ^
(static_cast<std::uint64_t>(iv[2]) * 0x165667B19E3779F9ULL);
fill(p, dV_sub, key, a + n * (b + n * c));
host.push_back(p);
}
}
}
auto &ptile = pc.GetParticles(0)[std::make_pair(mfi.index(), mfi.LocalTileIndex())];
auto old_size = ptile.numParticles();
ptile.resize(old_size + host.size());
amrex::Gpu::copy(amrex::Gpu::hostToDevice, host.begin(), host.end(), ptile.GetArrayOfStructs().begin() + old_size);
}
pc.Redistribute();
}
} // namespace
void VlasovParticles::init_uniform(const amrex::Geometry &geom, amrex::Real rho, int n_per_dir)
{
make_lattice(*this, geom, n_per_dir, m_lattice_buffer,
[=](ParticleType &p, amrex::Real dV_sub, std::uint64_t /*key*/, int /*sub*/) {
p.rdata(m) = rho * dV_sub;
p.rdata(Gamma) = 1.0;
p.rdata(tau) = m_tau0;
});
}
void VlasovParticles::init_lattice(const amrex::Geometry &geom, int n_per_dir)
{
make_lattice(*this, geom, n_per_dir, m_lattice_buffer,
[=](ParticleType &p, amrex::Real dV_sub, std::uint64_t /*key*/, int /*sub*/) {
p.rdata(m) = dV_sub;
p.rdata(Gamma) = 1.0;
p.rdata(tau) = m_tau0;
});
}
void VlasovParticles::init_from_table(const amrex::Geometry &geom, const std::vector<amrex::Real> &r_tab,
const std::vector<amrex::Real> &rho_tab, const std::vector<amrex::Real> &ur_tab,
const std::vector<amrex::Real> &chi_tab, int n_per_dir,
const std::vector<amrex::Real> &sr_tab, const std::vector<amrex::Real> &st_tab,
amrex::Real sigma, std::uint64_t seed)
{
const bool warm = sigma > 0.0 && sr_tab.size() == r_tab.size() && st_tab.size() == r_tab.size();
const auto plo = geom.ProbLoArray();
const auto phi = geom.ProbHiArray();
const amrex::Real cx = 0.5 * (plo[0] + phi[0]), cy = 0.5 * (plo[1] + phi[1]), cz = 0.5 * (plo[2] + phi[2]);
auto interp = [&](const std::vector<amrex::Real> &tab, amrex::Real r) {
if (r >= r_tab.back()) return tab.back();
const amrex::Real s = r / (r_tab[1] - r_tab[0]);
const int i = std::min(static_cast<int>(s), static_cast<int>(r_tab.size()) - 2);
const amrex::Real t = s - i;
return (1 - t) * tab[i] + t * tab[i + 1];
};
make_lattice(*this, geom, n_per_dir, m_lattice_buffer,
[&](ParticleType &p, amrex::Real dV_sub, std::uint64_t key, int sub) {
const amrex::Real X = p.pos(0) - cx, Y = p.pos(1) - cy, Z = p.pos(2) - cz;
const amrex::Real r = std::sqrt(X * X + Y * Y + Z * Z);
const amrex::Real rho = interp(rho_tab, r), chi = interp(chi_tab, r), ur = interp(ur_tab, r);
// momentum per unit mass in the orthonormal frame of the normal observer (gamma_ij = psi^4 delta_ij):
// bulk radial u_r plus, for warm particles, the thermal part with the local dispersions sigma_r, sigma_t
// (non-relativistic addition, as in the 1D code)
amrex::Real pv[3] = {0.0, 0.0, 0.0};
if (r > 0.0)
{
const amrex::Real nn[3] = {X / r, Y / r, Z / r};
for (int i = 0; i < 3; ++i)
pv[i] = ur * nn[i];
if (warm)
{
const amrex::Real sr = sigma * interp(sr_tab, r), st = sigma * interp(st_tab, r);
if (sr > 0.0 || st > 0.0)
{
amrex::Real v[3];
unit_velocity(key, sub, n_per_dir, seed, v);
const amrex::Real vr = v[0] * nn[0] + v[1] * nn[1] + v[2] * nn[2];
for (int i = 0; i < 3; ++i)
pv[i] += sr * vr * nn[i] + st * (v[i] - vr * nn[i]);
}
}
}
// rho is the normal-observer energy density = Gamma * rest-mass density, so the particle rest mass carries
// 1/Gamma of its own Gamma and the deposited energy density equals the constraint-satisfying rho;
// covariant components u_i = psi^2 p_i
const amrex::Real sqrtg = std::pow(chi, -1.5), psi2 = 1.0 / std::sqrt(chi);
const amrex::Real G = std::sqrt(1.0 + pv[0] * pv[0] + pv[1] * pv[1] + pv[2] * pv[2]);
p.rdata(m) = rho / G * sqrtg * dV_sub;
p.rdata(u1) = psi2 * pv[0];
p.rdata(u2) = psi2 * pv[1];
p.rdata(u3) = psi2 * pv[2];
p.rdata(Gamma) = G;
p.rdata(tau) = m_tau0;
});
}
namespace
{
// CIC deposit functor: comp 0 <- dep0 = m Gamma, comp k <- dep0 * dep_k (k = 1..9)
template <class PC>
struct DepositCIC
{
amrex::GpuArray<amrex::Real, 3> plo, dxi;
AMREX_GPU_DEVICE void operator()(const typename PC::ParticleTileType::ConstParticleTileDataType &ptd, int ip,
amrex::Array4<amrex::Real> const &arr) const
{
const auto &p = ptd.m_aos[ip];
amrex::ParticleInterpolator::Linear interp(p, plo, dxi);
const amrex::Real w0 = p.rdata(VlasovIdx::dep0);
interp.ParticleToMesh(p, arr, 0, 0, VlasovIdx::NDEP,
[=] AMREX_GPU_DEVICE(const typename PC::ParticleType &part, int comp) {
return comp == 0 ? w0 : w0 * part.rdata(VlasovIdx::dep0 + comp);
});
}
};
// fill the deposit slots of the particles of one level from (m, u_i, Gamma)
void fill_deposit_slots(VlasovParticles &pc, int lev)
{
using PC = VlasovParticles::PC;
for (PC::ParIterType pti(pc, lev); pti.isValid(); ++pti)
{
auto ptd = pti.GetParticleTile().getParticleTileData();
const int np = pti.numParticles();
amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(int ip) {
auto &p = ptd.m_aos[ip];
const amrex::Real mm = p.rdata(m), G = p.rdata(Gamma);
const amrex::Real u[3] = {p.rdata(u1), p.rdata(u2), p.rdata(u3)};
p.rdata(dep0) = mm * G;
for (int i = 0; i < 3; ++i)
p.rdata(dep0 + 1 + i) = u[i] / G;
static constexpr int ii[6] = {0, 0, 0, 1, 1, 2}, jj[6] = {0, 1, 2, 1, 2, 2};
for (int s = 0; s < 6; ++s)
p.rdata(dep0 + 4 + s) = u[ii[s]] * u[jj[s]] / (G * G);
});
}
amrex::Gpu::streamSynchronize();
}
// metric fields and first derivatives on a level: NFLD components, 1 ghost cell, from a state with >= 3 ghosts
void metric_fields(const amrex::MultiFab &a_state_gh, const amrex::Geometry &geom, amrex::MultiFab &fld)
{
const int comps[NMET] = {c_lapse, c_shift1, c_shift2, c_shift3, c_chi, c_h11, c_h12, c_h13, c_h22, c_h23, c_h33};
const auto dx = geom.CellSizeArray();
const auto &st = a_state_gh.const_arrays();
const auto &fa = fld.arrays();
amrex::GpuArray<int, NMET> comp_arr;
for (int i = 0; i < NMET; ++i)
comp_arr[i] = comps[i];
amrex::ParallelFor(fld, amrex::IntVect(1), [=] AMREX_GPU_DEVICE(int box_no, int ix, int iy, int iz) {
const auto &s = st[box_no];
for (int f = 0; f < NMET; ++f)
{
const int c = comp_arr[f];
fa[box_no](ix, iy, iz, 4 * f) = s(ix, iy, iz, c);
fa[box_no](ix, iy, iz, 4 * f + 1) =
(s(ix - 2, iy, iz, c) - 8 * s(ix - 1, iy, iz, c) + 8 * s(ix + 1, iy, iz, c) - s(ix + 2, iy, iz, c)) / (12 * dx[0]);
fa[box_no](ix, iy, iz, 4 * f + 2) =
(s(ix, iy - 2, iz, c) - 8 * s(ix, iy - 1, iz, c) + 8 * s(ix, iy + 1, iz, c) - s(ix, iy + 2, iz, c)) / (12 * dx[1]);
fa[box_no](ix, iy, iz, 4 * f + 3) =
(s(ix, iy, iz - 2, c) - 8 * s(ix, iy, iz - 1, c) + 8 * s(ix, iy, iz + 1, c) - s(ix, iy, iz + 2, c)) / (12 * dx[2]);
}
});
amrex::Gpu::streamSynchronize();
}
// linear (CIC) interpolation of the NFLD fields to a particle position. The stencil is clamped to the array box:
// the field MultiFab has one ghost cell, so a particle more than half a cell outside its box (possible at the
// midpoint of a step near a box boundary) would otherwise read out of bounds. Returns 0 (stencil inside), 1
// (clamped) or -1 (non-finite position; f untouched).
template <class P>
AMREX_GPU_DEVICE int interp_fields_at(const P &p, amrex::Array4<const amrex::Real> const &arr,
const amrex::GpuArray<amrex::Real, 3> &plo,
const amrex::GpuArray<amrex::Real, 3> &dxi, amrex::Real *f, int ncomp = NFLD)
{
const auto lo = amrex::lbound(arr);
const auto hi = amrex::ubound(arr);
const int alo[3] = {lo.x, lo.y, lo.z}, ahi[3] = {hi.x, hi.y, hi.z};
int i0[3], i1[3];
amrex::Real w1[3];
int status = 0;
for (int d = 0; d < 3; ++d)
{
const amrex::Real l = (p.pos(d) - plo[d]) * dxi[d] - 0.5;
if (!(l > -1.0e15 && l < 1.0e15))
return -1;
const int i = static_cast<int>(amrex::Math::floor(l));
w1[d] = l - i;
if (i < alo[d] || i + 1 > ahi[d])
status = 1;
i0[d] = amrex::min(amrex::max(i, alo[d]), ahi[d]);
i1[d] = amrex::min(amrex::max(i + 1, alo[d]), ahi[d]);
}
const amrex::Real wx0 = 1.0 - w1[0], wy0 = 1.0 - w1[1], wz0 = 1.0 - w1[2];
for (int c = 0; c < ncomp; ++c)
f[c] = wx0 * wy0 * wz0 * arr(i0[0], i0[1], i0[2], c) + w1[0] * wy0 * wz0 * arr(i1[0], i0[1], i0[2], c) +
wx0 * w1[1] * wz0 * arr(i0[0], i1[1], i0[2], c) + w1[0] * w1[1] * wz0 * arr(i1[0], i1[1], i0[2], c) +
wx0 * wy0 * w1[2] * arr(i0[0], i0[1], i1[2], c) + w1[0] * wy0 * w1[2] * arr(i1[0], i0[1], i1[2], c) +
wx0 * w1[1] * w1[2] * arr(i0[0], i1[1], i1[2], c) + w1[0] * w1[1] * w1[2] * arr(i1[0], i1[1], i1[2], c);
return status;
}
AMREX_GPU_DEVICE inline bool finite_rhs(const amrex::Real *kx, const amrex::Real *ku)
{
for (int i = 0; i < 3; ++i)
if (!(std::isfinite(kx[i]) && std::isfinite(ku[i])))
return false;
return true;
}
} // namespace
void VlasovParticles::set_from_fields(int lev, const amrex::MultiFab &fld, const amrex::Geometry &geom)
{
const auto plo = geom.ProbLoArray();
const auto dxi = geom.InvCellSizeArray();
amrex::MeshToParticle(*this, fld, lev,
[=] AMREX_GPU_DEVICE(ParticleType &p, amrex::Array4<const amrex::Real> const &arr) {
amrex::Real f[11];
interp_fields_at(p, arr, plo, dxi, f, 11);
const amrex::Real E = f[0], chi = f[4];
const amrex::Real h[3][3] = {{f[5], f[6], f[7]}, {f[6], f[8], f[9]}, {f[7], f[9], f[10]}};
const amrex::Real det = h[0][0] * (h[1][1] * h[2][2] - h[1][2] * h[2][1]) -
h[0][1] * (h[1][0] * h[2][2] - h[1][2] * h[2][0]) +
h[0][2] * (h[1][0] * h[2][1] - h[1][1] * h[2][0]);
amrex::Real hUU[3][3];
hUU[0][0] = (h[1][1] * h[2][2] - h[1][2] * h[2][1]) / det;
hUU[0][1] = (h[0][2] * h[2][1] - h[0][1] * h[2][2]) / det;
hUU[0][2] = (h[0][1] * h[1][2] - h[0][2] * h[1][1]) / det;
hUU[1][1] = (h[0][0] * h[2][2] - h[0][2] * h[2][0]) / det;
hUU[1][2] = (h[0][2] * h[1][0] - h[0][0] * h[1][2]) / det;
hUU[2][2] = (h[0][0] * h[1][1] - h[0][1] * h[1][0]) / det;
hUU[1][0] = hUU[0][1]; hUU[2][0] = hUU[0][2]; hUU[2][1] = hUU[1][2];
const amrex::Real v[3] = {f[1] / E, f[2] / E, f[3] / E};
amrex::Real v2 = 0.0;
for (int i = 0; i < 3; ++i)
for (int j = 0; j < 3; ++j)
v2 += chi * hUU[i][j] * v[i] * v[j];
const amrex::Real G = 1.0 / std::sqrt(1.0 - v2);
const amrex::Real dV_sub = p.rdata(m);
p.rdata(m) = E * std::pow(chi, -1.5) * dV_sub / G;
p.rdata(u1) = G * v[0];
p.rdata(u2) = G * v[1];
p.rdata(u3) = G * v[2];
p.rdata(Gamma) = G;
});
amrex::Gpu::streamSynchronize();
}
void VlasovParticles::apply_correction(int lev, const amrex::MultiFab &corr, const amrex::Geometry &geom)
{
const auto plo = geom.ProbLoArray();
const auto dxi = geom.InvCellSizeArray();
amrex::MeshToParticle(*this, corr, lev,
[=] AMREX_GPU_DEVICE(ParticleType &p, amrex::Array4<const amrex::Real> const &arr) {
amrex::Real f[4];
interp_fields_at(p, arr, plo, dxi, f, 4);
p.rdata(m) *= f[0];
p.rdata(u1) += f[1];
p.rdata(u2) += f[2];
p.rdata(u3) += f[3];
});
amrex::Gpu::streamSynchronize();
}
void VlasovParticles::deposit_level(int lev, const amrex::Vector<amrex::MultiFab *> &a_states,
const amrex::Vector<amrex::Geometry> &a_geoms)
{
const int finest = this->finestLevel();
fill_deposit_slots(*this, lev);
if (lev > 0)
fill_deposit_slots(*this, lev - 1);
if (lev < finest)
fill_deposit_slots(*this, lev + 1);
// 3 ghost cells: ghost particles may sit 2 cells outside the grids and their CIC clouds reach one more
amrex::MultiFab dep(this->ParticleBoxArray(lev), this->ParticleDistributionMap(lev), NDEP, 3);
const DepositCIC<PC> f{a_geoms[lev].ProbLoArray(), a_geoms[lev].InvCellSizeArray()};
if (m_verbose_deposit) amrex::Print() << " deposit: level " << lev << " own particles\n";
amrex::ParticleToMesh(*this, dep, lev, f, true);
if (lev > 0 && m_dep_ghosts)
{
// coarse (lev-1) particles within 1 coarse cell of the level-lev grids: their inward clouds
VlasovParticles ghosts;
ghosts.Define(this->GetParGDB());
ghosts.reserveData();
ghosts.resizeData();
PC::ParticleTileType ghost_tile;
this->CreateGhostParticles(lev - 1, 1, ghost_tile);
if (m_verbose_deposit) amrex::Print() << " deposit: level " << lev << " ghosts " << ghost_tile.numParticles() << " (rank 0)\n";
ghosts.AddParticlesAtLevel(ghost_tile, lev, 2);
amrex::ParticleToMesh(ghosts, dep, lev, f, false);
}
if (lev < finest && m_dep_virtuals)
{
// copies of the level-(lev+1) particles near the boundary of the fine grids, deposited here with this
// level's kernel: the cells just outside the fine grids see the same clouds as from own particles
const int lf = lev + 1;
const amrex::BoxArray &ba = this->ParticleBoxArray(lf);
const auto plo = a_geoms[lf].ProbLoArray();
const auto dxi = a_geoms[lf].InvCellSizeArray();
const int margin = 2 * this->GetParGDB()->refRatio(lev)[0]; // one coarse CIC cloud, in fine cells
amrex::Gpu::HostVector<ParticleType> host;
for (PC::ParIterType pti(*this, lf); pti.isValid(); ++pti)
{
const auto &aos = pti.GetArrayOfStructs();
for (const auto &p : aos)
{
const amrex::IntVect iv(static_cast<int>(std::floor((p.pos(0) - plo[0]) * dxi[0])),
static_cast<int>(std::floor((p.pos(1) - plo[1]) * dxi[1])),
static_cast<int>(std::floor((p.pos(2) - plo[2]) * dxi[2])));
const amrex::Box nb(iv - amrex::IntVect(margin), iv + amrex::IntVect(margin));
if (!ba.contains(nb))
host.push_back(p);
}
}
if (m_verbose_deposit) amrex::Print() << " deposit: level " << lev << " virtual copies from level " << lf << ": " << host.size() << " (rank 0)\n";
PC::ParticleTileType copies;
copies.resize(host.size());
amrex::Gpu::copy(amrex::Gpu::hostToDevice, host.begin(), host.end(), copies.GetArrayOfStructs().begin());
VlasovParticles virt;
virt.Define(this->GetParGDB());
virt.reserveData();
virt.resizeData();
virt.AddParticlesAtLevel(copies, lev, 0);
amrex::ParticleToMesh(virt, dep, lev, f, false);
}
// divide by sqrt(gamma) dV and copy into the state (interior cells)
const auto dx = a_geoms[lev].CellSizeArray();
const amrex::Real dV = dx[0] * dx[1] * dx[2];
amrex::MultiFab &st = *a_states[lev];
amrex::MultiFab dep_on_state(st.boxArray(), st.DistributionMap(), NDEP, 0);
dep_on_state.ParallelCopy(dep, 0, 0, NDEP);
const auto &dep_arrays = dep_on_state.const_arrays();
const auto &st_arrays = st.arrays();
amrex::ParallelFor(st, [=] AMREX_GPU_DEVICE(int box_no, int ix, int iy, int iz) {
const amrex::Real chi = st_arrays[box_no](ix, iy, iz, c_chi);
const amrex::Real inv = 1.0 / dV; // coordinate densities; ParticleMatter multiplies by chi^{3/2} = 1/sqrt(gamma)
amrex::ignore_unused(chi);
for (int c = 0; c < NDEP; ++c)
st_arrays[box_no](ix, iy, iz, c_rho_p + c) = dep_arrays[box_no](ix, iy, iz, c) * inv;
});
amrex::Gpu::streamSynchronize();
}
void VlasovParticles::deposit_all(const amrex::Vector<amrex::MultiFab *> &a_states,
const amrex::Vector<amrex::Geometry> &a_geoms)
{
const int finest = this->finestLevel();
for (int lev = 0; lev <= finest; ++lev)
deposit_level(lev, a_states, a_geoms);
for (int lev = finest; lev >= 1; --lev)
amrex::average_down(*a_states[lev], *a_states[lev - 1], a_geoms[lev], a_geoms[lev - 1], c_rho_p, NDEP,
this->GetParGDB()->refRatio(lev - 1));
}
void VlasovParticles::update_gamma(int lev, const amrex::MultiFab &a_state_gh, const amrex::Geometry &geom)
{
amrex::MultiFab fld(a_state_gh.boxArray(), a_state_gh.DistributionMap(), NFLD, 1);
metric_fields(a_state_gh, geom, fld);
const auto plo = geom.ProbLoArray();
const auto dxi = geom.InvCellSizeArray();
amrex::MeshToParticle(*this, fld, lev,
[=] AMREX_GPU_DEVICE(ParticleType &p, amrex::Array4<const amrex::Real> const &arr) {
amrex::Real f[NFLD];
interp_fields_at(p, arr, plo, dxi, f);
const MetricAtParticle M = MetricAtParticle::from(f);
const amrex::Real u[3] = {p.rdata(u1), p.rdata(u2), p.rdata(u3)};
p.rdata(Gamma) = M.Gamma(u);
});
amrex::Gpu::streamSynchronize();
}
void VlasovParticles::push(int lev, amrex::Real dt, const amrex::MultiFab &a_state_gh, const amrex::Geometry &geom)
{
amrex::MultiFab fld(a_state_gh.boxArray(), a_state_gh.DistributionMap(), NFLD, 1);
metric_fields(a_state_gh, geom, fld);
const auto plo = geom.ProbLoArray();
const auto dxi = geom.InvCellSizeArray();
const amrex::Real dx0 = geom.CellSize(0);
const amrex::Real cap_full = m_max_disp_cells * dx0; // cap on the displacement per step
const amrex::Real cap_half = 0.5 * cap_full;
// diagnostics: [0] clamped stencils, [1] capped displacements, [2] frozen particles; max displacement
amrex::Gpu::DeviceVector<long> cnt(3, 0L);
amrex::Gpu::DeviceVector<amrex::Real> dmax(1, 0.0);
long *pcnt = cnt.data();
amrex::Real *pdmax = dmax.data();
// stage 1: k1 at (x^n, u^n); move to the midpoint. The state is saved first so that stage 2 can restore it.
amrex::MeshToParticle(*this, fld, lev,
[=] AMREX_GPU_DEVICE(ParticleType &p, amrex::Array4<const amrex::Real> const &arr) {
p.rdata(px0) = p.pos(0); p.rdata(py0) = p.pos(1); p.rdata(pz0) = p.pos(2);
p.rdata(u01) = p.rdata(u1); p.rdata(u02) = p.rdata(u2); p.rdata(u03) = p.rdata(u3);
amrex::Real f[NFLD];
const int st = interp_fields_at(p, arr, plo, dxi, f);
if (st == 1) amrex::Gpu::Atomic::Add(pcnt + 0, 1L);
if (st < 0) { amrex::Gpu::Atomic::Add(pcnt + 2, 1L); return; }
const MetricAtParticle M = MetricAtParticle::from(f);
const amrex::Real u[3] = {p.rdata(u1), p.rdata(u2), p.rdata(u3)};
amrex::Real kx[3], ku[3];
M.rhs(u, kx, ku);
if (!finite_rhs(kx, ku)) { amrex::Gpu::Atomic::Add(pcnt + 2, 1L); return; }
amrex::Real d[3] = {0.5 * dt * kx[0], 0.5 * dt * kx[1], 0.5 * dt * kx[2]};
const amrex::Real dn = std::sqrt(d[0] * d[0] + d[1] * d[1] + d[2] * d[2]);
amrex::Gpu::Atomic::Max(pdmax, 2.0 * dn);
if (dn > cap_half)
{
const amrex::Real sc = cap_half / dn;
for (int i = 0; i < 3; ++i) d[i] *= sc;
amrex::Gpu::Atomic::Add(pcnt + 1, 1L);
}
for (int i = 0; i < 3; ++i)
p.pos(i) += d[i];
p.rdata(u1) = u[0] + 0.5 * dt * ku[0];
p.rdata(u2) = u[1] + 0.5 * dt * ku[1];
p.rdata(u3) = u[2] + 0.5 * dt * ku[2];
});
amrex::Gpu::streamSynchronize();
// stage 2: k2 at the midpoint, full step from the saved state; Gamma at the new state. A particle whose
// fields or rhs are not finite is restored to its saved state (frozen for this step) and counted.
amrex::MeshToParticle(*this, fld, lev,
[=] AMREX_GPU_DEVICE(ParticleType &p, amrex::Array4<const amrex::Real> const &arr) {
amrex::Real f[NFLD];
const int st = interp_fields_at(p, arr, plo, dxi, f);
if (st == 1) amrex::Gpu::Atomic::Add(pcnt + 0, 1L);
bool ok = st >= 0;
MetricAtParticle M;
amrex::Real kx[3] = {0.0, 0.0, 0.0}, ku[3] = {0.0, 0.0, 0.0}, dtau = 0.0;
if (ok)
{
M = MetricAtParticle::from(f);
const amrex::Real u[3] = {p.rdata(u1), p.rdata(u2), p.rdata(u3)};
M.rhs(u, kx, ku);
dtau = dt * M.alpha / M.Gamma(u); // d tau / dt = alpha / Gamma at the midpoint
ok = finite_rhs(kx, ku) && std::isfinite(dtau);
}
if (!ok)
{
p.pos(0) = p.rdata(px0); p.pos(1) = p.rdata(py0); p.pos(2) = p.rdata(pz0);
p.rdata(u1) = p.rdata(u01); p.rdata(u2) = p.rdata(u02); p.rdata(u3) = p.rdata(u03);
amrex::Gpu::Atomic::Add(pcnt + 2, 1L);
return;
}
amrex::Real d[3] = {dt * kx[0], dt * kx[1], dt * kx[2]};
const amrex::Real dn = std::sqrt(d[0] * d[0] + d[1] * d[1] + d[2] * d[2]);
amrex::Gpu::Atomic::Max(pdmax, dn);
if (dn > cap_full)
{
const amrex::Real sc = cap_full / dn;
for (int i = 0; i < 3; ++i) d[i] *= sc;
amrex::Gpu::Atomic::Add(pcnt + 1, 1L);
}
p.pos(0) = p.rdata(px0) + d[0];
p.pos(1) = p.rdata(py0) + d[1];
p.pos(2) = p.rdata(pz0) + d[2];
p.rdata(u1) = p.rdata(u01) + dt * ku[0];
p.rdata(u2) = p.rdata(u02) + dt * ku[1];
p.rdata(u3) = p.rdata(u03) + dt * ku[2];
const amrex::Real un[3] = {p.rdata(u1), p.rdata(u2), p.rdata(u3)};
p.rdata(Gamma) = M.Gamma(un);
p.rdata(tau) += dtau;
});
amrex::Gpu::streamSynchronize();
// diagnostics, reduced over the ranks
std::vector<long> c(3, 0L);
std::vector<amrex::Real> dm(1, 0.0);
amrex::Gpu::copy(amrex::Gpu::deviceToHost, cnt.begin(), cnt.end(), c.begin());
amrex::Gpu::copy(amrex::Gpu::deviceToHost, dmax.begin(), dmax.end(), dm.begin());
amrex::ParallelDescriptor::ReduceLongSum(c.data(), 3);
amrex::ParallelDescriptor::ReduceRealMax(dm[0]);
const amrex::Real dm_cells = dm[0] / dx0;
m_push_diag.max_disp_cells = std::max(m_push_diag.max_disp_cells, static_cast<double>(dm_cells));
m_push_diag.n_clamped += c[0];
m_push_diag.n_capped += c[1];
m_push_diag.n_frozen += c[2];
if (c[0] + c[1] + c[2] > 0 || dm_cells > 0.5)
amrex::Print() << "PBHVlasov push L" << lev << ": max displacement " << dm_cells << " cells, clamped " << c[0]
<< ", capped " << c[1] << ", frozen " << c[2] << "\n";
Redistribute();
}
VlasovParticles::tau_stats_t VlasovParticles::tau_in_shell(const amrex::Real *center, amrex::Real r_lo, amrex::Real r_hi)
{
tau_stats_t s;
double sum_mt = 0.0, sum_mg = 0.0, tmin = 1e300, tmax = -1e300;
for (int l = 0; l <= this->finestLevel(); ++l)
for (PC::ParIterType pti(*this, l); pti.isValid(); ++pti)
{
const auto &aos = pti.GetArrayOfStructs();
for (const auto &p : aos)
{
const amrex::Real X = p.pos(0) - center[0], Y = p.pos(1) - center[1], Z = p.pos(2) - center[2];
const amrex::Real r = std::sqrt(X * X + Y * Y + Z * Z);
if (r < r_lo || r >= r_hi)
continue;
const double mm = p.rdata(m), t = p.rdata(tau);
++s.n;
s.mass += mm;
sum_mt += mm * t;
sum_mg += mm * p.rdata(Gamma);
tmin = std::min(tmin, t);
tmax = std::max(tmax, t);
}
}
amrex::ParallelDescriptor::ReduceLongSum(s.n);
amrex::ParallelDescriptor::ReduceRealSum(s.mass);
amrex::ParallelDescriptor::ReduceRealSum(sum_mt);
amrex::ParallelDescriptor::ReduceRealSum(sum_mg);
amrex::ParallelDescriptor::ReduceRealMin(tmin);
amrex::ParallelDescriptor::ReduceRealMax(tmax);
if (s.n > 0 && s.mass > 0.0)
{
s.tau_mean = sum_mt / s.mass;
s.gamma_mean = sum_mg / s.mass;
s.tau_min = tmin;
s.tau_max = tmax;
}
return s;
}
void VlasovParticles::tau_profile(const amrex::Real *center, amrex::Real dr, int nbins, std::vector<double> &n,
std::vector<double> &mass, std::vector<double> &mtau, std::vector<double> &mgamma)
{
n.assign(nbins, 0.0); mass.assign(nbins, 0.0); mtau.assign(nbins, 0.0); mgamma.assign(nbins, 0.0);
for (int l = 0; l <= this->finestLevel(); ++l)
for (PC::ParIterType pti(*this, l); pti.isValid(); ++pti)
{
const auto &aos = pti.GetArrayOfStructs();
for (const auto &p : aos)
{
const amrex::Real X = p.pos(0) - center[0], Y = p.pos(1) - center[1], Z = p.pos(2) - center[2];
const int k = static_cast<int>(std::sqrt(X * X + Y * Y + Z * Z) / dr);
if (k >= nbins)
continue;
const double mm = p.rdata(m);
n[k] += 1.0; mass[k] += mm; mtau[k] += mm * p.rdata(tau); mgamma[k] += mm * p.rdata(Gamma);
}
}
amrex::ParallelDescriptor::ReduceRealSum(n.data(), nbins);
amrex::ParallelDescriptor::ReduceRealSum(mass.data(), nbins);
amrex::ParallelDescriptor::ReduceRealSum(mtau.data(), nbins);
amrex::ParallelDescriptor::ReduceRealSum(mgamma.data(), nbins);
}
amrex::Real VlasovParticles::total_mass(int lev)
{
amrex::Real sum = 0.0;
(void)lev;
for (int l = 0; l <= this->finestLevel(); ++l)
for (PC::ParIterType pti(*this, l); pti.isValid(); ++pti)
{
const auto &aos = pti.GetArrayOfStructs();
for (const auto &p : aos)
sum += p.rdata(m);
}
amrex::ParallelDescriptor::ReduceRealSum(sum);
return sum;
}