scene.hpp

Go to the documentation of this file.

// core — runtime-dispatched SDF scene nodes: the layer that turns the compile-time leaves into
// geometry a solver can be handed at run time.
//
// Layer 0 rung 1 of suite/docs/ANALYTIC_SDF_GEOMETRY.md. This header contains NO signed-distance
// formulas: every leaf case dispatches into peclet::core::geom::prim (geom/primitives.hpp), and the
// only arithmetic here is the conformal transform stack and the CSG combinators. voro keeps its
// zero-overhead compile-time provider path over the leaves; dem and flow get the runtime scene
// their Python-composed, per-particle heterogeneity requires.
//
// WHY RUNTIME DISPATCH IS FREE ENOUGH: dem already pays a per-probe `switch` on shape kind in its
// hottest kernel today (narrowphase.hpp sdfEvalShape), and a simulation carries few distinct
// shapes, so the branch is warp-coherent. flow's geometry derivation is setup-phase, where dispatch
// cost is irrelevant.
//
// TRANSFORMS ARE CONFORMAL ONLY (contract 1): translation + quaternion + ISOTROPIC scale. A
// non-uniform scale destroys the signed-distance property (d_world = s * d_canonical holds only for
// conformal maps); anisotropy belongs in primitive parameters (box half-extents, ellipsoid
// semi-axes). This is exactly dem's per-particle (pos, quat, scale x globalScale) contract, which
// is what makes the rung-3 port a relocation.
//
// C++17-clean so it can be pulled into CUDA translation units.
#ifndef PECLET_CORE_GEOM_SCENE_HPP
#define PECLET_CORE_GEOM_SCENE_HPP
 
#include "peclet/core/common/portable.hpp"
#include "peclet/core/geom/primitives.hpp"
 
namespace peclet::core::geom {
 
enum ShapeKind : int {
  kGrid = 0,                 
  kSphere = 1,               
  kHollowCylinder = 2,       
  kBox = 3,                  
  kHollowCylinderShell = 4,  
  kCapsule = 5,              
  kTorus = 6,                
  kCone = 7,                 
  kEllipsoid = 8,            
  kSuperquadric = 9,         
 
