block_octree.hpp

Go to the documentation of this file.

// core — per-block adaptive octree with block-local Morton codes.
//
// The suite's AMR primitive. Unlike the textbook linear octree (one global
// space-filling curve, partitioned by index ranges à la p4est/Dendro), the suite
// keeps the ORB *block* decomposition (peclet::core::decomp::BlockDecomposer) and gives
// *each block its own local Morton coordinate system*. BlockOctree is that
// per-block octree: leaves are addressed by `morton::Morton<Dim,Bits>` codes
// relative to the block origin, so codes stay narrow and the block is
// self-contained. The distributed glue (cross-block 2:1 balance + an owner-based
// ghost halo) is layered on top later (DistributedOctree / AmrHalo); a uniform,
// unrefined BlockOctree is bit-identical to the existing structured block grid.
//
// Conventions (morton's level convention; see also ../../../docs)
// -----------------------------------------------------------------------------
//  * Coordinates are in *fine units*: 1 fine unit = a leaf at level 0. A leaf at
//    `level` L covers a 2^L block per axis; its origin Morton code has the low
//    L*Dim bits zero (morton's level convention).
//  * Root cells sit at `level = lmax`; refinement decreases level toward 0
//    (finest). A block is a "root brick" of brick[i] root cells per axis, so it
//    spans brick[i]*2^lmax fine units; require 2^Bits >= brick[i]*2^lmax.
//  * Leaves tile the block without overlap and are stored sorted by code, i.e. in
//    Z-order — the same order morton's curve and the device leaf arrays use.
//
// checkout is present). The query helpers carry morton's MORTON_HD, so they are
// device-callable exactly when morton is (KOKKOS_FUNCTION under a Kokkos build);
// topology mutation (refine/coarsen/balance) is host-side and rebuilds the leaf
// arrays the device queries read — mirroring grid_halo.hpp (host) vs
// grid_halo_kokkos.hpp (device).
#ifndef PECLET_AMR_BLOCK_OCTREE_HPP
#define PECLET_AMR_BLOCK_OCTREE_HPP
 
#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <vector>
 
#include "morton/morton.hpp"
#include "peclet/amr/common.hpp"
#include "peclet/core/common/types.hpp"
 
namespace peclet::amr {
 
template <class M>
MORTON_HD inline Index amrLocate(const typename M::code_type* codes, const std::uint8_t* levels,
                                 Index n, typename M::code_type p) {
  // First index with codes[idx] > p, then step back one.
  Index lo = 0, hi = n;
  while (lo < hi) {
    Index mid = lo + ((hi - lo) >> 1);
    if (codes[mid] <= p)
      lo = mid + 1;
    else
      hi = mid;
  }
  Index idx = lo - 1;
  if (idx < 0)
    return -1;
  // p is inside leaf idx iff clearing the leaf's low bits gives its origin.
  if (M::from_code(p).ancestor(levels[idx]).code() == codes[idx])
    return idx;
  return -1;
}
 
template <int Dim, unsigned Bits = (Dim == 2 ? 32u : (Dim == 3 ? 21u : 16u))>
class BlockOctree {
 public:
  using M = morton::Morton<Dim, Bits>;
  using Code = typename M::code_type;
  using Coord = typename M::coord_type;
  static constexpr unsigned octants = M::octants;  // children per node (2^Dim)
 
  BlockOctree() = default;
 
  BlockOctree(IVec<Dim> brick, unsigned lmax, IVec<Dim> globalOrigin = IVec<Dim>{}) {
    init(brick, lmax, globalOrigin);
  }
 
  void init(IVec<Dim> brick, unsigned lmax, IVec<Dim> globalOrigin = IVec<Dim>{}) {
    brick_ = brick;
    lmax_ = lmax;
    globalOrigin_ = globalOrigin;
    const Coord rootSpan = Coord(Coord(1) << lmax_);  // fine units per root cell
    codes_.clear();
    levels_.clear();
    IVec<Dim> bgn{};
    IVec<Dim> end = brick_;
    forEachInBox<Dim>(bgn, end, [&](const IVec<Dim>& r) {
      std::array<Coord, Dim> o{};
      for (int i = 0; i < Dim; ++i)
        o[i] = static_cast<Coord>(r[i]) * rootSpan;
      codes_.push_back(M::encode(o).code());
      levels_.push_back(static_cast<std::uint8_t>(lmax_));
    });
    sortByCode();
  }
 
