particle_migrator.hpp

Go to the documentation of this file.

// core — Lagrangian particle migration over the same block decomposition.
//
// After particles move, each one may now belong to a different rank's block. migrate() reassigns
// every particle to the rank that owns the block containing its (periodically wrapped) position,
// shipping departing particles to their new owner via the NBX engine — the dynamic/sparse exchange
// the consensus protocol is built for. Positions and an opaque fixed-stride payload travel
// together, so the caller can carry velocity, orientation, id, etc. without this layer knowing the
// schema.
//
// This is the Lagrangian counterpart to peclet::core::halo::GridHaloTopology (Eulerian ghost
// cells): same decomposition, same async engine, different payload — the field-agnostic design in
// action.
#ifndef PECLET_CORE_HALO_PARTICLE_MIGRATOR_HPP
#define PECLET_CORE_HALO_PARTICLE_MIGRATOR_HPP
 
#include <algorithm>
#include <cmath>
#include <cstring>
#include <limits>
#include <map>
#include <utility>
#include <vector>
 
#include "peclet/core/common/mpi.hpp"
#include "peclet/core/common/types.hpp"
#include "peclet/core/decomp/block_decomposer.hpp"
#include "peclet/core/halo/nbx.hpp"
 
namespace peclet::core::halo {
 
template <int Dim>
struct DomainMap {
  Vec<Dim> origin{};                 
  Vec<Dim> cellSize{};               
  std::array<bool, Dim> periodic{};  
};
 
template <int Dim>
class ParticleMigrator {
 public:
  void init(const decomp::BlockDecomposer<Dim>& dec, int rank, DomainMap<Dim> map,
            MPI_Comm comm = MPI_COMM_WORLD) {
    dec_ = &dec;
    rank_ = rank;
    map_ = map;
    comm_ = comm;
  }
 
  IVec<Dim> cellOf(const Vec<Dim>& x) const {
    IVec<Dim> g{};
    const IVec<Dim>& gsize = dec_->globalSize();
    for (int i = 0; i < Dim; ++i) {
      double rel = (x[i] - map_.origin[i]) / map_.cellSize[i];
      Index c = static_cast<Index>(std::floor(rel));
      if (map_.periodic[i]) {
        c = wrap(c, gsize[i]);
      } else {
        c = std::clamp<Index>(c, 0, gsize[i] - 1);
      }
      g[i] = c;
    }
    return g;
  }
 
  int ownerOf(const Vec<Dim>& x) const { return dec_->ownerOf(cellOf(x)); }
 
  Vec<Dim> wrapPosition(Vec<Dim> x) const {
    const IVec<Dim>& gsize = dec_->globalSize();
    for (int i = 0; i < Dim; ++i) {
      if (!map_.periodic[i])
        continue;
      double L = map_.cellSize[i] * static_cast<double>(gsize[i]);
      double rel = x[i] - map_.origin[i];
      rel = std::fmod(std::fmod(rel, L) + L, L);
      x[i] = map_.origin[i] + rel;
    }
    return x;
  }
 
  std::size_t migrate(std::vector<Vec<Dim>>& pos, std::vector<char>& payload, std::size_t stride) {
    const std::size_t n = pos.size();
    const std::size_t posBytes = Dim * sizeof(double);
    const std::size_t recBytes = posBytes + stride;
 
    std::map<int, std::vector<char>> outbox;
    std::size_t keep = 0;
    sent_ = 0;
    for (std::size_t i = 0; i < n; ++i) {
      Vec<Dim> wx = wrapPosition(pos[i]);
      int owner = ownerOf(wx);
      if (owner == rank_) {
        pos[keep] = wx;
        if (stride)
          std::memmove(&payload[keep * stride], &payload[i * stride], stride);
        ++keep;
      } else {
        auto& buf = outbox[owner];
        std::size_t off = buf.size();
        buf.resize(off + recBytes);
        std::memcpy(buf.data() + off, wx.data(), posBytes);
        if (stride)
          std::memcpy(buf.data() + off + posBytes, &payload[i * stride], stride);
        ++sent_;
      }
    }
    pos.resize(keep);
    payload.resize(keep * stride);
 
    // Ship departing particles to their owners and absorb arrivals (NBX).
    std::vector<std::pair<int, const std::vector<char>*>> queue;
    queue.reserve(outbox.size());
    for (auto& [r, buf] : outbox)
      queue.emplace_back(r, &buf);
 
    received_ = 0;
    NbxEngine nbx(comm_);
    std::size_t q = 0;
    auto packNext = [&](std::vector<char>& out) -> int {
      if (q >= queue.size())
        return -1;
      int dest = queue[q].first;
      out = *queue[q].second;
      ++q;
      return dest;
    };
    auto onRecv = [&](int /*src*/, std::vector<char>& msg) {
      std::size_t cnt = msg.size() / recBytes;
      for (std::size_t i = 0; i < cnt; ++i) {
        const char* rec = msg.data() + i * recBytes;
        Vec<Dim> x{};
        std::memcpy(x.data(), rec, posBytes);
        pos.push_back(x);
        if (stride) {
          std::size_t off = payload.size();
          payload.resize(off + stride);
          std::memcpy(&payload[off], rec + posBytes, stride);
        }
        ++received_;
      }
    };
    nbx.exchange(packNext, onRecv, /*tag=*/7401);
    return pos.size();
  }
 
