Particles

These tutorials use AMReX particle containers: pure struct-of-arrays (SoA) particle data distributed over the same grids as mesh data, particle-mesh deposition and gather, tracer particles around an embedded boundary, sorting and filtering within tiles, neighbor lists, and a complete particle-in-cell code. All of them require AMReX built with particles (AMReX_PARTICLES, on by default); EBTracers also needs EB and the linear solvers, and ElectrostaticPIC the linear solvers. EBTracers, NeighborList and ElectrostaticPIC run in 2D and 3D only. Build only this topic with -DAMReXTutorials_TOPICS=Particles; see Getting Started.

Basics

Level:

★

Source:

Particles/Basics

Key APIs:

ParticleContainerPureSoA, SetSoACompileTimeNames, AddRealComp, DefineAndReturnParticleTile, ParticleTileData, SetParticleIDandCPU, ParIter, Redistribute, ParticleReduce, Increment, WritePlotFile, Checkpoint, Restart

A first particle code. Each particle stores its position, velocity and weight as separate arrays; the first AMREX_SPACEDIM real components are the positions. A component dist (distance traveled) is added at run time. Particles are created on the device with nppc random particles per cell, move in straight lines through a periodic domain, and are moved to their new grid and MPI rank by Redistribute. ParticleReduce computes the total weight, kinetic energy and path length in one pass. The program writes a plotfile with the particle count per cell and the particles, and a particle checkpoint that it reads back with Restart.

The test checks that the number of particles, total weight and kinetic energy are conserved, that the total path length is \(n_\text{steps}\,\Delta t \sum |v|\), and that the restarted particles match.

The components are named with an enum; the container stores all of them, including the positions and the id/cpu word, as separate arrays.

// Real components of each particle. In a pure struct-of-arrays (SoA)
// container, the first AMREX_SPACEDIM real components are the positions.
struct RealIdx
{
    enum {
        AMREX_D_DECL(x, y, z),
        AMREX_D_DECL(vx, vy, vz),
        w,      // weight, e.g., mass
        ncomps
    };
};

// All components, including position and id/cpu, are stored as separate arrays.
using PC = amrex::ParticleContainerPureSoA<RealIdx::ncomps, 0>;
using ParticleType = PC::ParticleType;

Each tile is filled directly on the device. A block of unique ids is reserved per tile, and ParallelForRNG provides a random engine.

// Reserve a block of unique particle ids on this process.
const amrex::Long id0 = ParticleType::NextID();
ParticleType::NextID(id0 + np);

// DefineAndReturnParticleTile also allocates the runtime components.
auto& tile = pc.DefineAndReturnParticleTile(lev, mfi);
tile.resize(np);
auto ptd = tile.getParticleTileData();

amrex::ParallelForRNG(np,
[=] AMREX_GPU_DEVICE (amrex::Long ip, amrex::RandomEngine const& engine) noexcept
{
    const amrex::IntVect iv = bx.atOffset(ip / nppc);
    ptd.idcpu(ip) = amrex::SetParticleIDandCPU(id0 + ip, myproc);
    for (int d = 0; d < AMREX_SPACEDIM; ++d) {
        ptd.pos(d, ip) = static_cast<amrex::ParticleReal>(
            plo[d] + (iv[d] + amrex::Random(engine)) * dx[d]);
        ptd.rdata(RealIdx::vx+d)[ip] = static_cast<amrex::ParticleReal>(
            vmax * (2.0_rt*amrex::Random(engine) - 1.0_rt));
    }
    ptd.rdata(RealIdx::w)[ip] = static_cast<amrex::ParticleReal>(
        0.5_rt + amrex::Random(engine));
    ptd.m_runtime_rdata[0][ip] = 0;
});

The push loops over the particle tiles with ParIter. Particles move less than one cell per step, so Redistribute only communicates with neighboring ranks.