  void assign(IVec<Dim> brick, unsigned lmax, IVec<Dim> globalOrigin, std::vector<Code> codes,
              std::vector<std::uint8_t> levels) {
    brick_ = brick;
    lmax_ = lmax;
    globalOrigin_ = globalOrigin;
    codes_ = std::move(codes);
    levels_ = std::move(levels);
    sortByCode();
  }
 
  // ---- introspection -----------------------------------------------------
  Index numLeaves() const { return static_cast<Index>(codes_.size()); }
  unsigned lmax() const { return lmax_; }
  static constexpr unsigned bits() { return Bits; }
  const IVec<Dim>& brick() const { return brick_; }
  const IVec<Dim>& globalOrigin() const { return globalOrigin_; }
  Code code(Index i) const { return codes_[static_cast<std::size_t>(i)]; }
  unsigned level(Index i) const { return levels_[static_cast<std::size_t>(i)]; }
  const std::vector<Code>& codes() const { return codes_; }
  const std::vector<std::uint8_t>& levels() const { return levels_; }
 
  // ---- queries -----------------------------------------------------------
 
  Index find(Code p) const { return amrLocate<M>(codes_.data(), levels_.data(), numLeaves(), p); }
  Index find(const std::array<Coord, Dim>& fine) const { return find(M::encode(fine).code()); }
 
  std::array<std::array<Coord, Dim>, 2> bounds(Index i) const {
    std::array<std::array<Coord, Dim>, 2> b{};
    auto o = M::from_code(codes_[static_cast<std::size_t>(i)]).decode();
    Coord extent = Coord((Coord(1) << levels_[static_cast<std::size_t>(i)]) - 1);
    for (int d = 0; d < Dim; ++d) {
      b[0][d] = o[d];
      b[1][d] = Coord(o[d] + extent);
    }
    return b;
  }
 
  Index faceNeighbor(Index i, int axis, int dir) const {
    M probe = M::from_code(codes_[static_cast<std::size_t>(i)]);
    const Coord step = Coord(Coord(1) << levels_[static_cast<std::size_t>(i)]);
    if (dir >= 0) {
      if (!probe.try_add(static_cast<unsigned>(axis), step))
        return -1;  // block edge
    } else {
      if (!probe.try_sub(static_cast<unsigned>(axis), 1))
        return -1;  // one past lower face
    }
    return find(probe.code());
  }
 
  // ---- refinement / coarsening ------------------------------------------
 
  template <class Pred>
  Index refineIf(Pred&& pred) {
    std::vector<Code> nc;
    std::vector<std::uint8_t> nl;
    nc.reserve(codes_.size());
    nl.reserve(levels_.size());
    Index nsplit = 0;
    for (std::size_t k = 0; k < codes_.size(); ++k) {
      const unsigned L = levels_[k];
      const Code c = codes_[k];
      if (L > 0 && pred(c, L)) {
        M parent = M::from_code(c);
        for (unsigned oct = 0; oct < octants; ++oct) {
          nc.push_back(parent.child(L, oct).code());
          nl.push_back(static_cast<std::uint8_t>(L - 1));
        }
        ++nsplit;
      } else {
        nc.push_back(c);
        nl.push_back(static_cast<std::uint8_t>(L));
      }
    }
    codes_.swap(nc);
    levels_.swap(nl);
    return nsplit;
  }
 
  bool refineLeaf(Index i) {
    const Code target = codes_[static_cast<std::size_t>(i)];
    return refineIf([&](Code c, unsigned) { return c == target; }) > 0;
  }
 
