Implement multi-cell finite element method for interpolated flexes.

This change introduces a `flex_cellcount` field to `mjModel` to specify the number of cells in each dimension for interpolated flexes. The stiffness computation, passive force calculation, and Jacobian derivatives are updated to operate on a per-cell basis, significantly improving performance by localizing computations to the nodes within each cell.

PiperOrigin-RevId: 901216393
Change-Id: Ic23132e609de11e71bb7fef8d1f139daad2ec264

src/user/user_flexcomp.cc

@@ -77,6 +77,7 @@ bool IsValidElementOrNodeHeader22(const std::string& line) {
 mjCFlexcomp::mjCFlexcomp(void) {
   type = mjFCOMPTYPE_GRID;
   count[0] = count[1] = count[2] = 10;
+  cellcount[0] = cellcount[1] = cellcount[2] = -1;
   mjuu_setvec(spacing, 0.02, 0.02, 0.02);
   mjuu_setvec(scale, 1, 1, 1);
   mass = 1;
@@ -269,10 +270,19 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
 
   // construct pinned array
   int nnode = 0;
-  if (doftype == mjFCOMPDOF_TRILINEAR) {
-    nnode = 8;
-  } else if (doftype == mjFCOMPDOF_QUADRATIC) {
-    nnode = 27;
+  if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) {
+    int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
+    // multi-cell count for mesh/direct/gmsh, else single cell
+    int cx = 1, cy = 1, cz = 1;
+    if (type == mjFCOMPTYPE_MESH || type == mjFCOMPTYPE_DIRECT ||
+        type == mjFCOMPTYPE_GMSH) {
+      if (cellcount[0] >= 0) {
+        cx = cellcount[0];
+        cy = cellcount[1];
+        cz = cellcount[2];
+      }
+    }
+    nnode = (cx*order+1) * (cy*order+1) * (cz*order+1);
   }
   pinned = vector<bool>(std::max(npnt, nnode), rigid);
 
@@ -562,36 +572,81 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
     }
   }
 