for (PC::ParIterType pti(pc, lev); pti.isValid(); ++pti)
{
    const int np = pti.numParticles();
    auto ptd = pti.GetParticleTile().getParticleTileData();
    // Components can also be looked up as arrays, e.g., by index or by name.
    amrex::ParticleReal* dist = pti.GetStructOfArrays().GetRealData(dist_comp).data();

    amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) noexcept
    {
        amrex::ParticleReal v2 = 0;
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            const amrex::ParticleReal v = ptd.rdata(RealIdx::vx+d)[i];
            ptd.pos(d, i) += static_cast<amrex::ParticleReal>(v * dt);
            v2 += v*v;
        }
        dist[i] += static_cast<amrex::ParticleReal>(std::sqrt(v2) * dt);
    });
}

// Particles moved less than one cell, so only neighboring ranks need to
// communicate (local = 1). Redistribute also applies periodic wrapping.
pc.Redistribute(0, -1, 0, 1);

Run (from the top of the repository, see Getting Started):

cmake --build build --target Particles_Basics -j 8
cd build/Particles/Basics
mpiexec -n 2 ./Basics inputs

Use n_cell, nppc and nsteps for a larger problem, and particles.do_tiling=1 for better OpenMP performance on CPUs. The particles in plt00020/particles can be viewed with yt or ParaView (see Visualization, Debugging and Profiling).

ParticleMesh

Level:

★★

Source:

Particles/ParticleMesh

Key APIs:

ParticleToMesh, MeshToParticle, ParticleInterpolator::Linear, InitRandom, ReduceSum, ReduceMax, WritePlotFile

Deposit a particle quantity onto the mesh and gather a mesh field to the particles with cloud-in-cell (linear) weights. In 1D, a particle at \(x\) between the cell centers \(x_i\) and \(x_{i+1}\), with \(f = (x - x_i)/\Delta x\), gives a fraction \(1-f\) of its mass to cell \(i\) and \(f\) to cell \(i+1\); the gather uses the same weights. In 2D and 3D the weights are products of the 1D weights.

The particles and the mesh use different BoxArrays (particle_max_grid_size and max_grid_size); ParticleToMesh and MeshToParticle copy data between the two layouts. The test checks that the deposited mass equals the total particle mass and that the gather reproduces a linear field to round-off.

The deposition callable runs once per particle and adds its weighted contributions with atomic adds. Contributions to ghost cells are summed into the valid cells that own them, including across periodic boundaries.

// Deposit the mass density with cloud-in-cell (linear) weights. A
// particle touches 2^AMREX_SPACEDIM cells, so rho needs one ghost cell;
// ParticleToMesh sums ghost cell contributions into the valid cells.
amrex::MultiFab rho(mesh_ba, mesh_dm, 1, 1);
amrex::ParticleToMesh(pc, rho, 0,
    [=] AMREX_GPU_DEVICE (const ConstPTD& ptd, int i,
                          amrex::Array4<amrex::Real> const& rho_arr) noexcept
    {
        const ConstParticle p = ptd[i];
        amrex::ParticleInterpolator::Linear interp(p, plo, dxi);
        // The callable returns the value to deposit (uses atomic adds).
        interp.ParticleToMesh(p, rho_arr, 0, 0, 1,
            [=] (ConstParticle const&, int) noexcept
            {
                return ptd.rdata(RealIdx::mass)[i] * inv_cell_volume;
            });
    });

The gather needs a mesh field with filled ghost cells. The first callable reads the mesh, the second accumulates each weighted value on the particle.

// Gather the field to the particles with the same linear weights.
amrex::MeshToParticle(pc, efield, 0,
    [=] AMREX_GPU_DEVICE (const PTD& ptd, int i,
                          amrex::Array4<const amrex::Real> const& e_arr) noexcept
    {
        auto p = ptd[i];
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            ptd.rdata(RealIdx::Ex+d)[i] = 0;
        }
        amrex::ParticleInterpolator::Linear interp(p, plo, dxi);
        // The first callable reads the mesh, the second updates the
        // particle with each weighted contribution.
        interp.MeshToParticle(p, e_arr, 0, 0, AMREX_SPACEDIM,
            [=] (amrex::Array4<const amrex::Real> const& arr,
                 int ii, int jj, int kk, int comp) noexcept
            {
                return arr(ii, jj, kk, comp);
            },
            [=] (PC::ParticleType&, int comp, amrex::Real val) noexcept
            {
                ptd.rdata(RealIdx::Ex+comp)[i] += static_cast<amrex::ParticleReal>(val);
            });
    });

