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

Author: quagla

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();
 };
 
 

test/user/user_mesh_test.cc

@@ -1366,6 +1366,167 @@ TEST_F(MjCMeshTest, OctreeIsBalanced) {
   mj_deleteModel(model);
 }
 
+TEST_F(MjCMeshTest, OctreeHangingNodeInterpolation) {
+  const std::string xml_path = GetTestDataFilePath(kTorusPath);
+  std::array<char, 1024> error;
+  mjSpec* spec = mj_parseXML(xml_path.c_str(), 0, error.data(), error.size());
+  mjsGeom* geom = mjs_asGeom(mjs_firstElement(spec, mjOBJ_GEOM));
+  geom->type = mjGEOM_SDF;
+  mjModel* model = mj_compile(spec, 0);
+  ASSERT_THAT(model, NotNull()) << error.data();
+  EXPECT_GT(model->mesh_octnum[0], 0);
+  double kEps = 1e-6;
+
+  const int octree_adr = model->mesh_octadr[0];
+  const int noct = model->mesh_octnum[0];
+  const mjtNum* sdf = model->oct_coeff + octree_adr * 8;
+
+  // find all leaves in the octree
+  std::vector<int> leaves;
+  for (int i = 0; i < noct; ++i) {
+    bool is_leaf = true;
+    for (int j = 0; j < 8; ++j) {
+      if (model->oct_child[(octree_adr + i) * 8 + j] != -1) {
+        is_leaf = false;
+        break;
+      }
+    }
+    if (is_leaf) {
+      leaves.push_back(i);
+    }
+  }
+
+  // do a n^2 check of all pairs of leaves in the octree
+  // for each pair, check if they are adjacent and if so, check that all hanging
+  // nodes within the octree can be interpolated from their parent nodes
+  int hanging_nodes_checked = 0;
+  int interpolation_failures = 0;
+  for (int i = 0; i < leaves.size(); ++i) {
+    for (int j = i + 1; j < leaves.size(); ++j) {
+      const int node1_idx = leaves[i];
+      const int node2_idx = leaves[j];
+      const mjtNum* aabb1 = &model->oct_aabb[(octree_adr + node1_idx) * 6];
+      const mjtNum* aabb2 = &model->oct_aabb[(octree_adr + node2_idx) * 6];
+
+      if (AreAabbsAdjacent(aabb1, aabb2)) {
+        const int level1 = model->oct_depth[octree_adr + node1_idx];
+        const int level2 = model->oct_depth[octree_adr + node2_idx];
+
+        if (level1 == level2) {
+          continue;
+        }
+
+        // decide which node is finer and which is coarser
+        const int finer_node_idx = (level1 > level2) ? node1_idx : node2_idx;
+        const int coarser_node_idx =
+            (level1 > level2) ? node2_idx : node1_idx;
+        const mjtNum* coarser_aabb =
+            &model->oct_aabb[(octree_adr + coarser_node_idx) * 6];
+        const mjtNum* finer_aabb =
+            &model->oct_aabb[(octree_adr + finer_node_idx) * 6];
+
+        mjtNum coarser_corners[8][3];
+        for (int c = 0; c < 8; ++c) {
+          int sx = (c & 1) ? 1 : -1;
+          int sy = (c & 2) ? 1 : -1;
+          int sz = (c & 4) ? 1 : -1;
+          coarser_corners[c][0] = coarser_aabb[0] + sx * coarser_aabb[3];
+          coarser_corners[c][1] = coarser_aabb[1] + sy * coarser_aabb[4];
+          coarser_corners[c][2] = coarser_aabb[2] + sz * coarser_aabb[5];
+        }
+
+        // for all vertices in the finer node, check if they are hanging and
+        // can be interpolated
+        for (int v_idx = 0; v_idx < 8; ++v_idx) {
+          mjtNum v_pos[3];
+          int sx = (v_idx & 1) ? 1 : -1;
+          int sy = (v_idx & 2) ? 1 : -1;
+          int sz = (v_idx & 4) ? 1 : -1;
+          v_pos[0] = finer_aabb[0] + sx * finer_aabb[3];
+          v_pos[1] = finer_aabb[1] + sy * finer_aabb[4];
+          v_pos[2] = finer_aabb[2] + sz * finer_aabb[5];
+
+          // skip finer vertices that are also coarse corners
+          bool is_coarse_corner = false;
+          for (int c = 0; c < 8; ++c) {
+            if (mju_dist3(v_pos, coarser_corners[c]) < 1e-6) {
+              is_coarse_corner = true;
+              break;
+            }
+          }
+          if (is_coarse_corner) {
+            continue;
+          }
+
+          // skip vertices that are not on the boundary
+          mjtNum p_local[3];
+          bool outside = false;
+          for (int d = 0; d < 3; ++d) {
+            p_local[d] = (v_pos[d] - coarser_aabb[d]) / coarser_aabb[d + 3];
+            if (std::abs(p_local[d]) > 1.0 + kEps) {
+              outside = true;
+              break;
+            }
+          }
+          if (outside) {
+            continue;
+          }
+
+          // count the number of dimensions that are on the boundary
+          int num_dim = 0;
+          for (int d = 0; d < 3; ++d) {
+            if (std::abs(p_local[d] - 1.0) < kEps) {
+              num_dim++;
+            } else if (std::abs(p_local[d] + 1.0) < kEps) {
+              num_dim++;
+            }
+          }
+
+          double interpolated_sdf = 0;
+          const mjtNum* coarser_sdf = sdf + coarser_node_idx * 8;
+
+          // for edge or face nodes, try to interpolate the hanging nodes
+          if (num_dim == 1 || num_dim == 2) {
+            for (int k = 0; k < 8; ++k) {
+              int sx = (k & 1) ? 1 : -1;
+              int sy = (k & 2) ? 1 : -1;
+              int sz = (k & 4) ? 1 : -1;
+              double weight = (1 + p_local[0] * sx) / 2.0 *
+                              (1 + p_local[1] * sy) / 2.0 *
+                              (1 + p_local[2] * sz) / 2.0;
+              if (weight > kEps && num_dim == 1) {
+                ASSERT_NEAR(weight, 0.25, kEps);
+              } else if (weight > kEps && num_dim == 2) {
+                ASSERT_NEAR(weight, 0.5, kEps);
+              }
+              interpolated_sdf += weight * coarser_sdf[k];
+            }
+          } else {
+            continue;
+          }
+
+          // if the values do not match, log an error
+          const mjtNum* finer_sdf = sdf + finer_node_idx * 8;
+          if (std::abs(finer_sdf[v_idx] - interpolated_sdf) > kEps) {
+            if (interpolation_failures < 10) {
+              EXPECT_NEAR(finer_sdf[v_idx], interpolated_sdf, kEps);
+            }
+            interpolation_failures++;
+          }
+          hanging_nodes_checked++;
+        }
+      }
+    }
+  }
+  EXPECT_GT(hanging_nodes_checked, 0);
+  EXPECT_EQ(interpolation_failures, 0)
+      << "Found " << interpolation_failures
+      << " hanging node interpolation failures.";
+
+  mj_deleteSpec(spec);
+  mj_deleteModel(model);
+}
+
 TEST_F(MjCMeshTest, OctreeNotComputedForNonSDF) {
   const std::string xml_path = GetTestDataFilePath(kTorusPath);
   std::array<char, 1024> error;