Make the Octree interpolation continuous.

This is done by detecting the hanging nodes and compute the function at those location by interpolating the corresponding coarse vertices.

PiperOrigin-RevId: 807700228
Change-Id: I27dcb85361ca445c2099dd07b280cae5c92d3131

src/user/user_mesh.cc

@@ -20,6 +20,7 @@
 #include <cstddef>
 #include <cstdio>
 #include <cstring>
+#include <deque>
 #include <functional>
 #include <limits>
 #include <map>
@@ -784,11 +785,44 @@ void mjCMesh::TryCompile(const mjVFS* vfs) {
       tmd::TriangleMeshDistance sdf(vert_.data(), nvert(), face_.data(), nface());
 
       std::vector<double> coeffs(octree_.NumVerts());
-      for (int i = 0; i < octree_.NumVerts(); ++i) {
-        coeffs[i] = sdf.signed_distance(octree_.Vert(i)).distance;
+      std::vector<bool> processed(octree_.NumVerts(), false);
+      std::deque<int> queue;
+
+      if (octree_.NumNodes() > 0) {
+        queue.push_back(0);  // start traversal from the root node
+      }
+
+      while (!queue.empty()) {
+        int node_idx = queue.front();
+        queue.pop_front();
+
+        for (int j = 0; j < 8; ++j) {
+          int vert_id = octree_.VertId(node_idx, j);
+          if (processed[vert_id]) {
+            continue;
+          }
+          if (octree_.Hang(vert_id).empty()) {
+            coeffs[vert_id] = sdf.signed_distance(octree_.Vert(vert_id)).distance;
+          } else {
+            double sum_coeff = 0;
+            for (int dep_id : octree_.Hang(vert_id)) {
+              sum_coeff += coeffs[dep_id];
+              if (!processed[dep_id]) {
+                throw mjCError(this, "sdf coefficient computation failed");
+              }
+            }
+            coeffs[vert_id] = sum_coeff / octree_.Hang(vert_id).size();
+          }
+          processed[vert_id] = true;
+        }
+
+        for (int child_idx : octree_.Children(node_idx)) {
+          if (child_idx != -1) {
+            queue.push_back(child_idx);
+          }
+        }
       }
 
-      // TODO: the value at hanging vertices should be computed from the parent
       for (int i = 0; i < octree_.NumNodes(); ++i) {
         for (int j = 0; j < 8; j++) {
             octree_.AddCoeff(i, j, coeffs[octree_.VertId(i, j)]);

src/user/user_objects.cc

@@ -629,6 +629,7 @@ void mjCOctree::CreateOctree(const double aamm[6]) {
                  [](Triangle& triangle) { return ▵ });
   std::unordered_map<Point, int> vert_map;
   MakeOctree(elements_ptrs, box, vert_map);
+  MarkHangingNodes();
 }
 
 
@@ -871,6 +872,101 @@ void mjCOctree::BalanceOctree(std::unordered_map<Point, int>& vert_map) {
 }
 
 
+// mark all hanging vertices in the octree
+void mjCOctree::MarkHangingNodes() {
+  hang_.assign(nvert_, std::vector<int>());
+
+  std::vector<int> leaves;
+  for (int i = 0; i < nnode_; ++i) {
+    if (node_[i].child[0] == -1) {
+      leaves.push_back(i);
+    }
+  }
+
+  for (int leaf_idx : leaves) {
+    for (int dir = 0; dir < 6; ++dir) {
+      int neighbor_idx = FindNeighbor(leaf_idx, dir);
+      if (neighbor_idx == -1 ||
+          node_[neighbor_idx].level >= node_[leaf_idx].level) {
+        continue;
+      }
+
+      // coarser neighbor found, this leaf's face has hanging nodes
+      int dim = dir / 2;
+      int side = dir % 2;
+
+      // iterate over the 4 vertices of the leaf's face
+      for (int i = 0; i < 4; ++i) {
+        // construct vertex index on the face
+        int v_idx = side << dim;
+        int d1 = (dim + 1) % 3;
+        int d2 = (dim + 2) % 3;
+        v_idx |= (i & 1) << d1;
+        v_idx |= ((i >> 1) & 1) << d2;
+
+        int hv_id = node_[leaf_idx].vertid[v_idx];
+        if (!hang_[hv_id].empty()) {
+          continue;  // already processed
+        }
+
+        const double* hv_pos = vert_[hv_id].p.data();
+        const auto& neighbor_aamm = node_[neighbor_idx].aamm;
+
+        bool is_min[3], is_max[3];
+        int on_boundary_planes = 0;
+        for (int d = 0; d < 3; ++d) {
+          is_min[d] = std::abs(hv_pos[d] - neighbor_aamm[d]) < 1e-9;
+          is_max[d] = std::abs(hv_pos[d] - neighbor_aamm[d + 3]) < 1e-9;
+          if (is_min[d] || is_max[d]) {
+            on_boundary_planes++;
+          }
+        }
+
+        if (on_boundary_planes == 2) {  // edge hanging
+          int d_mid = -1;
+          for (int d = 0; d < 3; ++d) {
+            if (!is_min[d] && !is_max[d]) {
+              d_mid = d;
+              break;
+            }
+          }
+
+          int bits[3];
+          bits[d_mid] = 0;  // this will be toggled
+          bits[(d_mid + 1) % 3] = is_max[(d_mid + 1) % 3];
+          bits[(d_mid + 2) % 3] = is_max[(d_mid + 2) % 3];
+
+          int nv_idx1 = (bits[2] << 2) | (bits[1] << 1) | bits[0];
+          bits[d_mid] = 1;
+          int nv_idx2 = (bits[2] << 2) | (bits[1] << 1) | bits[0];
+
+          hang_[hv_id].push_back(node_[neighbor_idx].vertid[nv_idx1]);
+          hang_[hv_id].push_back(node_[neighbor_idx].vertid[nv_idx2]);
+        } else if (on_boundary_planes == 1) {  // face hanging
+          int d_face = -1;
+          for (int d = 0; d < 3; ++d) {
+            if (is_min[d] || is_max[d]) {
+              d_face = d;
+              break;
+            }
+          }
+
+          int bits[3];
+          bits[d_face] = is_max[d_face];
+
+          for (int j = 0; j < 4; ++j) {
+            bits[(d_face + 1) % 3] = j & 1;
+            bits[(d_face + 2) % 3] = (j >> 1) & 1;
+            int nv_idx = (bits[2] << 2) | (bits[1] << 1) | bits[0];
+            hang_[hv_id].push_back(node_[neighbor_idx].vertid[nv_idx]);
+          }
+        }
+      }
+    }
+  }
+}
+
+
 void mjCOctree::MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6],
                            std::unordered_map<Point, int>& vert_map) {
   std::deque<OctreeTask> queue;

src/user/user_objects.h

@@ -270,8 +270,9 @@ struct mjCOctree_ {
   int nnode_ = 0;
   int nvert_ = 0;
   std::vector<OctNode> node_;
-  std::vector<Triangle> face_;       // mesh faces                (nmeshface x 3)
-  std::vector<Point> vert_;          // octree vertices           (nvert x 3)
+  std::vector<Triangle> face_;          // mesh faces                (nmeshface x 3)
+  std::vector<Point> vert_;             // octree vertices           (nvert x 3)
+  std::vector<std::vector<int>> hang_;  // hanging nodes status      (nvert x 1)
   double ipos_[3] = {0, 0, 0};
   double iquat_[4] = {1, 0, 0, 0};
 };
@@ -287,7 +288,9 @@ class mjCOctree : public mjCOctree_ {
   void CopyAabb(mjtNum* aabb) const;
   void CopyCoeff(mjtNum* coeff) const;
   const double* Vert(int i) const { return vert_[i].p.data(); }
+  const std::vector<int>& Hang(int i) const { return hang_[i]; }
   int VertId(int n, int v) const { return node_[n].vertid[v]; }
+  const std::array<int, 8>& Children(int i) const { return node_[i].child; }
   void SetFace(const std::vector<double>& vert, const std::vector<int>& face);
   int Size() const {
     return sizeof(OctNode) * node_.size() + sizeof(Triangle) * face_.size() +
@@ -310,6 +313,7 @@ class mjCOctree : public mjCOctree_ {
   int FindNeighbor(int node_idx, int dir);
   int FindCoarseNeighbor(int node_idx, int dir);
   void BalanceOctree(std::unordered_map<Point, int>& vert_map);
+  void MarkHangingNodes();
 };