Run (from the top of the repository, see Getting Started):

cmake --build build --target Particles_ParticleMesh -j 8
cd build/Particles/ParticleMesh
mpiexec -n 2 ./ParticleMesh inputs

The plotfile plt00000 holds the density rho, the field, and the particles.

EBTracers

Level:

★★

Source:

Particles/EBTracers

Key APIs:

mac_interpolate, FillSignedDistance, ParticleToMesh, ParticleInterpolator::Linear, ReduceSum/ReduceMax over particles, MLEBABecLap

Tracer particles carry a scalar through a flow around an embedded boundary: each particle has a fixed weight, and depositing the weights on the mesh gives the scalar field. It is the Lagrangian counterpart of Advection, and builds the same kind of flow: a uniform velocity made to go around a sphere in a periodic domain with a MAC projection. The velocity is kept at face centers, where mac_interpolate expects it.

Each step moves the particles with the midpoint method, with the face velocity interpolated to the particles by mac_interpolate. The interpolation does not know about the EB, so a few particles near the sphere end up inside it, up to about a cell deep. They are reflected back across the EB using the signed distance s to the EB, filled on the nodes by FillSignedDistance: the depth is -s and the normal is the gradient of s. This works for any EB geometry.

// Inside the body (s < 0): move 2|s| along the normal grad(s)/|grad(s)|.
Real grad[AMREX_SPACEDIM];
const Real s = signed_distance(sda, p.x, plo, dxi, grad);
if (s < 0.0_rt) {
    Real gnorm = 0.0_rt;
    for (int d = 0; d < AMREX_SPACEDIM; ++d) { gnorm += grad[d]*grad[d]; }
    gnorm = std::sqrt(gnorm);
    if (gnorm > 0.0_rt) {
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            p.x[d] += ParticleReal(-2.0_rt * s * grad[d] / gnorm);
        }
        ptd.idata(IntIdx::reflected)[i] += 1;
    }
}

The particles move forward to stop_time and back with the velocity reversed. The program checks that the deposited weight equals the total particle weight, that no particle stays more than 0.2 cells inside the sphere (measured with the exact sphere), and that the particles that were never reflected return, on average, to within 0.001 cells of where they started.

Run:

cmake --build build --target Particles_EBTracers -j 8
cd build/Particles/EBTracers
mpiexec -n 2 ./EBTracers inputs

Back at t = 0:
  particles reflected off the EB: 210
  max depth inside the sphere after any step: 0.07412058333 cells
  distance of the other particles from their start: mean 3.902734996e-05, max 0.0415392124 cells
  relative error of the deposited weight: 2.526048763e-14
PASSED

Sorting

Level:

★★

Source:

Particles/Sorting

Key APIs:

DenseBins, getParticleCell, SortParticlesByCell, SortParticlesByBin, partitionParticles, filterParticles, transformParticles, DefineAndReturnParticleTile

Organize the particles of a tile:

1. DenseBins builds a counting sort of the particles by cell without moving them. permutationPtr() lists particle indices bin by bin and offsetsPtr()[b] is where bin b starts, so a kernel over cells can loop over the particles of each cell.

2. SortParticlesByCell and SortParticlesByBin reorder the particle data so that particles in the same cell (or group of cells) are adjacent in memory, which improves memory access in deposition kernels.

3. partitionParticles moves the particles that satisfy a predicate (here: speed above vcut) to the front of the tile.

4. filterParticles copies the particles that satisfy a predicate to another tile, and transformParticles copies particles while changing them (here: reversing the velocity).

The test checks each step: per-cell counts against ParticleContainer::Increment, sorted order, the partition, the number of filtered particles, and the reversed velocities.