  std::size_t lastSent() const { return sent_; }
  std::size_t lastReceived() const { return received_; }
 
  int rank() const { return rank_; }
  MPI_Comm comm() const { return comm_; }
  const decomp::BlockDecomposer<Dim>& decomposer() const { return *dec_; }
 
  bool withinRcutOfBlock(const Vec<Dim>& x, int r, double rcut, Vec<Dim>& img) const {
    const auto& o = dec_->origins()[r];
    const auto& s = dec_->sizes()[r];
    const IVec<Dim>& gsize = dec_->globalSize();
    double d2 = 0.0;
    for (int i = 0; i < Dim; ++i) {
      double lo = map_.origin[i] + o[i] * map_.cellSize[i];
      double hi = map_.origin[i] + (o[i] + s[i]) * map_.cellSize[i];
      double L = map_.cellSize[i] * static_cast<double>(gsize[i]);
      double cands[3];
      int nc = 0;
      cands[nc++] = x[i];
      if (map_.periodic[i]) {
        cands[nc++] = x[i] - L;
        cands[nc++] = x[i] + L;
      }
      double bestGap = std::numeric_limits<double>::infinity();
      double bestImg = x[i];
      for (int c = 0; c < nc; ++c) {
        double p = cands[c];
        double gap = (p < lo) ? (lo - p) : (p > hi) ? (p - hi) : 0.0;
        if (gap < bestGap) {
          bestGap = gap;
          bestImg = p;
        }
      }
      img[i] = bestImg;
      d2 += bestGap * bestGap;
    }
    return d2 < rcut * rcut;
  }
 
  bool boxWithinRcutOfBlock(const Vec<Dim>& bmin, const Vec<Dim>& bmax, int r, double rcut) const {
    const auto& o = dec_->origins()[r];
    const auto& s = dec_->sizes()[r];
    const IVec<Dim>& gsize = dec_->globalSize();
    double d2 = 0.0;
    for (int d = 0; d < Dim; ++d) {
      const double lo = map_.origin[d] + o[d] * map_.cellSize[d];
      const double hi = map_.origin[d] + (o[d] + s[d]) * map_.cellSize[d];
      auto gap = [&](double shift) {
        const double a = bmin[d] + shift, b = bmax[d] + shift;
        return std::max(0.0, std::max(lo - b, a - hi));
      };
      double g = gap(0.0);
      if (map_.periodic[d]) {
        const double L = map_.cellSize[d] * static_cast<double>(gsize[d]);
        g = std::min(g, std::min(gap(-L), gap(L)));
      }
      d2 += g * g;
    }
    return d2 < rcut * rcut;
  }
 
