Reduce flex stiffness and bending memory size to the exact amount needed.

PiperOrigin-RevId: 911774147
Change-Id: I3b18eba23e75ec91ece8342945f25e0cd7c63a1a

src/user/user_mesh.cc

@@ -4918,6 +4918,14 @@ void mjCFlex::Compile(const mjVFS* vfs) {
     stiffness_cached = LoadCachedStiffness();
   }
 
+  // check if any strain equality references this flex
+  for (auto* equality : model->Equalities()) {
+    if (equality->spec.type == mjEQ_FLEXSTRAIN && *equality->spec.name1 == name) {
+      has_strain_eq = true;
+      break;
+    }
+  }
+
   if (!stiffness_cached && interpolated && (young > 0 || has_strain_eq)) {
     // use young=1 for strain constraints (eigenvectors are geometry-only)
     double K_young = has_strain_eq ? 1e1 : young;

src/user/user_model.cc

@@ -2106,21 +2106,6 @@ static size_t getpathslength(std::vector<T> list) {
   return result;
 }
 
-// compute extra stiffness/bending array size for an interpolated flex
-static int flexInterpExtraSize(int order, const int cellcount[3], bool shell) {
-  int npe, nelem;
-  int cx = cellcount[0], cy = cellcount[1], cz = cellcount[2];
-  if (shell) {
-    npe = (int)pow(order + 1, 2);
-    nelem = 2*(cy*cz + cx*cz + cx*cy);
-  } else {
-    npe = (int)pow(order + 1, 3);
-    nelem = cx * cy * cz;
-  }
-  int ndof_elem = 3 * npe;
-  return nelem * ndof_elem * ndof_elem;
-}
-
 // set array sizes
 void mjCModel::SetSizes() {
   // set from object list sizes
@@ -2194,8 +2179,6 @@ void mjCModel::SetSizes() {
   }
   nbvh = nbvhstatic + nbvhdynamic;
 
-  int extra_stiffness_size = 0;
-  int extra_bending_size = 0;
   // flex counts
   for (int i=0; i < nflex; i++) {
     nflexnode += flexes_[i]->nnode;
@@ -2207,13 +2190,8 @@ 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]->spec.order != 0) {
-      int extra_size = flexInterpExtraSize(
-          flexes_[i]->spec.order, flexes_[i]->spec.cellcount,
-          flexes_[i]->elastic2d != 0);
-      extra_stiffness_size += extra_size;
-      extra_bending_size += extra_size;
-    }
+    nflexstiffness += flexes_[i]->stiffness.size();
+    nflexbending += flexes_[i]->bending.size();
     if (flexes_[i]->interpolated || flexes_[i]->rigid) {
       continue;
     }
@@ -2263,10 +2241,6 @@ void mjCModel::SetSizes() {
       }
     }
   }
-  // TODO: This can be compacted further when we update mjwarp to not rely on
-  // 21*elem_adr for non-interpolated flexes and 17*edge_adr for bending.
-  nflexstiffness = nflexelem * 21 + extra_stiffness_size;
-  nflexbending = nflexedge * 17 + extra_bending_size;
 
   // mesh counts
   for (int i=0; i < nmesh; i++) {
@@ -3351,6 +3325,7 @@ int mjCModel::CountNJten(const mjModel* m) {
 // copy objects outside kinematic tree
 void mjCModel::CopyObjects(mjModel* m) {
   mjtSize adr, bone_adr, vert_adr, node_adr, normal_adr, face_adr, texcoord_adr, oct_adr;
+  mjtSize stiffness_adr, bending_adr;
   mjtSize edge_adr, elem_adr, elemdata_adr, elemedge_adr, shelldata_adr, evpair_adr;
   mjtSize bonevert_adr, graph_adr, data_adr, bvh_adr;
   mjtSize poly_adr, polymap_adr, polyvert_adr;
@@ -3465,10 +3440,8 @@ void mjCModel::CopyObjects(mjModel* m) {
   shelldata_adr = 0;
   evpair_adr = 0;
   texcoord_adr = 0;
-  int standard_stiffness_size = 21 * m->nflexelem;
-  int current_extra_stiffness_adr = standard_stiffness_size;
-  int standard_bending_size = 17 * m->nflexedge;
-  int current_extra_bending_adr = standard_bending_size;
+  stiffness_adr = 0;
+  bending_adr = 0;
   for (int i=0; i < nflex; i++) {
     // get pointer
     mjCFlex* pfl = flexes_[i];
@@ -3491,46 +3464,19 @@ void mjCModel::CopyObjects(mjModel* m) {
     mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4);
 
     // elasticity
-    if (pfl->spec.order == 0) {
-      m->flex_stiffnessadr[i] = 21 * elem_adr;
+    if (pfl->stiffness.empty()) {
+      m->flex_stiffnessadr[i] = -1;
     } else {
-      m->flex_stiffnessadr[i] = current_extra_stiffness_adr;
-      current_extra_stiffness_adr += flexInterpExtraSize(
-          pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
-    }
-
-    if (!pfl->stiffness.empty()) {
+      m->flex_stiffnessadr[i] = stiffness_adr;
       mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i],
                    pfl->stiffness.data(), pfl->stiffness.size());
-    } else {
-      int stiff_size;
-      if (pfl->spec.order == 0) {
-        stiff_size = 21 * pfl->nelem;
-      } else {
-        stiff_size = flexInterpExtraSize(
-            pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
-      }
-      mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], stiff_size);
     }
-    if (pfl->spec.order == 0) {
-      m->flex_bendingadr[i] = 17 * edge_adr;
+    if (pfl->bending.empty()) {
+      m->flex_bendingadr[i] = -1;
     } else {
-      m->flex_bendingadr[i] = current_extra_bending_adr;
-      current_extra_bending_adr += flexInterpExtraSize(
-          pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
-    }
-
-    if (!pfl->bending.empty()) {
-      mjuu_copyvec(m->flex_bending + m->flex_bendingadr[i], pfl->bending.data(), pfl->bending.size());
-    } else {
-      int bending_size;
-      if (pfl->spec.order == 0) {
-        bending_size = 17 * pfl->nedge;
-      } else {
-        bending_size = flexInterpExtraSize(
-            pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
-      }
-      mjuu_zerovec(m->flex_bending + m->flex_bendingadr[i], bending_size);
+      m->flex_bendingadr[i] = bending_adr;
+      mjuu_copyvec(m->flex_bending + m->flex_bendingadr[i],
+                   pfl->bending.data(), pfl->bending.size());
     }
     m->flex_damping[i] = (mjtNum)pfl->damping;
 
@@ -3704,6 +3650,8 @@ void mjCModel::CopyObjects(mjModel* m) {
     evpair_adr += (int)pfl->evpair.size()/2;
     texcoord_adr += (int)pfl->texcoord_.size()/2;
     bvh_adr += pfl->tree.Nbvh();
+    stiffness_adr += pfl->stiffness.size();
+    bending_adr += pfl->bending.size();
   }
 
   // skins