-  // create nodal mesh for trilinear interpolation
+  // create nodal mesh for trilinear/quadratic interpolation
   if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) {
-    int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
-    flex->SetOrder(order);
-    std::vector<double> node(3*(order+1)*(order+1)*(order+1), 0);
+    flex->spec.order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
+
+    if (cellcount[0] >= 0) {
+      flex->spec.cellcount[0] = cellcount[0];
+      flex->spec.cellcount[1] = cellcount[1];
+      flex->spec.cellcount[2] = cellcount[2];
+    }
+
+    // total number of nodes with shared boundaries
+    int nx = flex->spec.cellcount[0] * flex->spec.order + 1;
+    int ny = flex->spec.cellcount[1] * flex->spec.order + 1;
+    int nz = flex->spec.cellcount[2] * flex->spec.order + 1;
+    int nnode = nx * ny * nz;
+
+    std::vector<double> node(3 * nnode, 0);
     int idx = 0;
-    double step = 1.0 / (double)order;
+
+    // Simpson's rule weights for quadratic mass distribution
     double massP2[3] = {1. / 6., 2. / 3., 1. / 6.};
-    for (int i=0; i <= order; i++) {
-      for (int j=0; j <= order; j++) {
-        for (int k=0; k <= order; k++) {
+
+    // compute per-node mass for trilinear:
+    //   mass / nnode (uniform), or use Simpson for quadratic
+    double node_mass_uniform = mass / nnode;
+
+    for (int gi = 0; gi < nx; gi++) {
+      for (int gj = 0; gj < ny; gj++) {
+        for (int gk = 0; gk < nz; gk++) {
+          // parametric position in [0, 1]^3
+          double s = (double)gi / (flex->spec.cellcount[0] * flex->spec.order);
+          double t = (double)gj / (flex->spec.cellcount[1] * flex->spec.order);
+          double u = (double)gk / (flex->spec.cellcount[2] * flex->spec.order);
+
+          // physical position
+          double px = minmax[0] + s * (minmax[3] - minmax[0]);
+          double py = minmax[1] + t * (minmax[4] - minmax[1]);
+          double pz = minmax[2] + u * (minmax[5] - minmax[2]);
+
           if (pinned[idx]) {
-            node[3*idx+0] = minmax[0] + i * step * (minmax[3] - minmax[0]);
-            node[3*idx+1] = minmax[1] + j * step * (minmax[4] - minmax[1]);
-            node[3*idx+2] = minmax[2] + k * step * (minmax[5] - minmax[2]);
-            mjs_appendString(pf->nodebody, mjs_getName(body->element)->c_str());
+            node[3*idx+0] = px;
+            node[3*idx+1] = py;
+            node[3*idx+2] = pz;
+            mjs_appendString(pf->nodebody,
+                             mjs_getName(body->element)->c_str());
             idx++;
             continue;
           }
 
           mjsBody* pb = mjs_addBody(body, 0);
-          pb->pos[0] = minmax[0] + i * step * (minmax[3] - minmax[0]);
-          pb->pos[1] = minmax[1] + j * step * (minmax[4] - minmax[1]);
-          pb->pos[2] = minmax[2] + k * step * (minmax[5] - minmax[2]);
+          pb->pos[0] = px;
+          pb->pos[1] = py;
+          pb->pos[2] = pz;
           mjuu_zerovec(pb->ipos, 3);
+
+          // mass distribution
           if (doftype == mjFCOMPDOF_TRILINEAR) {
-            pb->mass = mass / 8;
+            pb->mass = node_mass_uniform;
           } else {
-            pb->mass = mass * massP2[i] * massP2[j] * massP2[k];
+            // local index within the cell for mass computation
+            int li = gi % flex->spec.order;
+            int lj = gj % flex->spec.order;
+            int lk = gk % flex->spec.order;
+            // boundary nodes: average mass contribution
+            int ncells_i = (gi > 0 && gi < nx-1 && li == 0) ? 2 : 1;
+            int ncells_j = (gj > 0 && gj < ny-1 && lj == 0) ? 2 : 1;
+            int ncells_k = (gk > 0 && gk < nz-1 && lk == 0) ? 2 : 1;
+            // use Simpson weights scaled by cell count
+            double wi = massP2[li == 0 ? 0 : li];
+            double wj = massP2[lj == 0 ? 0 : lj];
+            double wk = massP2[lk == 0 ? 0 : lk];
+            pb->mass = mass * wi * wj * wk * ncells_i * ncells_j * ncells_k
+                       / (flex->spec.cellcount[0] * flex->spec.cellcount[1] * flex->spec.cellcount[2]);
           }
+
           pb->inertia[0] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
           pb->inertia[1] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
           pb->inertia[2] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
@@ -607,7 +662,7 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
 
           // construct node name, add to nodebody
           char txt[100];
-          mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), i, j, k);
+          mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), gi, gj, gk);
           mjs_setName(pb->element, txt);
           mjs_appendString(pf->nodebody, mjs_getName(pb->element)->c_str());
 

src/user/user_flexcomp.h

@@ -78,6 +78,7 @@ class mjCFlexcomp {
   std::string name;               // flex name
   mjtFcompType type;              // flexcomp type
   int count[3];                   // grid count in each dimension
+  int cellcount[3];               // number of cells for interpolation
   double spacing[3];              // spacing between grid elements
   double scale[3];                // scaling for mesh and direct
   double origin[3];               // origin for generating a 3D mesh from a convex 2D mesh

src/user/user_init.c

@@ -224,6 +224,9 @@ void mjs_defaultFlex(mjsFlex* flex) {
   // set other defaults
   flex->dim = 2;
   flex->radius = 0.005;
+  flex->cellcount[0] = 1;
+  flex->cellcount[1] = 1;
+  flex->cellcount[2] = 1;
   flex->internal = 0;
   flex->selfcollide = mjFLEXSELF_AUTO;
   flex->activelayers = 1;

src/user/user_mesh.cc

@@ -3961,7 +3961,10 @@ std::string mjCFlex::ComputeStiffnessCacheKey() const {
 
   combine(std::hash<double>{}(young));
   combine(std::hash<double>{}(poisson));
-  combine(std::hash<int>{}(order_));
+  combine(std::hash<int>{}(spec.order));
+  combine(std::hash<int>{}(spec.cellcount[0]));
+  combine(std::hash<int>{}(spec.cellcount[1]));
+  combine(std::hash<int>{}(spec.cellcount[2]));
 
   // compute bounding box from vertex positions
   if (!vert_.empty()) {
@@ -4086,10 +4089,20 @@ void mjCFlex::Compile(const mjVFS* vfs) {
 
   // set nnode
   nnode = static_cast<int>(nodebody_.size());
-  if (nnode && !order_) {
-    order_ = std::pow(nnode, 1.0 / 3) - 1;
-    if (nnode != std::pow(order_ + 1, 3)) {
-      throw mjCError(this, "number of nodes must be %d^3 but it is %d", nullptr, order_, nnode);
+  if (nnode && !spec.order) {
+    throw mjCError(this, "Interpolation order must be explicitly specified (dof is missing)");
+  }
+
+  // check node compatibility with count and dof
+  if (spec.order > 0) {
+    int expected_nodes = (spec.cellcount[0] * spec.order + 1) *
+                         (spec.cellcount[1] * spec.order + 1) *
+                         (spec.cellcount[2] * spec.order + 1);
+    if (nnode != expected_nodes) {
+      std::string msg = "number of nodes (" + std::to_string(nnode) +
+                        ") does not match cellcount and dof expected (" +
+                        std::to_string(expected_nodes) + ")";
+      throw mjCError(this, msg.c_str());
     }
   }
 
@@ -4329,12 +4342,48 @@ void mjCFlex::Compile(const mjVFS* vfs) {
   }
 
   if (!stiffness_cached && young > 0 && interpolated) {
-    int n = pow(order_ + 1, 3);
-    int ndof = 3 * n;
-    if (stiffness.size() < ndof * ndof) {
-      stiffness.resize(ndof * ndof, 0);
+    int npc = pow(spec.order + 1, 3);  // nodes per cell
+    int ndof_cell = 3 * npc;
+    int cx = spec.cellcount[0], cy = spec.cellcount[1], cz = spec.cellcount[2];
+    int ncells = cx * cy * cz;
+    int ny_global = cy * spec.order + 1;
+    int nz_global = cz * spec.order + 1;
+
+    // total stiffness = ncells * ndof_cell^2
+    stiffness.resize(ncells * ndof_cell * ndof_cell, 0);
+
+    // compute stiffness per cell
+    for (int ci = 0; ci < cx; ci++) {
+      for (int cj = 0; cj < cy; cj++) {
+        for (int ck = 0; ck < cz; ck++) {
+          int cell_idx = ci * cy * cz + cj * cz + ck;
+
+          // gather cell's local node positions
+          std::vector<double> cell_pos(3 * npc);
+          int local = 0;
+          for (int li = 0; li <= spec.order; li++) {
+            for (int lj = 0; lj <= spec.order; lj++) {
+              for (int lk = 0; lk <= spec.order; lk++) {
+                int gi = ci * spec.order + li;
+                int gj = cj * spec.order + lj;
+                int gk = ck * spec.order + lk;
+                int global = gi * ny_global * nz_global + gj * nz_global + gk;
+                mjuu_copyvec(cell_pos.data() + 3*local, nodexpos.data() + 3*global, 3);
+                local++;
+              }
+            }
+          }
+
+          // compute per-cell stiffness
+          std::vector<double> K_cell(ndof_cell * ndof_cell, 0);
+          ComputeLinearStiffness(K_cell, cell_pos.data(), young, poisson, spec.order);
+
+          // copy into global stiffness array
+          mjuu_copyvec(stiffness.data() + cell_idx * ndof_cell * ndof_cell,
+                       K_cell.data(), ndof_cell * ndof_cell);
+        }
+      }
     }
-    ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_);
   }
 
   // create bounding volume hierarchy

src/user/user_model.cc

@@ -2190,8 +2190,13 @@ void mjCModel::SetSizes() {
     nflexshelldata += (int)flexes_[i]->shell.size();
     nflexevpair += (int)flexes_[i]->evpair.size()/2;
     nflextexcoord += (flexes_[i]->HasTexcoord() ? flexes_[i]->get_texcoord().size()/2 : 0);
-    if (flexes_[i]->order_ != 0) {
-      extra_stiffness_size += (3 * flexes_[i]->nnode) * (3 * flexes_[i]->nnode);
+    if (flexes_[i]->spec.order != 0) {
+      int npc = (int)pow(flexes_[i]->spec.order + 1, 3);
+      int ndof_cell = 3 * npc;
+      int ncells = flexes_[i]->spec.cellcount[0] *
+                   flexes_[i]->spec.cellcount[1] *
+                   flexes_[i]->spec.cellcount[2];
+      extra_stiffness_size += ncells * ndof_cell * ndof_cell;
     }
     if (flexes_[i]->interpolated || flexes_[i]->rigid) {
       continue;
@@ -3467,18 +3472,30 @@ void mjCModel::CopyObjects(mjModel* m) {
     mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4);
 
     // elasticity
-    if (pfl->order_ == 0) {
+    if (pfl->spec.order == 0) {
       m->flex_stiffnessadr[i] = 21 * elem_adr;
     } else {
       m->flex_stiffnessadr[i] = current_extra_stiffness_adr;
-      current_extra_stiffness_adr += (3 * pfl->nnode) * (3 * pfl->nnode);
+      int npc = (int)pow(pfl->spec.order + 1, 3);
+      int ndof_cell = 3 * npc;
+      int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2];
+      current_extra_stiffness_adr += ncells * ndof_cell * ndof_cell;
     }
 
     if (!pfl->stiffness.empty()) {
-      mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i], pfl->stiffness.data(), pfl->stiffness.size());
+      mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i],
+                   pfl->stiffness.data(), pfl->stiffness.size());
     } else {
-      int size = (pfl->order_ == 0) ? 21 * pfl->nelem : (3 * pfl->nnode) * (3 * pfl->nnode);
-      mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], size);
+      int stiff_size;
+      if (pfl->spec.order == 0) {
+        stiff_size = 21 * pfl->nelem;
+      } else {
+        int npc = (int)pow(pfl->spec.order + 1, 3);
+        int ndof_cell = 3 * npc;
+        int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2];
+        stiff_size = ncells * ndof_cell * ndof_cell;
+      }
+      mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], stiff_size);
     }
     if (!pfl->bending.empty()) {
       mjuu_copyvec(m->flex_bending + 17 * edge_adr, pfl->bending.data(), pfl->bending.size());
@@ -3613,7 +3630,12 @@ void mjCModel::CopyObjects(mjModel* m) {
     }
 
     // set interpolation type, only two types for now
-    m->flex_interp[i] = pfl->order_;
+    m->flex_interp[i] = pfl->spec.order;
+
+    // set cell count for multi-cell finite cell method
+    m->flex_cellnum[3*i+0] = pfl->spec.cellcount[0];
+    m->flex_cellnum[3*i+1] = pfl->spec.cellcount[1];
+    m->flex_cellnum[3*i+2] = pfl->spec.cellcount[2];
 
     // convert edge pairs to int array, set edge rigid
     for (int k=0; k < pfl->nedge; k++) {

src/user/user_objects.h

@@ -1042,7 +1042,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
 
   static constexpr int kNumEdges[3] = {1, 3, 6};  // number of edges per element indexed by dim
 
-  void SetOrder(int order) { order_ = order; }  // set interpolation order
+
 
  private:
   void Compile(const mjVFS* vfs);         // compiler
@@ -1052,7 +1052,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
   std::vector<double> vert0_;             // vertex positions in [0, 1]^d in the bounding box
   std::vector<double> node0_;             // node Cartesian positions
 
-  int order_ = 0;                         // interpolation order
+
 
   // stiffness caching
   std::string ComputeStiffnessCacheKey() const;