Binning by cell and looping over the particles of each cell:

amrex::DenseBins<ConstPTD> bins;
for (PC::ParIterType pti(pc, lev); pti.isValid(); ++pti)
{
    // With particle tiling, a grid holds several tiles.
    const amrex::Box bx = pti.tilebox();
    const auto ptd = pti.GetParticleTile().getConstParticleTileData();
    const amrex::IntVect lo = bx.smallEnd();
    const auto domain_box = geom.Domain();
    // Map each particle to its cell, relative to the lower corner of bx.
    bins.build(pti.numParticles(), ptd, bx,
        [=] AMREX_GPU_DEVICE (const ConstPTD& p, int i) noexcept -> amrex::IntVect
        {
            return amrex::getParticleCell(p, i, plo, dxi, domain_box) - lo;
        });

    const auto* offsets = bins.offsetsPtr();
    const auto* perm = bins.permutationPtr();
    auto const& cnt = count_bins.array(pti);
    amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
    {
        const int b = static_cast<int>(bx.index(amrex::IntVect(AMREX_D_DECL(i, j, k))));
        // The particles in cell (i,j,k) are perm[offsets[b]:offsets[b+1]].
        int n = 0;
        for (int m = offsets[b]; m < offsets[b+1]; ++m) {
            const amrex::IntVect iv = amrex::getParticleCell(ptd, perm[m], plo, dxi, domain_box);
            if (iv == amrex::IntVect(AMREX_D_DECL(i, j, k))) { ++n; }
        }
        cnt(i,j,k) = amrex::Real(n);
    });
    amrex::Gpu::streamSynchronize(); // bins is reused for the next tile
}

Copying a subset of the particles, and transformed copies, into other containers. The destination tiles must be large enough.

PC slow_pc(geom, dm, ba);
PC reversed_pc(geom, dm, ba);
for (PC::ParIterType pti(pc, lev); pti.isValid(); ++pti)
{
    const auto& src = pti.GetParticleTile();
    const int np = src.numParticles();

    auto& slow = slow_pc.DefineAndReturnParticleTile(lev, pti);
    slow.resize(np);
    const int nslow = amrex::filterParticles(slow, src,
        [=] AMREX_GPU_DEVICE (const ConstPTD& ptd, int i) noexcept
        {
            return speed2(ptd, i) <= vcut2;
        });
    slow.resize(nslow);

    auto& reversed = reversed_pc.DefineAndReturnParticleTile(lev, pti);
    reversed.resize(np);
    amrex::transformParticles(reversed, src,
        [=] AMREX_GPU_HOST_DEVICE (const PTD& dst, const ConstPTD& s, int si, int di) noexcept
        {
            dst.idcpu(di) = s.idcpu(si);
            for (int d = 0; d < AMREX_SPACEDIM; ++d) {
                dst.pos(d, di) = s.pos(d, si);
                dst.rdata(RealIdx::vx+d)[di] = -s.rdata(RealIdx::vx+d)[si];
            }
        });
}

Run (from the top of the repository, see Getting Started):

cmake --build build --target Particles_Sorting -j 8
cd build/Particles/Sorting
mpiexec -n 2 ./Sorting inputs

NeighborList

Level:

★★★

Source:

Particles/NeighborList

Key APIs:

NeighborParticleContainerPureSoA, fillNeighbors, updateNeighbors, buildNeighborList, NeighborList::data, getNeighbors, Redistribute, ParticleReduce, RandomNormal

A small molecular dynamics code with short-range forces. Particles are soft spheres of diameter \(\sigma\) with the potential \(U(r) = \frac{k}{2}(\sigma - r)^2\) for \(r < \sigma\).

MDParticleContainer derives from NeighborParticleContainerPureSoA. fillNeighbors copies particles within one cell of each tile (including periodic images) into that tile, after its real particles. buildNeighborList then lists, for each real particle, the nearby real and neighbor particles that pass a CheckPair test. The list radius is one cell, larger than \(\sigma\), so the list stays valid for num_rebuild steps; in between, updateNeighbors only refreshes the positions of the neighbor copies.

