CD_ParticleContainerImplem.H

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: 2021-2026 SINTEF Energy Research
 *
 * SPDX-License-Identifier: GPL-3.0-or-later
 */
 
#ifndef CD_PARTICLECONTAINERIMPLEM_H
#define CD_PARTICLECONTAINERIMPLEM_H
 
// Std includes
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <map>
#include <unordered_map>
#include <utility>
#include <vector>
 
// Chombo includes
#include <CH_Timer.H>
#include <MayDay.H>
#include <CH_assert.H>
#include <BoxIterator.H>
#include <parstream.H>
 
// Our includes
#include <CD_ParticleContainer.H>
#include <CD_ParticleMemory.H>
#include <CD_NamespaceHeader.H>
 
template <typename P, typename Traits>
inline LevelTiles::LevelAndBox
ParticleContainer<P, Traits>::findDestination(const RealVect& a_pos) const
{
  return LevelTiles::findDestination(a_pos, m_probLo, m_dx, m_minBlockSize, m_levelTiles, m_finestLevel);
}
 
template <typename P, typename Traits>
inline typename ParticleContainer<P, Traits>::MoverPool
ParticleContainer<P, Traits>::gatherToPool(AMRParticlesSoA<P, Traits>&      a_source,
                                           const Vector<DisjointBoxLayout>& a_sourceGrids)
{
  CH_TIME("ParticleContainer::gatherToPool");
 
  MoverPool pool(numProc());
 
  // Bin every particle of one leaf into a destination pool, returning the number of off-domain
  // (outcast) particles. Shared by the serial and OpenMP paths so the binning logic lives once.
  // findDestination() is const and reads only the (immutable-during-gather) tile maps, so it is
  // safe to call concurrently; each thread bins into its own thread-local pool.
  auto binLeaf = [this](MoverPool& a_target, const Leaf& a_leaf) -> unsigned long long {
    unsigned long long outcast = 0;
 
    for (std::size_t i = 0; i < a_leaf.size(); i++) {
      const RealVect pos = a_leaf.position(i);
      const auto     dst = this->findDestination(pos);
 
      if (!dst.valid) {
        outcast++;
        continue;
      }
 
      Leaf& bin = a_target[dst.rank][PoolKey(dst.level, dst.gridIndex)];
 
      bin.append(pos, a_leaf.weight(i), a_leaf.gather(i));
 
      const std::size_t k = bin.size() - 1;
 
      bin.particleID(k) = a_leaf.particleID(i); // preserve id
      bin.rankID(k)     = dst.rank;             // new owner
    }
 
    return outcast;
  };
 
  unsigned long long numOutcast = 0;
 
  for (int lvl = 0; lvl < a_source.size(); lvl++) {
    LevelParticles&     level = *a_source[lvl];
    const DataIterator& dit   = a_sourceGrids[lvl].dataIterator();
    const int           nbox  = dit.size();
 
#ifdef _OPENMP
#pragma omp parallel reduction(+ : numOutcast)
    {
      MoverPool threadPool(numProc()); // thread-local; no locking while binning
 
#pragma omp for schedule(runtime)
      for (int mybox = 0; mybox < nbox; mybox++) {
        Leaf& leaf = level[dit[mybox]];
 
        numOutcast += binLeaf(threadPool, leaf);
 
        // The pool now holds this leaf's particles, and the source is discarded by every caller, so
        // release the arena here rather than at the end. Otherwise the source and the pool are both
        // fully live until the caller returns, which doubles the particle footprint at exactly the
        // point in a regrid where it is already highest. Each thread owns its own boxes, so this is
        // race-free.
        leaf.clear();
        leaf.shrinkToFit();
      }
 
      // Merge the thread-local pool into the shared pool. First writer to a (rank,key) slot gets the
      // O(1) arena swap inside catenate(); later writers append.
#pragma omp critical(ParticleContainer_gatherMerge)
      {
        for (int r = 0; r < numProc(); r++) {
          for (auto& kv : threadPool[r]) {
            pool[r][kv.first].catenate(kv.second);
          }
        }
      }
    }
#else
    for (int mybox = 0; mybox < nbox; mybox++) {
      Leaf& leaf = level[dit[mybox]];
 
      numOutcast += binLeaf(pool, leaf);
 
      // See the OpenMP branch: drain as we go so the source and the pool are never both fully live.
      leaf.clear();
      leaf.shrinkToFit();
    }
#endif
  }
 
  m_numOutcastLocal = numOutcast;
 
  return pool;
}
 