  void imagesWithinRcutOfBlock(const Vec<Dim>& x, int r, double rcut, bool allowIdentity,
                               std::vector<Vec<Dim>>& outShifts) const {
    const auto& o = dec_->origins()[r];
    const auto& s = dec_->sizes()[r];
    const IVec<Dim>& gsize = dec_->globalSize();
    double lo[Dim], hi[Dim], cand[Dim][3];
    int nc[Dim];
    for (int d = 0; d < Dim; ++d) {
      lo[d] = map_.origin[d] + o[d] * map_.cellSize[d];
      hi[d] = map_.origin[d] + (o[d] + s[d]) * map_.cellSize[d];
      const double L = map_.cellSize[d] * static_cast<double>(gsize[d]);
      cand[d][0] = 0.0;
      nc[d] = 1;
      if (map_.periodic[d]) {
        cand[d][1] = -L;
        cand[d][2] = L;
        nc[d] = 3;
      }
    }
    // Per-axis prefilter (exact): an image qualifies only if its total squared gap is below rc2,
    // so every one of its per-axis gaps is -- keep, per axis and in candidate order, the shifts
    // whose own gap passes, then enumerate their product in the same nesting (axis 0 fastest) as
    // the full 3^Dim walk. Same set, same order, typically one candidate per axis instead of 27
    // combinations (this runs for every owned particle x every rank at every topology build).
    const double rc2 = rcut * rcut;
    int keep[Dim][3];
    double keepGap2[Dim][3];
    int nk[Dim];
    for (int d = 0; d < Dim; ++d) {
      nk[d] = 0;
      for (int k = 0; k < nc[d]; ++k) {
        const double p = x[d] + cand[d][k];
        const double gap = (p < lo[d]) ? (lo[d] - p) : (p > hi[d]) ? (p - hi[d]) : 0.0;
        if (gap * gap < rc2) {
          keep[d][nk[d]] = k;
          keepGap2[d][nk[d]] = gap * gap;
          ++nk[d];
        }
      }
      if (nk[d] == 0)
        return;
    }
    int total = 1;
    for (int d = 0; d < Dim; ++d)
      total *= nk[d];
    for (int idx = 0; idx < total; ++idx) {
      Vec<Dim> shift{};
      int t = idx;
      bool identity = true;
      double d2 = 0.0;
      for (int d = 0; d < Dim; ++d) {
        const int j = t % nk[d];
        t /= nk[d];
        const int k = keep[d][j];
        shift[d] = cand[d][k];
        if (k != 0)
          identity = false;
        d2 += keepGap2[d][j];
      }
      if (d2 < rc2 && (allowIdentity || !identity))
        outShifts.push_back(shift);
    }
  }
 
  std::size_t gatherGhosts(const std::vector<Vec<Dim>>& pos, const std::vector<char>& payload,
                           std::size_t stride, double rcut, std::vector<Vec<Dim>>& ghostPos,
                           std::vector<char>& ghostPayload) {
    ghostPos.clear();
    ghostPayload.clear();
    const std::size_t posBytes = Dim * sizeof(double);
    const std::size_t recBytes = posBytes + stride;
    int nranks = 0;
    MPI_Comm_size(comm_, &nranks);
 
    std::map<int, std::vector<char>> outbox;
    Vec<Dim> img;
    for (std::size_t i = 0; i < pos.size(); ++i) {
      for (int r = 0; r < nranks; ++r) {
        if (r == rank_)
          continue;
        if (!withinRcutOfBlock(pos[i], r, rcut, img))
          continue;
        auto& buf = outbox[r];
        std::size_t off = buf.size();
        buf.resize(off + recBytes);
        std::memcpy(buf.data() + off, img.data(), posBytes);
        if (stride)
          std::memcpy(buf.data() + off + posBytes, &payload[i * stride], stride);
      }
    }
 
    std::vector<std::pair<int, const std::vector<char>*>> queue;
    queue.reserve(outbox.size());
    for (auto& kv : outbox)
      queue.emplace_back(kv.first, &kv.second);
    std::size_t q = 0, recv = 0;
    NbxEngine nbx(comm_);
    auto packNext = [&](std::vector<char>& out) -> int {
      if (q >= queue.size())
        return -1;
      int dest = queue[q].first;
      out = *queue[q].second;
      ++q;
      return dest;
    };
    auto onRecv = [&](int /*src*/, std::vector<char>& msg) {
      std::size_t cnt = msg.size() / recBytes;
      for (std::size_t i = 0; i < cnt; ++i) {
        const char* rec = msg.data() + i * recBytes;
        Vec<Dim> x{};
        std::memcpy(x.data(), rec, posBytes);
        ghostPos.push_back(x);
        if (stride) {
          std::size_t off = ghostPayload.size();
          ghostPayload.resize(off + stride);
          std::memcpy(&ghostPayload[off], rec + posBytes, stride);
        }
        ++recv;
      }
    };
    nbx.exchange(packNext, onRecv, /*tag=*/7402);
    return recv;
  }
 
 private:
  const decomp::BlockDecomposer<Dim>* dec_ = nullptr;
  int rank_ = 0;
  DomainMap<Dim> map_{};
  MPI_Comm comm_ = MPI_COMM_WORLD;
  std::size_t sent_ = 0, received_ = 0;
};
 
}  // namespace peclet::core::halo
 
#endif  // PECLET_CORE_HALO_PARTICLE_MIGRATOR_HPP