The test checks that the total momentum is conserved to round-off (pair forces are equal and opposite, also across tiles and MPI ranks) and that the relative change in total energy is below 1%.

The time loop uses velocity Verlet and rebuilds the neighbor list every num_rebuild steps:

// Velocity Verlet time stepping.
for (int step = 1; step <= nsteps; ++step)
{
    pc.kickDrift(0.5_rt*dt, dt);

    if (step % num_rebuild == 0) {
        // Particles moved less than one cell since the last rebuild.
        pc.Redistribute(0, -1, 0, 1);
        pc.fillNeighbors();
        pc.buildNeighborList(check_pair);
    } else {
        // Only refresh the data of the existing neighbor copies.
        pc.updateNeighbors();
    }

    pc.computeForces(params);
    pc.kickDrift(0.5_rt*dt, 0.0_rt);

The force kernel loops over the neighbors of each real particle:

// For each real particle i, the list holds the indices of nearby
// particles in this tile, which may be real particles or neighbor copies.
const auto nbor_data = m_neighbor_list[lev].at(index).data();

amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) noexcept
{
    amrex::Real acc[AMREX_SPACEDIM] = {AMREX_D_DECL(0._rt, 0._rt, 0._rt)};
    amrex::Real pe = 0;
    for (const auto& p2 : nbor_data.getNeighbors(i))
    {
        amrex::Real dr[AMREX_SPACEDIM];
        amrex::Real r2 = 0;
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            dr[d] = ptd.pos(d, i) - p2.pos(d);
            r2 += dr[d]*dr[d];
        }
        if (r2 >= sigma*sigma) { continue; }

        // Repulsive force k (sigma - r) along the unit vector dr/r.
        const amrex::Real r = amrex::max(std::sqrt(r2), 1.e-6_rt*sigma);
        const amrex::Real overlap = sigma - r;
        const amrex::Real coef = k * overlap / r;
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            acc[d] += coef * dr[d];
        }
        // Each pair is visited from both sides; take half the energy.
        pe += 0.25_rt * k * overlap * overlap;
    }
    for (int d = 0; d < AMREX_SPACEDIM; ++d) {
        ptd.rdata(RealIdx::ax+d)[i] = static_cast<amrex::ParticleReal>(acc[d]);
    }
    ptd.rdata(RealIdx::pe)[i] = static_cast<amrex::ParticleReal>(pe);
});

Run (from the top of the repository, see Getting Started):

cmake --build build --target Particles_NeighborList -j 8
cd build/Particles/NeighborList
mpiexec -n 2 ./NeighborList inputs

Use n_cell and nsteps for a larger run. The time step is cfl*sigma/max_speed, and 2*num_rebuild*cfl*sigma must be smaller than the skin 1 - sigma.

ElectrostaticPIC

Level:

★★★ / Blueprint

Source:

Particles/ElectrostaticPIC

Key APIs:

ParticleContainerPureSoA, ParticleToMesh, MeshToParticle, ParticleInterpolator::Linear, MLPoisson, MLMG, ParReduce, Redistribute, WriteSingleLevelPlotfile, WritePlotFile

An electrostatic particle-in-cell (PIC) code for electrons moving through a fixed, neutralizing ion background in a periodic box, organized to be copied as the starting point of a PIC code. Each step:

1. Deposit the electron charge with cloud-in-cell weights (ParticleToMesh) and add the ion charge.

2. Solve \(\nabla^2\phi = -\rho/\epsilon_0\) with MLPoisson and MLMG, and compute \(E = -\nabla\phi\) at cell centers with centered differences.

3. Interpolate \(E\) to the particles (MeshToParticle).

4. Leapfrog push: \(v \mathrel{+}= (q/m) E\,\Delta t\), then \(x \mathrel{+}= v\,\Delta t\), then Redistribute.