template <typename P, typename Traits>
inline typename ParticleContainer<P, Traits>::MoverPool
ParticleContainer<P, Traits>::gatherMoversToPool()
{
  CH_TIME("ParticleContainer::gatherMoversToPool");
 
  CH_assert(m_validCells != nullptr);
 
  MoverPool          pool(numProc());
  unsigned long long numOutcast = 0;
 
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    LevelParticles&                 level     = *m_particles[lvl];
    const LevelData<BaseFab<bool>>* maskLevel = (*m_validCells)[lvl].isNull() ? nullptr : &(*(*m_validCells)[lvl]);
    const DataIterator&             dit       = m_grids[lvl].dataIterator();
    const int                       nbox      = dit.size();
 
    // Partition one leaf in place: keep stayers (still inside their box AND in a cell flagged valid =
    // not covered by a finer level), pool the movers and swap-pop them out. a_mask == nullptr forces
    // every particle to be treated as a mover (no usable mask for this leaf). Returns outcast movers.
    // findDestination()/cellIndex() are const reads, so this is safe to call from many threads as long
    // as each leaf is touched by a single thread (the omp-for partitions boxes disjointly).
    auto partitionLeaf = [this, lvl](MoverPool&                        a_target,
                                     [[maybe_unused]] const DataIndex& a_din,
                                     Leaf&                             a_leaf,
                                     const BaseFab<bool>*              a_mask) -> unsigned long long {
      unsigned long long outcast = 0;
 
      const Box maskBox = (a_mask != nullptr) ? a_mask->box() : Box();
 
      std::size_t i = 0;
 
      while (i < a_leaf.size()) {
        const RealVect pos = a_leaf.position(i);
        if (a_mask != nullptr) {
          const IntVect iv = this->cellIndex(lvl, pos);
 
          if (maskBox.contains(iv) && (*a_mask)(iv, 0)) {
            // Stayer: in its own box and not under a finer grid, so it cannot have changed owner.
            // Debug self-check of that invariant -- the destination the gather-all path would compute
            // must be exactly this (level, grid, rank). If this ever fires, the skip is unsound.
#ifndef NDEBUG
            const auto chk = this->findDestination(pos);
            CH_assert(chk.valid && chk.level == lvl && chk.rank == procID() &&
                      m_levelTiles[lvl]->getMyGrids().at(chk.gridIndex) == a_din);
#endif
            i++;
 
            continue;
          }
        }
 
        const auto dst = this->findDestination(pos); // mover: route it
 
        if (dst.valid) {
          Leaf& bin = a_target[dst.rank][PoolKey(dst.level, dst.gridIndex)];
 
          bin.append(pos, a_leaf.weight(i), a_leaf.gather(i));
 
          const std::size_t k = bin.size() - 1;
 
          bin.particleID(k) = a_leaf.particleID(i); // preserve id
          bin.rankID(k)     = dst.rank;             // new owner
        }
        else {
          outcast++;
        }
 
        a_leaf.remove(i); // swap-pop; do NOT advance i (a new particle now occupies slot i)
      }
 
      return outcast;
    };
 
#ifdef _OPENMP
#pragma omp parallel reduction(+ : numOutcast)
    {
      MoverPool threadPool(numProc());
 
#pragma omp for schedule(runtime)
      for (int mybox = 0; mybox < nbox; mybox++) {
        const DataIndex&     din  = dit[mybox];
        const BaseFab<bool>* mask = (maskLevel != nullptr) ? &(*maskLevel)[din] : nullptr;
 
        if (mask != nullptr && !mask->isUsable()) {
          mask = nullptr;
        }
 
        numOutcast += partitionLeaf(threadPool, din, level[din], mask);
      }
 
#pragma omp critical(ParticleContainer_gatherMerge)
      {
        for (int r = 0; r < numProc(); r++) {
          for (auto& kv : threadPool[r]) {
            pool[r][kv.first].catenate(kv.second);
          }
        }
      }
    }
#else
    for (int mybox = 0; mybox < nbox; mybox++) {
      const DataIndex&     din  = dit[mybox];
      const BaseFab<bool>* mask = (maskLevel != nullptr) ? &(*maskLevel)[din] : nullptr;
 
      if (mask != nullptr && !mask->isUsable()) {
        mask = nullptr;
      }
 
      numOutcast += partitionLeaf(pool, din, level[din], mask);
    }
#endif
  }
 
  m_numOutcastLocal = numOutcast;
 
  return pool;
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::distributeFromPool(MoverPool& a_pool, const bool a_setOwnerToReceiver)
{
  CH_TIME("ParticleContainer::distributeFromPool");
 
  const int myRank   = procID();
  const int numRanks = numProc();
 
  // ---- same-rank movers: catenate each pooled leaf into its destination leaf ----
  // On the gather-all path the destination leaves are empty (remap clears them; regrid allocates them
  // fresh) so catenate() takes the O(1) arena-swap fast path. On the mask fast path the destinations
  // still hold their stayers, so catenate() appends (O(movers)); a mover never returns to its own
  // origin leaf, so there is no self-append.
  for (auto& kv : a_pool[myRank]) {
    const int          lvl  = kv.first.first;
    const unsigned int gidx = kv.first.second;
    const DataIndex    din  = m_levelTiles[lvl]->getMyGrids().at(gidx);
 
    (*m_particles[lvl])[din].catenate(kv.second);
  }
 
#ifdef CH_MPI
  // ---- cross-rank movers: one Alltoallv of linearized particles ----
  constexpr std::size_t pp        = Leaf::bytesPerParticle();
  constexpr std::size_t headerLen = 2 * sizeof(std::uint32_t) + sizeof(std::uint64_t);
 
  // Size the send side before packing anything. A destination's byte count is arithmetic over the
  // pool -- a header plus the payload for each pooled leaf -- so computing it up front costs nothing,
  // and knowing it lets every particle be linearized straight into the flat buffer that MPI reads.
  // Packing into per-destination buffers first and copying those into the flat one afterwards would
  // hold a second full copy of every outbound particle for the duration of the copy.
  std::vector<int>          sendCounts(numRanks, 0);
  std::vector<int>          recvCounts(numRanks, 0);
  std::vector<std::int64_t> sendBytes(numRanks, 0);
 
  std::int64_t stot = 0;
 
  for (int r = 0; r < numRanks; r++) {
    if (r == myRank) {
      continue;
    }
 
    std::int64_t bytes = 0;
 
    for (const auto& kv : a_pool[r]) {
      bytes += static_cast<std::int64_t>(headerLen + kv.second.size() * pp);
    }
 
    sendBytes[r] = bytes;
    stot += bytes;
  }
 
  // MPI counts and displacements are int, so the byte totals have to fit in one. This is a real limit
  // at scale rather than a formality -- a rank holding 2 GB of outbound particles reaches it -- and it
  // has to be enforced in every build. CH_assert compiles out under OPT, which is precisely where the
  // scale is, so a guard that used it would let the displacements overflow and corrupt the exchange
  // silently instead of stopping it.
  if (stot > static_cast<std::int64_t>(std::numeric_limits<int>::max())) {
    MayDay::Abort("ParticleContainer::distributeFromPool - outbound particle bytes exceed MPI's int limit");
  }
 
  for (int r = 0; r < numRanks; r++) {
    sendCounts[r] = static_cast<int>(sendBytes[r]);
  }
 
  MPI_Alltoall(sendCounts.data(), 1, MPI_INT, recvCounts.data(), 1, MPI_INT, Chombo_MPI::comm);
 
  std::vector<int> sdispl(numRanks, 0);
  std::vector<int> rdispl(numRanks, 0);
 
  std::int64_t sacc = 0;
  std::int64_t rtot = 0;
 
  for (int r = 0; r < numRanks; r++) {
    sdispl[r] = static_cast<int>(sacc);
    sacc += sendCounts[r];
 
    rdispl[r] = static_cast<int>(rtot);
    rtot += recvCounts[r];
  }
 
  if (rtot > static_cast<std::int64_t>(std::numeric_limits<int>::max())) {
    MayDay::Abort("ParticleContainer::distributeFromPool - inbound particle bytes exceed MPI's int limit");
  }
 
  std::vector<char> sflat(stot);
  std::vector<char> rflat(rtot);
 
  // The flat buffers are particle payload in flight, so account for them alongside the arenas. They
  // are part of the peak this rank has to fit, and on a particle-heavy step they are comparable in
  // size to the arenas themselves -- leaving them out would report a peak that understates the point
  // in the step where an out-of-memory failure actually happens.
  const std::size_t flatBytes = static_cast<std::size_t>(stot) + static_cast<std::size_t>(rtot);
 
  ParticleMemory::addBytes(ParticleMemory::Kind::Buffer, flatBytes);
 
  // Each destination's slice of sflat is disjoint and linearizeParticle is a const read, so packing
  // different destinations in parallel is race-free. A pooled leaf is spent once it has been
  // linearized, so release it there and then: holding it until the pool goes out of scope would keep
  // a second copy of every outbound particle alive across the exchange itself, which is the point in
  // the step where the footprint peaks.
#ifdef _OPENMP
#pragma omp parallel for schedule(runtime)
#endif
  for (int r = 0; r < numRanks; r++) {
    if (r == myRank) {
      continue;
    }
 
    char* p = sflat.data() + sdispl[r];
 
    for (auto& kv : a_pool[r]) {
      Leaf& src = kv.second;
 
      const std::uint32_t lvl   = static_cast<std::uint32_t>(kv.first.first);
      const std::uint32_t gidx  = static_cast<std::uint32_t>(kv.first.second);
      const std::uint64_t count = static_cast<std::uint64_t>(src.size());
 
      std::memcpy(p, &lvl, sizeof(std::uint32_t));
      p += sizeof(std::uint32_t);
 
      std::memcpy(p, &gidx, sizeof(std::uint32_t));
      p += sizeof(std::uint32_t);
 
      std::memcpy(p, &count, sizeof(std::uint64_t));
      p += sizeof(std::uint64_t);
 
      for (std::size_t i = 0; i < count; i++) {
        src.linearizeParticle(p, i);
        p += pp;
      }
 
      src.clear();
      src.shrinkToFit();
    }
  }
 
  MPI_Alltoallv(sflat.data(),
                sendCounts.data(),
                sdispl.data(),
                MPI_BYTE,
                rflat.data(),
                recvCounts.data(),
                rdispl.data(),
                MPI_BYTE,
                Chombo_MPI::comm);
 
  // ---- unpack: group received chunks by destination leaf, then append in parallel ----
  // Pass 1 (serial, header-only walk): record each chunk's payload pointer + count, grouped by its
  // destination (level, gridIndex). Pass 2 appends in parallel with one destination leaf per thread,
  // so no two threads ever touch the same leaf (delinearizeAndAppend mutates the leaf arena).
  struct RecvChunk
  {
    const char*   payload; // first particle of the chunk within rflat
    std::uint64_t count;
  };
  struct DestGroup
  {
    int                    level;
    unsigned int           gridIndex;
    std::uint64_t          total;
    std::vector<RecvChunk> chunks;
  };
 
  std::vector<DestGroup>         groups;
  std::map<PoolKey, std::size_t> keyToGroup;
  {
    const char* p   = rflat.data();
    const char* end = p + rtot;
 
    while (p < end) {
      std::uint32_t lvl;
      std::uint32_t gidx;
      std::uint64_t count;
 
      std::memcpy(&lvl, p, sizeof(std::uint32_t));
      p += sizeof(std::uint32_t);
 
      std::memcpy(&gidx, p, sizeof(std::uint32_t));
      p += sizeof(std::uint32_t);
 
      std::memcpy(&count, p, sizeof(std::uint64_t));
      p += sizeof(std::uint64_t);
 
      const PoolKey key(static_cast<int>(lvl), gidx);
      const auto    it = keyToGroup.find(key);
 
      if (it == keyToGroup.end()) {
        keyToGroup.emplace(key, groups.size());
        groups.push_back(DestGroup{static_cast<int>(lvl), gidx, count, {RecvChunk{p, count}}});
      }
      else {
        DestGroup& g = groups[it->second];
        g.total += count;
        g.chunks.push_back(RecvChunk{p, count});
      }
 
      p += count * pp;
    }
  }
 
  const int numGroups = static_cast<int>(groups.size());
 
#ifdef _OPENMP
#pragma omp parallel for schedule(runtime)
#endif
  for (int d = 0; d < numGroups; d++) {
    const DestGroup& g   = groups[d];
    const DataIndex  din = m_levelTiles[g.level]->getMyGrids().at(g.gridIndex);
 
    Leaf& tgt = (*m_particles[g.level])[din];
 
    tgt.reserve(tgt.size() + g.total); // one allocation; avoids geometric regrow during the appends
    for (const RecvChunk& chunk : g.chunks) {
      const char* q = chunk.payload;
 
      for (std::uint64_t c = 0; c < chunk.count; c++) {
        tgt.delinearizeAndAppend(q);
        q += pp;
        if (a_setOwnerToReceiver) {
          tgt.rankID(tgt.size() - 1) = myRank; // remap/regrid: owner is the receiving rank
        }
        // else: keep the rankID carried in the buffer (ghost particles record their true owner)
      }
    }
  }
 
  ParticleMemory::removeBytes(ParticleMemory::Kind::Buffer, flatBytes);
#endif
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::remap()
{
  CH_TIME("ParticleContainer::remap");
 
  CH_assert(m_isDefined);
 
  // Ghost particles are transient, non-owned copies; drop them so they are never routed as if owned.
  this->clearGhostParticles();
 
  if (m_validCells != nullptr && static_cast<int>(m_validCells->size()) >= m_finestLevel + 1) {
    // Mask fast path: a particle still in its box and not under a finer grid cannot change owner, so we
    // leave those "stayers" in place and pool only the movers. Destinations therefore keep their stayers
    // -- do NOT clearParticles() here; distributeFromPool() appends the movers onto the existing leaves.
    MoverPool pool = this->gatherMoversToPool();
 
    this->distributeFromPool(pool, true); // remap: the receiving rank becomes the new owner
  }
  else {
    // Fallback: re-pool every valid particle (used when no mask is available, e.g. addParticlesDestructive
    // staging, or before the realm's valid-cell mask has been built).
    MoverPool pool = this->gatherToPool(m_particles, m_grids);
 
    this->clearParticles();
    this->distributeFromPool(pool, true); // remap: the receiving rank becomes the new owner
  }
 
  m_isOrganizedByCell = false;
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::preRegrid()
{
  CH_TIME("ParticleContainer::preRegrid");
 
  CH_assert(m_isDefined);
 
  // Ghost particles are transient, non-owned copies; drop them so they are never cached/regridded.
  this->clearGhostParticles();
 
  m_cacheParticles = m_particles;
  m_cacheGrids     = m_grids;
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::regrid(const Vector<DisjointBoxLayout>&         a_grids,
                                     const Vector<ProblemDomain>&             a_domains,
                                     const Vector<Real>&                      a_dx,
                                     const Vector<int>&                       a_refRat,
                                     const int                                a_minBlockSize,
                                     const Vector<RefCountedPtr<LevelTiles>>& a_levelTiles,
                                     const int                                a_newFinestLevel)
{
  CH_TIME("ParticleContainer::regrid");
 
  CH_assert(m_isDefined);
  CH_assert(a_levelTiles.size() > a_newFinestLevel);
 
  // ---- adopt the new layout ----
  m_grids        = a_grids;
  m_domains      = a_domains;
  m_refRat       = a_refRat;
  m_minBlockSize = a_minBlockSize;
  m_finestLevel  = a_newFinestLevel;
 
  // Alias the Realm's per-level tile->box maps (the single source of truth); never build our own.
  m_levelTiles = a_levelTiles;
 
  m_dx.resize(a_newFinestLevel + 1);
  for (int lvl = 0; lvl <= a_newFinestLevel; lvl++) {
    m_dx[lvl] = a_dx[lvl] * RealVect::Unit;
  }
 
  // ---- re-allocate the (empty) holders over the new grids ----
  m_particles.resize(a_newFinestLevel + 1);
  m_maskParticles.resize(a_newFinestLevel + 1);
 
  for (int lvl = 0; lvl <= a_newFinestLevel; lvl++) {
    m_particles[lvl]     = RefCountedPtr<LevelParticles>(new LevelParticles(a_grids[lvl]));
    m_maskParticles[lvl] = RefCountedPtr<LevelParticles>(new LevelParticles(a_grids[lvl]));
  }
 
  this->setupGrownGrids();
 
  m_bufferParticles.resize(a_newFinestLevel + 1);
 
  for (int lvl = 0; lvl <= a_newFinestLevel; lvl++) {
    m_bufferParticles[lvl] = RefCountedPtr<LevelParticles>(new LevelParticles(m_grownGrids[lvl]));
  }
 
  // ---- redistribute the cached particles onto the new layout ----
  MoverPool pool = this->gatherToPool(m_cacheParticles, m_cacheGrids);
  this->distributeFromPool(pool, true); // regrid: the receiving rank becomes the new owner
 
  // ---- release the cache ----
  m_cacheParticles.resize(0);
  m_cacheGrids.resize(0);
 
  m_isOrganizedByCell = false;
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::clearGhostParticles()
{
  CH_TIME("ParticleContainer::clearGhostParticles");
 
  CH_assert(m_isDefined);
 
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    LevelParticles&          level = *m_particles[lvl];
    const DisjointBoxLayout& dbl   = m_grids[lvl];
    const DataIterator&      dit   = dbl.dataIterator();
 
    const int nbox = dit.size();
 
#pragma omp parallel for schedule(runtime)
    for (int mybox = 0; mybox < nbox; mybox++) {
      const DataIndex& din  = dit[mybox];
      Leaf&            leaf = level[din];
 
      // Swap-and-pop every ghost; the swapped-in particle now occupies index i, so do NOT advance on a hit.
      std::size_t i = 0;
      while (i < leaf.size()) {
        if (leaf.isGhost(i)) {
          leaf.remove(i);
        }
        else {
          i++;
        }
      }
    }
  }
 
  m_isOrganizedByCell = false; // swap-pop reorders the surviving valid particles
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::claimOwnership(const bool a_onlyValidRegion)
{
  CH_TIME("ParticleContainer::claimOwnership");
 
  CH_assert(m_isDefined);
 
  const RankID myRank = static_cast<RankID>(procID());
 
  // Ghosts are left untouched (they carry their true owner rank). Each non-ghost particle in the valid
  // region (or all non-ghost particles when !a_onlyValidRegion) is stamped with this rank. Per-particle
  // and box-local -> embarrassingly parallel.
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    LevelParticles&          level = *m_particles[lvl];
    const DisjointBoxLayout& dbl   = m_grids[lvl];
    const DataIterator&      dit   = dbl.dataIterator();
    const int                nbox  = dit.size();
 
#pragma omp parallel for schedule(runtime)
    for (int mybox = 0; mybox < nbox; mybox++) {
      const DataIndex& din  = dit[mybox];
      const Box        box  = dbl[din];
      Leaf&            leaf = level[din];
 
      for (std::size_t i = 0; i < leaf.size(); i++) {
        if (leaf.isGhost(i)) {
          continue;
        }
        if (!a_onlyValidRegion || this->inValidBox(lvl, box, leaf.position(i))) {
          leaf.rankID(i) = myRank;
        }
      }
    }
  }
}
 
template <typename P, typename Traits>
inline long long
ParticleContainer<P, Traits>::resetParticleIDs(const bool a_onlyValidRegion, const ParticleID a_startID)
{
  CH_TIME("ParticleContainer::resetParticleIDs");
 
  CH_assert(m_isDefined);
 
  // A non-ghost particle qualifies for a new id iff (!a_onlyValidRegion) or its cell lies in the leaf's
  // grid box. Ghosts never qualify (and are left untouched); non-qualifying non-ghost strays get
  // s_invalidID. The predicate must be identical in the count and assign passes below.
 
  // Flatten the local boxes in (level, box) order -- this fixes the local id ordering.
  std::vector<std::pair<int, DataIndex>> boxes;
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    const DataIterator& dit  = m_grids[lvl].dataIterator();
    const int           nbox = dit.size();
    for (int mybox = 0; mybox < nbox; mybox++) {
      boxes.emplace_back(lvl, dit[mybox]);
    }
  }
  const int numBoxes = static_cast<int>(boxes.size());
 
  // Pass 1: count qualifying particles per box (parallel; each box independent).
  std::vector<long long> boxCount(numBoxes, 0);
#pragma omp parallel for schedule(runtime)
  for (int b = 0; b < numBoxes; b++) {
    const int        lvl  = boxes[b].first;
    const DataIndex& din  = boxes[b].second;
    const Box        box  = m_grids[lvl][din];
    const Leaf&      leaf = (*m_particles[lvl])[din];
 
    long long count = 0;
    for (std::size_t i = 0; i < leaf.size(); i++) {
      if (!leaf.isGhost(i) && (!a_onlyValidRegion || this->inValidBox(lvl, box, leaf.position(i)))) {
        count++;
      }
    }
    boxCount[b] = count;
  }
 
  // Local exclusive prefix over boxes + local total.
  std::vector<long long> boxLocalOffset(numBoxes, 0);
  long long              localTotal = 0;
  for (int b = 0; b < numBoxes; b++) {
    boxLocalOffset[b] = localTotal;
    localTotal += boxCount[b];
  }
 
  // Per-rank exclusive offset: all-gather each rank's local total, then sum the ranks before this one.
  const Vector<long long> perRankCount = ParallelOps::gather(localTotal);
 
  ParticleID rankOffset = a_startID;
  for (int r = 0; r < procID(); r++) {
    rankOffset += perRankCount[r];
  }
 
  // Pass 2: assign ids (parallel; each box starts at its own offset). Ghosts untouched; qualifying get a
  // contiguous id; non-qualifying non-ghost strays get s_invalidID.
#pragma omp parallel for schedule(runtime)
  for (int b = 0; b < numBoxes; b++) {
    const int        lvl  = boxes[b].first;
    const DataIndex& din  = boxes[b].second;
    const Box        box  = m_grids[lvl][din];
    Leaf&            leaf = (*m_particles[lvl])[din];
 
    ParticleID id = rankOffset + static_cast<ParticleID>(boxLocalOffset[b]);
    for (std::size_t i = 0; i < leaf.size(); i++) {
      if (leaf.isGhost(i)) {
        continue;
      }
      if (!a_onlyValidRegion || this->inValidBox(lvl, box, leaf.position(i))) {
        leaf.particleID(i) = id++;
      }
      else {
        leaf.particleID(i) = ParticleSoA<P, Traits>::s_invalidID;
      }
    }
  }
 
  return ParallelOps::sum(localTotal);
}
 
template <typename P, typename Traits>
inline typename ParticleContainer<P, Traits>::MoverPool
ParticleContainer<P, Traits>::gatherGhostsFromMasks(const AMRParticleGhostMask& a_maskSame,
                                                    const AMRParticleGhostMask& a_maskCoarToFine,
                                                    const AMRParticleGhostMask& a_maskFineToCoar)
{
  CH_TIME("ParticleContainer::gatherGhostsFromMasks");
 
  MoverPool pool(numProc());
 
  // Bin one source leaf's valid particles as ghosts of every target box listed by its cell's masks. The
  // three masks are all indexed by the SOURCE cell on the source level, so a single cell lookup drives all
  // directions. GhostType is from the RECEIVER's view: same level -> SameLevel, coarse source scattered to
  // a finer patch -> Coarse, fine source scattered to a coarser patch -> Fine.
  auto binLeaf = [&](MoverPool& a_target, const Leaf& a_leaf, const int a_srcLvl, const DataIndex& a_din) {
    const bool haveSame = (a_srcLvl < a_maskSame.size()) && !a_maskSame[a_srcLvl].isNull();
    const bool haveC2F  = (a_srcLvl < m_finestLevel) && (a_srcLvl < a_maskCoarToFine.size()) &&
                         !a_maskCoarToFine[a_srcLvl].isNull();
    const bool haveF2C = (a_srcLvl >= 1) && (a_srcLvl < a_maskFineToCoar.size()) &&
                         !a_maskFineToCoar[a_srcLvl].isNull();
 
    const ParticleGhostMask* same = haveSame ? &((*a_maskSame[a_srcLvl])[a_din]) : nullptr;
    const ParticleGhostMask* c2f  = haveC2F ? &((*a_maskCoarToFine[a_srcLvl])[a_din]) : nullptr;
    const ParticleGhostMask* f2c  = haveF2C ? &((*a_maskFineToCoar[a_srcLvl])[a_din]) : nullptr;
 
    if (same == nullptr && c2f == nullptr && f2c == nullptr) {
      return;
    }
 
    for (std::size_t i = 0; i < a_leaf.size(); i++) {
      const RealVect pos = a_leaf.position(i);
      const IntVect  iv  = this->cellIndex(a_srcLvl, pos);
 
      auto scatter =
        [&](const ParticleGhostMask* a_mask, const int a_dstLvl, const GhostType a_gt, const bool a_pruneByTargetBox) {
          if (a_mask == nullptr || !a_mask->box().contains(iv)) {
            return;
          }
          const int nt = a_mask->numTargets(iv);
 
          // Coarse-to-fine only: the source (coarse) cell is coarser than the fine ghost width, so a whole
          // coarse cell would over-scatter. Keep the particle for a target only if its DESTINATION (fine)
          // cell lies in that target's acceptance box (the exact fine ghost shell).
          const IntVect dstCell = a_pruneByTargetBox ? this->cellIndex(a_dstLvl, pos) : IntVect::Zero;
 
          for (int k = 0; k < nt; k++) {
            if (a_pruneByTargetBox && !a_mask->targetBox(iv, k).contains(dstCell)) {
              continue;
            }
 
            Leaf& bin = a_target[a_mask->targetRank(iv, k)][PoolKey(a_dstLvl, a_mask->targetGridIndex(iv, k))];
 
            bin.append(pos, a_leaf.weight(i), a_leaf.gather(i));
 
            const std::size_t j = bin.size() - 1;
 
            bin.particleID(j) = a_leaf.particleID(i); // preserve id
            bin.rankID(j)     = procID();             // owner is this rank (it owns its valid particles)
            bin.ghost(j)      = a_gt;
          }
        };
 
      scatter(same, a_srcLvl, GhostType::SameLevel, false);
      scatter(c2f, a_srcLvl + 1, GhostType::Coarse, true);
      scatter(f2c, a_srcLvl - 1, GhostType::Fine, false);
    }
  };
 
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    const LevelParticles& level = *m_particles[lvl];
    const DataIterator&   dit   = m_grids[lvl].dataIterator();
    const int             nbox  = dit.size();
 
#ifdef _OPENMP
#pragma omp parallel
    {
      MoverPool threadPool(numProc()); // thread-local; no locking while binning
 
#pragma omp for schedule(runtime)
      for (int mybox = 0; mybox < nbox; mybox++) {
        binLeaf(threadPool, level[dit[mybox]], lvl, dit[mybox]);
      }
 
#pragma omp critical(ParticleContainer_gatherGhostsFromMasksMerge)
      {
        for (int r = 0; r < numProc(); r++) {
          for (auto& kv : threadPool[r]) {
            pool[r][kv.first].catenate(kv.second);
          }
        }
      }
    }
#else
    for (int mybox = 0; mybox < nbox; mybox++) {
      binLeaf(pool, level[dit[mybox]], lvl, dit[mybox]);
    }
#endif
  }
 
  return pool;
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::fillGhostParticles(const AMRParticleGhostMask& a_maskSame,
                                                 const AMRParticleGhostMask& a_maskCoarToFine,
                                                 const AMRParticleGhostMask& a_maskFineToCoar)
{
  CH_TIME("ParticleContainer::fillGhostParticles");
 
  CH_assert(m_isDefined);
 
  // Refill from a clean slate: drop any existing ghost halo first.
  this->clearGhostParticles();
 
  MoverPool pool = this->gatherGhostsFromMasks(a_maskSame, a_maskCoarToFine, a_maskFineToCoar);
 
  // a_setOwnerToReceiver == false: ghost copies keep their true owner's rankID (not the receiver's).
  this->distributeFromPool(pool, false);
 
  m_isOrganizedByCell = false; // ghosts appended to the leaves invalidate any cell sort
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::setupGrownGrids()
{
  CH_TIME("ParticleContainer::setupGrownGrids");
 
  m_grownGrids.resize(m_finestLevel + 1);
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    const DisjointBoxLayout& dbl = m_grids[lvl];
 
    Vector<Box> boxes = dbl.boxArray();
 
    if (lvl > 0) {
      const Box domainBox = m_domains[lvl].domainBox();
      for (int i = 0; i < boxes.size(); i++) {
        boxes[i].grow(m_refRat[lvl - 1]);
        boxes[i] &= domainBox;
      }
    }
 
    m_grownGrids[lvl] = BoxLayout(boxes, dbl.procIDs());
  }
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::copyMaskParticles(const Vector<RefCountedPtr<LevelData<BaseFab<bool>>>>& a_mask)
{
  CH_TIME("ParticleContainer::copyMaskParticles(amr)");
 
  CH_assert(m_isDefined);
 
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    if (!a_mask[lvl].isNull()) {
      this->copyMaskParticles(lvl, *a_mask[lvl]);
    }
  }
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::copyMaskParticles(const int a_level, const LevelData<BaseFab<bool>>& a_mask)
{
  CH_TIME("ParticleContainer::copyMaskParticles(level)");
 
  CH_assert(m_isDefined);
 
  LevelParticles&       maskLevel = *m_maskParticles[a_level];
  const LevelParticles& srcLevel  = *m_particles[a_level];
 
  const DisjointBoxLayout& dbl = m_grids[a_level];
  const DataIterator&      dit = dbl.dataIterator();
 
  const int nbox = dit.size();
 
#pragma omp parallel for schedule(runtime)
  for (int mybox = 0; mybox < nbox; mybox++) {
    const DataIndex& din = dit[mybox];
 
    maskLevel[din].clear();
 
    const BaseFab<bool>& mask = a_mask[din];
    if (!mask.isUsable()) {
      continue;
    }
    const Box maskBox = mask.box();
 
    const Leaf& src = srcLevel[din];
    Leaf&       dst = maskLevel[din];
 
    for (std::size_t i = 0; i < src.size(); i++) {
      const IntVect iv = this->cellIndex(a_level, src.position(i));
 
      if (maskBox.contains(iv) && mask(iv, 0)) {
        dst.append(src.position(i), src.weight(i), src.gather(i));
 
        const std::size_t k = dst.size() - 1;
 
        dst.particleID(k) = src.particleID(i);
        dst.rankID(k)     = src.rankID(i);
      }
    }
  }
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::transferMaskParticles(const Vector<RefCountedPtr<LevelData<BaseFab<bool>>>>& a_mask)
{
  CH_TIME("ParticleContainer::transferMaskParticles(amr)");
 
  CH_assert(m_isDefined);
 
  for (int lvl = 0; lvl <= m_finestLevel; lvl++) {
    if (!a_mask[lvl].isNull()) {
      this->transferMaskParticles(lvl, *a_mask[lvl]);
    }
  }
}
 
template <typename P, typename Traits>
inline void
ParticleContainer<P, Traits>::transferMaskParticles(const int a_level, const LevelData<BaseFab<bool>>& a_mask)
{
  CH_TIME("ParticleContainer::transferMaskParticles(level)");
 
  CH_assert(m_isDefined);
 
  LevelParticles& maskLevel = *m_maskParticles[a_level];
  LevelParticles& srcLevel  = *m_particles[a_level];
 
  const DisjointBoxLayout& dbl = m_grids[a_level];
  const DataIterator&      dit = dbl.dataIterator();
 
  const int nbox = dit.size();
 
#pragma omp parallel for schedule(runtime)
  for (int mybox = 0; mybox < nbox; mybox++) {
    const DataIndex& din = dit[mybox];
 
    const BaseFab<bool>& mask = a_mask[din];
    if (!mask.isUsable()) {
      continue;
    }
    const Box maskBox = mask.box();
 
    Leaf& src = srcLevel[din];
    Leaf& dst = maskLevel[din];
 
    std::size_t i = 0;
    while (i < src.size()) {
      const IntVect iv = this->cellIndex(a_level, src.position(i));
 
      if (maskBox.contains(iv) && mask(iv, 0)) {
        dst.append(src.position(i), src.weight(i), src.gather(i));
 
        const std::size_t k = dst.size() - 1;
 
        dst.particleID(k) = src.particleID(i);
        dst.rankID(k)     = src.rankID(i);
 
        src.remove(i); // swap-and-pop; the swapped-in particle now occupies index i
      }
      else {
        ++i;
      }
    }
  }
}
 
#include <CD_NamespaceFooter.H>
 
#endif