  template <class Pred>
  Index coarsenIf(Pred&& pred) {
    std::vector<Code> nc;
    std::vector<std::uint8_t> nl;
    nc.reserve(codes_.size());
    nl.reserve(levels_.size());
    Index nmerged = 0;
    std::size_t k = 0;
    while (k < codes_.size()) {
      const unsigned L = levels_[k];
      bool merged = false;
      if (L < lmax_ && k + octants <= codes_.size()) {
        const Code parent = M::from_code(codes_[k]).ancestor(L + 1).code();
        bool full = (codes_[k] == parent);  // octant-0 child shares the parent origin
        for (unsigned oct = 0; full && oct < octants; ++oct) {
          full = levels_[k + oct] == L &&
                 M::from_code(codes_[k + oct]).ancestor(L + 1).code() == parent &&
                 M::from_code(parent).child(L + 1, oct).code() == codes_[k + oct];
        }
        if (full && pred(parent, L + 1)) {
          nc.push_back(parent);
          nl.push_back(static_cast<std::uint8_t>(L + 1));
          k += octants;
          ++nmerged;
          merged = true;
        }
      }
      if (!merged) {
        nc.push_back(codes_[k]);
        nl.push_back(static_cast<std::uint8_t>(L));
        ++k;
      }
    }
    codes_.swap(nc);
    levels_.swap(nl);
    return nmerged;
  }
 
  // ---- root lift (multigrid below the root brick) -------------------------
 
  bool canLiftRoot() const {
    for (int d = 0; d < Dim; ++d)
      if ((brick_[d] % 2) != 0 || (globalOrigin_[d] % 2) != 0)
        return false;
    return true;
  }
 
  void liftRoot() {
    assert(canLiftRoot());
    for (int d = 0; d < Dim; ++d) {
      brick_[d] /= 2;
      globalOrigin_[d] /= 2;
    }
    ++lmax_;
  }
 
  Index balance2to1() {
    Index total = 0;
    for (;;) {
      std::vector<Code> toRefine;  // origins of leaves to split this round
      for (std::size_t k = 0; k < codes_.size(); ++k) {
        const unsigned Lf = levels_[k];
        for (int axis = 0; axis < Dim; ++axis) {
          for (int dir = -1; dir <= 1; dir += 2) {
            Index nb = faceNeighbor(static_cast<Index>(k), axis, dir);
            if (nb >= 0 && levels_[static_cast<std::size_t>(nb)] >= Lf + 2)
              toRefine.push_back(codes_[static_cast<std::size_t>(nb)]);
          }
        }
      }
      if (toRefine.empty())
        break;
      std::sort(toRefine.begin(), toRefine.end());
      toRefine.erase(std::unique(toRefine.begin(), toRefine.end()), toRefine.end());
      total += refineIf([&](Code c, unsigned) {
        return std::binary_search(toRefine.begin(), toRefine.end(), c);
      });
    }
    return total;
  }
 
  bool isBalanced() const {
    for (std::size_t k = 0; k < codes_.size(); ++k) {
      const unsigned Lf = levels_[k];
      for (int axis = 0; axis < Dim; ++axis)
        for (int dir = -1; dir <= 1; dir += 2) {
          Index nb = faceNeighbor(static_cast<Index>(k), axis, dir);
          if (nb >= 0) {
            unsigned Ln = levels_[static_cast<std::size_t>(nb)];
            unsigned hi = Lf > Ln ? Lf : Ln, lo = Lf > Ln ? Ln : Lf;
            if (hi - lo > 1)
              return false;
          }
        }
    }
    return true;
  }
 
 private:
  void sortByCode() {
    std::vector<Index> ord(codes_.size());
    for (std::size_t i = 0; i < ord.size(); ++i)
      ord[i] = static_cast<Index>(i);
    std::sort(ord.begin(), ord.end(), [&](Index a, Index b) { return codes_[a] < codes_[b]; });
    std::vector<Code> nc(codes_.size());
    std::vector<std::uint8_t> nl(levels_.size());
    for (std::size_t i = 0; i < ord.size(); ++i) {
      nc[i] = codes_[static_cast<std::size_t>(ord[i])];
      nl[i] = levels_[static_cast<std::size_t>(ord[i])];
    }
    codes_.swap(nc);
    levels_.swap(nl);
  }
 
  IVec<Dim> brick_{};
  IVec<Dim> globalOrigin_{};
  unsigned lmax_ = 0;
  std::vector<Code> codes_;           // leaf origin codes, ascending (Z-order)
  std::vector<std::uint8_t> levels_;  // parallel to codes_
};
 
}  // namespace peclet::amr
 
#endif  // PECLET_AMR_BLOCK_OCTREE_HPP