The problem is a Langmuir (plasma) oscillation. In normalized units (\(\epsilon_0 = m = n_0 = 1\), \(q = -1\)) the plasma frequency is \(\omega_p = \sqrt{n_0 q^2/(\epsilon_0 m)} = 1\). The electrons start at rest, displaced by \(\xi_0 \sin(kx)\), and linear theory gives \(E_x = -(q n_0 \xi_0/\epsilon_0) \sin(kx) \cos(\omega_p t)\). The code measures the frequency from the zero crossings of the \(\sin(kx)\) mode of \(E_x\). It checks that the frequency and the initial amplitude agree with theory to 2%, and that the total (field plus kinetic) energy changes by less than 2%.

Code structure

File	Contents
PICParticleContainer.H/.cpp	Pure SoA particle species: initialization, charge deposition, field gather, push
FieldSolver.H/.cpp	Periodic Poisson solve with MLMG and \(E = -\nabla\phi\)
main.cpp	Setup, time loop, diagnostics and plotfiles

The time loop in main.cpp:

// Charge deposition, field solve and gather at time t^n.
electrons.depositCharge(rho);
rho.plus(rho_ion, 0, 1, 0);
solver.solvePhi(phi, rho);
solver.computeE(efield, phi);
electrons.gatherField(efield);

// Leapfrog: velocities live at half steps. Start with v^{-1/2}.
if (step == 0) { electrons.pushVelocity(-0.5_rt*dt); }

Charge deposition in PICParticleContainer.cpp:

// ParticleToMesh zeros rho, then adds ghost cell contributions to the
// valid cells they overlap, including across periodic boundaries.
amrex::ParticleToMesh(*this, rho, lev,
    [=] AMREX_GPU_DEVICE (const ConstPTDType& ptd, int i,
                          amrex::Array4<amrex::Real> const& rho_arr) noexcept
    {
        const ConstParticle p = ptd[i];
        amrex::ParticleInterpolator::Linear interp(p, plo, dxi);
        interp.ParticleToMesh(p, rho_arr, 0, 0, 1,
            [=] (ConstParticle const&, int) noexcept
            {
                return q_over_vol * ptd.rdata(PIdx::w)[i];
            });
    });

The field solver sets up MLPoisson and MLMG once, in its constructor, and reuses them every step:

FieldSolver::FieldSolver (const amrex::Geometry& geom, const amrex::BoxArray& ba,
                          const amrex::DistributionMapping& dm, amrex::Real epsilon0)
    : m_geom(geom),
      m_epsilon0(epsilon0),
      m_linop({geom}, {ba}, {dm}),
      m_mlmg(m_linop)
{
    // All-periodic domain. The operator is singular; MLMG removes the mean
    // of the right-hand side to make the problem solvable.
    m_linop.setDomainBC({AMREX_D_DECL(amrex::LinOpBCType::Periodic,
                                      amrex::LinOpBCType::Periodic,
                                      amrex::LinOpBCType::Periodic)},
                        {AMREX_D_DECL(amrex::LinOpBCType::Periodic,
                                      amrex::LinOpBCType::Periodic,
                                      amrex::LinOpBCType::Periodic)});
    m_linop.setLevelBC(0, nullptr);
    m_mlmg.setVerbose(0);
}

Run (from the top of the repository, see Getting Started):

cmake --build build --target Particles_ElectrostaticPIC -j 8
cd build/Particles/ElectrostaticPIC
mpiexec -n 2 ./ElectrostaticPIC inputs

n_cell and nppc take three values; 2D uses the first two. The plotfiles plt00000, plt00050 and plt00100 contain rho, phi, the electric field and the electrons.

Initial E mode amplitude: 0.009936866131 (linear theory 0.01)
Measured frequency: 0.9956005796 (plasma frequency 1)
Relative errors: frequency 0.004399420394, amplitude 0.006313386911, max energy change 0.001909800398
PASSED

To add species, create one container per species and sum their charge densities before the solve. For other boundary conditions, change setDomainBC in FieldSolver and the particle boundary handling; for periodic problems, Poisson can replace MLMG when AMReX is built with FFT support.