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

@@ -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;
     }

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

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 ---------------------------------------------------
 

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)

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