From ffdfe542cfd5296f7cfdd05a026299169b1b6c3a Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 16 Oct 2025 04:38:15 -0700
Subject: [PATCH] Add mju_interpolate3D.

This is a more general API in preparation of a higher interpolation order

PiperOrigin-RevId: 820174741
Change-Id: I81bd52ab44e0485d72579e21d1c8a48a1b35ec3f
---
 src/engine/engine_core_constraint.c |  4 +---
 src/engine/engine_core_smooth.c     | 13 ++++++-------
 src/engine/engine_util_misc.c       | 17 +++++++++++++++++
 src/engine/engine_util_misc.h       |  6 ++++++
 src/engine/engine_vis_interact.c    | 19 ++++++++++++++-----
 5 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c
index 45296ecc..7ebacddb 100644
--- a/src/engine/engine_core_constraint.c
+++ b/src/engine/engine_core_constraint.c
@@ -186,9 +186,7 @@ static int mj_vertBodyWeight(const mjModel* m, const mjData* d, int f, int v,
   int nb = 0;
 
   for (int i = nstart; i < nend; i++) {
-    mjtNum w = ((i-nstart)&1 ? coord[2] : 1-coord[2]) *
-               ((i-nstart)&2 ? coord[1] : 1-coord[1]) *
-               ((i-nstart)&4 ? coord[0] : 1-coord[0]);
+    mjtNum w = mju_evalBasis(coord, i-nstart, m->flex_interp[f]);
     if (w < 1e-5) {
       continue;
     }
diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
index d55b9483..513886bb 100644
--- a/src/engine/engine_core_smooth.c
+++ b/src/engine/engine_core_smooth.c
@@ -475,15 +475,14 @@ void mj_flex(const mjModel* m, mjData* d) {
         }
       }
 
+      int order = m->flex_interp[f];
+      if (nend - nstart != (order + 1) * (order + 1) * (order + 1)) {
+        mjERROR("flex_interp_order mismatch");
+      }
+
       for (int i=vstart; i < vend; i++) {
         mju_zero3(d->flexvert_xpos+3*i);
-        mjtNum* coord = m->flex_vert0 + 3*i;
-        for (int j=0; j < nend-nstart; j++) {
-          mjtNum coef = (j&1 ? coord[2] : 1-coord[2]) *
-                        (j&2 ? coord[1] : 1-coord[1]) *
-                        (j&4 ? coord[0] : 1-coord[0]);
-          mju_addToScl3(d->flexvert_xpos+3*i, nodexpos+3*j, coef);
-        }
+        mju_interpolate3D(d->flexvert_xpos+3*i, m->flex_vert0 + 3*i, nodexpos, order);
       }
     }
   }
diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c
index 010b25c0..fec5cf01 100644
--- a/src/engine/engine_util_misc.c
+++ b/src/engine/engine_util_misc.c
@@ -537,6 +537,23 @@ void mju_defGradient(mjtNum res[9], const mjtNum p[3], const mjtNum* dof, int or
   }
 }
 
+// evaluate the basis function at x for the i-th node
+mjtNum mju_evalBasis(const mjtNum x[3], int i, int order) {
+  if (order > 1) {
+    mjERROR("mju_evalBasis: order must be <= 1");
+    return -1;
+  }
+  return phi(x[2], i&1) * phi(x[1], i&2) * phi(x[0], i&4);
+}
+
+// interpolate a function at x with given interpolation coefficients and order n
+void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order) {
+  int npoint = (order + 1) * (order + 1) * (order + 1);
+  for (int j=0; j < npoint; j++) {
+    mju_addToScl3(res, coeff+3*j, mju_evalBasis(x, j, order));
+  }
+}
+
 
 //------------------------------ actuator models ---------------------------------------------------
 
diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h
index b682c4d0..e59aaed9 100644
--- a/src/engine/engine_util_misc.h
+++ b/src/engine/engine_util_misc.h
@@ -62,6 +62,12 @@ int mju_insideGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3
 // evaluate the deformation gradient at p using the nodal dof values
 MJAPI void mju_defGradient(mjtNum res[9], const mjtNum p[3], const mjtNum* dof, int order);
 
+// evaluate the basis function at x for the i-th node
+MJAPI mjtNum mju_evalBasis(const mjtNum x[3], int i, int order);
+
+// interpolate a function at x with given interpolation coefficients and order n
+MJAPI void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order);
+
 // ----------------------------- Base64 -----------------------------------------------------------
 
 // encode data as Base64 into buf (including padding and null char)
diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c
index 537efb47..b89cb45a 100644
--- a/src/engine/engine_vis_interact.c
+++ b/src/engine/engine_vis_interact.c
@@ -861,11 +861,20 @@ int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
       if (newdist >= 0 && (newdist < flexdist || flexdist < 0)) {
         flexdist = newdist;
         if (m->flex_interp[i]) {
-          mjtNum* vert0 = m->flex_vert0 + 3*(m->flex_vertadr[i] + vertid);
-          int l = vert0[0] > 0.5 ? 1 : 0;
-          int j = vert0[1] > 0.5 ? 1 : 0;
-          int k = vert0[2] > 0.5 ? 1 : 0;
-          int nodeid = 4*l+2*j+k;
+          mjtNum* coord = m->flex_vert0 + 3*(m->flex_vertadr[i] + vertid);
+          int nodeid = -1;
+          int nstart = m->flex_nodeadr[i];
+          int nend = nstart + m->flex_nodenum[i];
+          mjtNum w = 0;
+          for (int j = nstart; j < nend; j++) {
+            if (mju_evalBasis(coord, j-nstart, m->flex_interp[i]) > w) {
+              w = mju_evalBasis(coord, j-nstart, m->flex_interp[i]);
+              nodeid = j;
+            }
+          }
+          if (nodeid < 0) {
+            mjERROR("flex %d: node closest to vertex %d not found", i, vertid);
+          }
           flexbodyid = m->flex_nodebodyid[m->flex_nodeadr[i] + nodeid];
           if (m->flex_centered[i]) {
             mju_copy3(flexpnt, d->xpos + 3*flexbodyid);