  kCsgBase = 32,
  kUnion = 32,         
  kIntersection = 33,  
  kDifference = 34,    
};
 
 
enum class GridExtension : int { kObject = 0, kContainer = 1, kClamp = 2 };
 
template <class Real>
struct GridDesc {
  int nx = 0, ny = 0, nz = 0;
  int offset = 0;  
  Vec3<Real> origin{0, 0, 0};
  Vec3<Real> invSpacing{0, 0, 0};  
  GridExtension extension = GridExtension::kObject;
};
 
template <class Real>
struct Transform {
  Vec3<Real> translation{0, 0, 0};
  Quat<Real> rotation{0, 0, 0, 1};
  Real scale = Real(1);
};
 
template <class Real>
PECLET_HD Vec3<Real> toCanonical(const Transform<Real>& tr, Vec3<Real> p) {
  return scale(invRotate(tr.rotation, sub(p, tr.translation)), Real(1) / tr.scale);
}
 
template <class Real>
struct ShapeNode {
  int kind = kSphere;
  int aux0 = -1;  
  int aux1 = -1;  
  Real params[8] = {Real(1), 0, 0, 0, 0, 0, 0, 0};
  Transform<Real> transform{};
};
 
template <class T>
struct PoolPtr {
  const T* data = nullptr;
  PECLET_HD T operator()(long i) const { return data[i]; }
};
template <class T>
struct TablePtr {
  const T* data = nullptr;
  PECLET_HD const T& operator()(long i) const { return data[i]; }
};
 
template <class Real, class Pool>
PECLET_HD Real sampleGrid(Vec3<Real> p, const GridDesc<Real>& d, const Pool& pool) {
  const Real fx = (p.x - d.origin.x) * d.invSpacing.x;
  const Real fy = (p.y - d.origin.y) * d.invSpacing.y;
  const Real fz = (p.z - d.origin.z) * d.invSpacing.z;
  const Real cx = detail::hdMin(detail::hdMax(fx, Real(0)), Real(d.nx - 1));
  const Real cy = detail::hdMin(detail::hdMax(fy, Real(0)), Real(d.ny - 1));
  const Real cz = detail::hdMin(detail::hdMax(fz, Real(0)), Real(d.nz - 1));
  const int ix = (int)cx, iy = (int)cy, iz = (int)cz;
  const int ix1 = ix < d.nx - 1 ? ix + 1 : ix;
  const int iy1 = iy < d.ny - 1 ? iy + 1 : iy;
  const int iz1 = iz < d.nz - 1 ? iz + 1 : iz;
  const Real tx = cx - ix, ty = cy - iy, tz = cz - iz;
  const long nxny = (long)d.nx * d.ny;
  const int off = d.offset;
  const auto at = [&](int x, int y, int z) {
    return (Real)pool(off + (long)z * nxny + (long)y * d.nx + x);
  };
  const Real c00 = at(ix, iy, iz) * (1 - tx) + at(ix1, iy, iz) * tx;
  const Real c10 = at(ix, iy1, iz) * (1 - tx) + at(ix1, iy1, iz) * tx;
  const Real c01 = at(ix, iy, iz1) * (1 - tx) + at(ix1, iy, iz1) * tx;
  const Real c11 = at(ix, iy1, iz1) * (1 - tx) + at(ix1, iy1, iz1) * tx;
  const Real c0 = c00 * (1 - ty) + c10 * ty;
  const Real c1 = c01 * (1 - ty) + c11 * ty;
  const Real val = c0 * (1 - tz) + c1 * tz;
  if (d.extension == GridExtension::kClamp)
    return val;  // flat at the box face -- skip the residual (and its three divisions) entirely
  const Real rx = (d.invSpacing.x > Real(0)) ? (fx - cx) / d.invSpacing.x : Real(0);
  const Real ry = (d.invSpacing.y > Real(0)) ? (fy - cy) / d.invSpacing.y : Real(0);
  const Real rz = (d.invSpacing.z > Real(0)) ? (fz - cz) / d.invSpacing.z : Real(0);
  const Real residual = detail::hdSqrt(rx * rx + ry * ry + rz * rz);
  return (d.extension == GridExtension::kObject) ? val + residual : val - residual;
}
 
template <class Real, class Grids, class Pool>
PECLET_HD Real evalLeaf(const ShapeNode<Real>& n, Vec3<Real> p, const Grids& grids,
                        const Pool& pool) {
  switch (n.kind) {
    case kSphere:
      return prim::Sphere<Real>{n.params[0]}.eval(p);
    case kBox:
      return prim::Box<Real>{n.params[0], n.params[1], n.params[2]}.eval(p);
    case kHollowCylinder:
      return prim::HollowCylinder<Real>{n.params[0], n.params[1], n.params[2]}.eval(p);
    case kHollowCylinderShell:
      return prim::HollowCylinderShell<Real>{n.params[0], n.params[1], n.params[2]}.eval(p);
    case kCapsule:
      return prim::Capsule<Real>{n.params[0], n.params[1]}.eval(p);
    case kTorus:
      return prim::Torus<Real>{n.params[0], n.params[1]}.eval(p);
    case kCone:
      return prim::Cone<Real>{n.params[0], n.params[1], n.params[2]}.eval(p);
    case kEllipsoid:
      return prim::Ellipsoid<Real>{n.params[0], n.params[1], n.params[2]}.eval(p);
    case kSuperquadric:
      return prim::Superquadric<Real>{n.params[0], n.params[1], n.params[2], n.params[3]}.eval(p);
    case kGrid:
      return sampleGrid(p, grids(n.aux0), pool);
    default:
      return Real(1e9);  // unknown kind: same "infinitely far" sentinel dem's sdfEval returns
  }
}
 
template <class Real>
struct Instance {
  int shapeRoot = 0;    
  int materialId = -1;  
  Transform<Real> transform{};
  Vec3<Real> linVel{0, 0, 0};
  Vec3<Real> angVel{0, 0, 0};
  Vec3<Real> center{0, 0, 0};  
  bool centerPinned =
      false;  
};
 
template <class Real>
PECLET_HD Vec3<Real> instanceCentre(const Instance<Real>& inst) {
  return inst.centerPinned ? inst.center : inst.transform.translation;
}
template <class Real>
PECLET_HD Vec3<Real> instanceVelocity(const Instance<Real>& inst, Vec3<Real> p) {
  return add(inst.linVel, cross(inst.angVel, sub(p, instanceCentre(inst))));
}
 
template <class Real>
struct SceneView {
  const ShapeNode<Real>* nodes = nullptr;
  int nodeCount = 0;
  const GridDesc<Real>* grids = nullptr;
  int gridCount = 0;
  const float* samples = nullptr;
  long sampleCount = 0;
  const Instance<Real>* instances = nullptr;
  int instanceCount = 0;
};
 
constexpr int kMaxTreeDepth = 16;
 
template <class Real, class Nodes, class Grids, class Pool>
PECLET_HD Real evalTree(const Nodes& nodes, int nodeCount, int root, Vec3<Real> p,
                        const Grids& grids, const Pool& pool) {
  struct Frame {
    int node;
    int stage;  
    Vec3<Real> can;
    Real scale;
    Real left;
  };
  Frame stack[kMaxTreeDepth];
  int sp = 0;
  Real result = Real(1e9);
 
  if (root < 0 || root >= nodeCount)
    return result;
 
  {
    const auto& n = nodes(root);
    stack[0] = Frame{root, 0, toCanonical(n.transform, p), n.transform.scale, Real(0)};
    sp = 1;
  }
 
  while (sp > 0) {
    Frame& f = stack[sp - 1];
    const auto& n = nodes(f.node);
 
    if (n.kind < kCsgBase) {  // leaf: evaluate and unwind
      result = f.scale * evalLeaf(n, f.can, grids, pool);
      --sp;
    } else if (f.stage == 0) {  // descend into the left child
      if (n.aux0 < 0 || n.aux0 >= nodeCount || sp >= kMaxTreeDepth)
        return Real(1e9);
      f.stage = 1;
      const auto& c = nodes(n.aux0);
      stack[sp] = Frame{n.aux0, 0, toCanonical(c.transform, f.can), c.transform.scale, Real(0)};
      ++sp;
      continue;
    } else if (f.stage == 1) {  // left value in hand; descend into the right child
      if (n.aux1 < 0 || n.aux1 >= nodeCount || sp >= kMaxTreeDepth)
        return Real(1e9);
      f.left = result;
      f.stage = 2;
      const auto& c = nodes(n.aux1);
      stack[sp] = Frame{n.aux1, 0, toCanonical(c.transform, f.can), c.transform.scale, Real(0)};
      ++sp;
      continue;
    } else {  // both children evaluated: combine
      const Real l = f.left, r = result;
      Real combined;
      if (n.kind == kUnion)
        combined = detail::hdMin(l, r);
      else if (n.kind == kIntersection)
        combined = detail::hdMax(l, r);
      else  // kDifference: L minus R
        combined = detail::hdMax(l, -r);
      result = f.scale * combined;
      --sp;
    }
  }
  return result;
}
 
template <class Real, class Grids, class Pool>
PECLET_HD Vec3<Real> leafGrad(const ShapeNode<Real>& n, Vec3<Real> p, const Grids& grids,
                              const Pool& pool) {
  switch (n.kind) {
    case kSphere:
      return prim::Sphere<Real>{n.params[0]}.grad(p);
    case kBox:
      return prim::Box<Real>{n.params[0], n.params[1], n.params[2]}.grad(p);
    case kHollowCylinder:
      return prim::HollowCylinder<Real>{n.params[0], n.params[1], n.params[2]}.grad(p);
    case kHollowCylinderShell:
      return prim::HollowCylinderShell<Real>{n.params[0], n.params[1], n.params[2]}.grad(p);
    case kCapsule:
      return prim::Capsule<Real>{n.params[0], n.params[1]}.grad(p);
    case kTorus:
      return prim::Torus<Real>{n.params[0], n.params[1]}.grad(p);
    case kCone:
      return prim::Cone<Real>{n.params[0], n.params[1], n.params[2]}.grad(p);
    case kEllipsoid:
      return prim::Ellipsoid<Real>{n.params[0], n.params[1], n.params[2]}.grad(p);
    case kSuperquadric:
      return prim::Superquadric<Real>{n.params[0], n.params[1], n.params[2], n.params[3]}.grad(p);
    case kGrid: {
      const GridDesc<Real>& d = grids(n.aux0);
      const Real hx = Real(0.5) / d.invSpacing.x, hy = Real(0.5) / d.invSpacing.y,
                 hz = Real(0.5) / d.invSpacing.z;
      return Vec3<Real>{(sampleGrid(Vec3<Real>{p.x + hx, p.y, p.z}, d, pool) -
                         sampleGrid(Vec3<Real>{p.x - hx, p.y, p.z}, d, pool)) /
                            (Real(2) * hx),
                        (sampleGrid(Vec3<Real>{p.x, p.y + hy, p.z}, d, pool) -
                         sampleGrid(Vec3<Real>{p.x, p.y - hy, p.z}, d, pool)) /
                            (Real(2) * hy),
                        (sampleGrid(Vec3<Real>{p.x, p.y, p.z + hz}, d, pool) -
                         sampleGrid(Vec3<Real>{p.x, p.y, p.z - hz}, d, pool)) /
                            (Real(2) * hz)};
    }
    default:
      return Vec3<Real>{0, 0, 0};
  }
}
 
template <class Real, class Nodes, class Grids, class Pool>
PECLET_HD Real evalTreeGrad(const Nodes& nodes, int nodeCount, int root, Vec3<Real> p,
                            const Grids& grids, const Pool& pool, Vec3<Real>& grad) {
  struct Frame {
    int node;
    int stage;
    Vec3<Real> can;
    Real scale;
    Real left;
    Vec3<Real> leftGrad;
  };
  Frame stack[kMaxTreeDepth];
  int sp = 0;
  Real result = Real(1e9);
  Vec3<Real> resultGrad{0, 0, 0};
  grad = resultGrad;
 
  if (root < 0 || root >= nodeCount)
    return result;
 
  {
    const auto& n = nodes(root);
    stack[0] = Frame{
        root, 0, toCanonical(n.transform, p), n.transform.scale, Real(0), Vec3<Real>{0, 0, 0}};
    sp = 1;
  }
 
  while (sp > 0) {
    Frame& f = stack[sp - 1];
    const auto& n = nodes(f.node);
 
    if (n.kind < kCsgBase) {  // leaf: value + closed-form gradient, rotated into the parent frame
      result = f.scale * evalLeaf(n, f.can, grids, pool);
      resultGrad = rotate(n.transform.rotation, leafGrad(n, f.can, grids, pool));
      --sp;
    } else if (f.stage == 0) {
      if (n.aux0 < 0 || n.aux0 >= nodeCount || sp >= kMaxTreeDepth) {
        grad = Vec3<Real>{0, 0, 0};
        return Real(1e9);
      }
      f.stage = 1;
      const auto& c = nodes(n.aux0);
      stack[sp] = Frame{n.aux0,
                        0,
                        toCanonical(c.transform, f.can),
                        c.transform.scale,
                        Real(0),
                        Vec3<Real>{0, 0, 0}};
      ++sp;
      continue;
    } else if (f.stage == 1) {
      if (n.aux1 < 0 || n.aux1 >= nodeCount || sp >= kMaxTreeDepth) {
        grad = Vec3<Real>{0, 0, 0};
        return Real(1e9);
      }
      f.left = result;
      f.leftGrad = resultGrad;
      f.stage = 2;
      const auto& c = nodes(n.aux1);
      stack[sp] = Frame{n.aux1,
                        0,
                        toCanonical(c.transform, f.can),
                        c.transform.scale,
                        Real(0),
                        Vec3<Real>{0, 0, 0}};
      ++sp;
      continue;
    } else {  // combine: same value expressions as evalTree; gradient = the active branch's
      const Real l = f.left, r = result;
      Real combined;
      Vec3<Real> g;
      if (n.kind == kUnion) {
        combined = detail::hdMin(l, r);
        g = (l <= r) ? f.leftGrad : resultGrad;
      } else if (n.kind == kIntersection) {
        combined = detail::hdMax(l, r);
        g = (l >= r) ? f.leftGrad : resultGrad;
      } else {  // kDifference: max(l, -r); the right branch's gradient enters NEGATED
        combined = detail::hdMax(l, -r);
        g = (l >= -r) ? f.leftGrad : Vec3<Real>{-resultGrad.x, -resultGrad.y, -resultGrad.z};
      }
      result = f.scale * combined;
      resultGrad = rotate(n.transform.rotation, g);
      --sp;
    }
  }
  grad = resultGrad;
  return result;
}
 
template <class Real>
PECLET_HD Real evalInstance(const SceneView<Real>& sc, int i, Vec3<Real> p) {
  if (i < 0 || i >= sc.instanceCount)
    return Real(1e9);
  const Instance<Real>& inst = sc.instances[i];
  const Vec3<Real> q = toCanonical(inst.transform, p);
  const Real d = evalTree<Real>(TablePtr<ShapeNode<Real>>{sc.nodes}, sc.nodeCount, inst.shapeRoot,
                                q, TablePtr<GridDesc<Real>>{sc.grids}, PoolPtr<float>{sc.samples});
  return inst.transform.scale * d;
}
 
template <class Real>
PECLET_HD Real evalCandidates(const SceneView<Real>& sc, Vec3<Real> p, const int* candidates,
                              int nCandidates) {
  Real m = Real(1e9);
  for (int k = 0; k < nCandidates; ++k) {
    const Real d = evalInstance(sc, candidates[k], p);
    if (d < m)
      m = d;
  }
  return m;
}
 
template <class Real>
PECLET_HD Real evalScene(const SceneView<Real>& sc, Vec3<Real> p) {
  Real m = Real(1e9);
  for (int i = 0; i < sc.instanceCount; ++i) {
    const Real d = evalInstance(sc, i, p);
    if (d < m)
      m = d;
  }
  return m;
}
 
}  // namespace peclet::core::geom
 
#endif  // PECLET_CORE_GEOM_SCENE_HPP