Add flex elasticity stiffness matrix to mjModel.

PiperOrigin-RevId: 675101645
Change-Id: Ic169c108b9eece3657ea4ec211df357abe615b7e

doc/includes/references.h

@@ -1168,6 +1168,7 @@ struct mjModel_ {
   mjtNum*   flexedge_length0;     // edge lengths in qpos0                    (nflexedge x 1)
   mjtNum*   flexedge_invweight0;  // edge inv. weight in qpos0                (nflexedge x 1)
   mjtNum*   flex_radius;          // radius around primitive element          (nflex x 1)
+  mjtNum*   flex_stiffness;       // finite element stiffness matrix          (nflexelem x 21)
   mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
   mjtByte*  flex_edgeequality;    // is edge equality constraint defined      (nflex x 1)

include/mujoco/mjmodel.h

@@ -878,6 +878,7 @@ struct mjModel_ {
   mjtNum*   flexedge_length0;     // edge lengths in qpos0                    (nflexedge x 1)
   mjtNum*   flexedge_invweight0;  // edge inv. weight in qpos0                (nflexedge x 1)
   mjtNum*   flex_radius;          // radius around primitive element          (nflex x 1)
+  mjtNum*   flex_stiffness;       // finite element stiffness matrix          (nflexelem x 21)
   mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
   mjtByte*  flex_edgeequality;    // is edge equality constraint defined      (nflex x 1)

include/mujoco/mjxmacro.h

@@ -347,6 +347,7 @@
     X   ( mjtNum,  flexedge_length0,      nflexedge,     1                    ) \
     X   ( mjtNum,  flexedge_invweight0,   nflexedge,     1                    ) \
     XMJV( mjtNum,  flex_radius,           nflex,         1                    ) \
+    X   ( mjtNum,  flex_stiffness,        nflexelem,     21                   ) \
     X   ( mjtNum,  flex_edgestiffness,    nflex,         1                    ) \
     X   ( mjtNum,  flex_edgedamping,      nflex,         1                    ) \
     X   ( mjtByte, flex_edgeequality,     nflex,         1                    ) \

introspect/structs.py

@@ -2449,6 +2449,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='radius around primitive element          (nflex x 1)',
              ),
+             StructFieldDecl(
+                 name='flex_stiffness',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum'),
+                 ),
+                 doc='finite element stiffness matrix          (nflexelem x 21)',  # pylint: disable=line-too-long
+             ),
              StructFieldDecl(
                  name='flex_edgestiffness',
                  type=PointerType(

plugin/elasticity/elasticity.h

@@ -66,13 +66,18 @@ void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3],
 template <typename T>
 inline void ComputeForce(std::vector<mjtNum>& qfrc_passive,
                          const std::vector<T>& elements,
-                         const std::vector<mjtNum>& metric,
                          const std::vector<mjtNum>& elongationglob,
-                         const mjModel* m,
+                         const mjModel* m, int flex,
                          const mjtNum* xpos) {
   mju_zero(qfrc_passive.data(), qfrc_passive.size());
+  mjtNum* k = m->flex_stiffness + 21 * m->flex_elemadr[flex];
 
-  for (int t = 0; t < elements.size(); t++)  {
+  if (elements.size() != m->flex_elemnum[flex]) {
+    mju_error("plugin stencil does not match flex stencil");
+  }
+
+  // compute force element-by-element
+  for (int t = 0; t < m->flex_elemnum[flex]; t++)  {
     const int* v = elements[t].vertices;
 
     // compute length gradient with respect to dofs
@@ -86,20 +91,30 @@ inline void ComputeForce(std::vector<mjtNum>& qfrc_passive,
       elongation[e] = elongationglob[idx];
     }
 
+    // unpack triangular representation
+    mjtNum metric[T::kNumEdges*T::kNumEdges];
+
+    int id = 0;
+    for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
+      for (int ed2 = ed1; ed2 < T::kNumEdges; ed2++) {
+        metric[T::kNumEdges*ed1 + ed2] = k[21*t + id];
+        metric[T::kNumEdges*ed2 + ed1] = k[21*t + id++];
+      }
+    }
+
     // we now multiply the elongations by the precomputed metric tensor,
     // notice that if metric=diag(1/reference) then this would yield a
     // mass-spring model
 
     // compute local force
     mjtNum force[T::kNumVerts*3] = {0};
-    int offset = T::kNumEdges*T::kNumEdges;
     for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
       for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
         for (int i = 0; i < 2; i++) {
           for (int x = 0; x < 3; x++) {
             force[3 * T::edge[ed2][i] + x] -=
                 elongation[ed1] * gradient[ed2][i][x] *
-                metric[offset * t + T::kNumEdges * ed1 + ed2];
+                metric[T::kNumEdges * ed1 + ed2];
           }
         }
       }
@@ -139,10 +154,11 @@ inline void AddFlexForce(mjtNum* qfrc,
 
 // compute metric tensor of edge lengths inner product
 template <typename T>
-void inline MetricTensor(std::vector<mjtNum>& metric, int idx, mjtNum mu,
+void inline MetricTensor(mjtNum* metric, int idx, mjtNum mu,
                          mjtNum la, const mjtNum basis[T::kNumEdges][9]) {
   mjtNum trE[T::kNumEdges] = {0};
   mjtNum trEE[T::kNumEdges*T::kNumEdges] = {0};
+  mjtNum k[T::kNumEdges*T::kNumEdges];
 
   // compute first invariant i.e. trace(strain)
   for (int e = 0; e < T::kNumEdges; e++) {
@@ -165,11 +181,22 @@ void inline MetricTensor(std::vector<mjtNum>& metric, int idx, mjtNum mu,
   // assembly of strain metric tensor
   for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
     for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
-      int index = T::kNumEdges*T::kNumEdges*idx + T::kNumEdges*ed1 + ed2;
+      k[T::kNumEdges*ed1 + ed2] = mu * trEE[T::kNumEdges * ed1 + ed2] +
-      metric[index] = mu * trEE[T::kNumEdges * ed1 + ed2] +
+                                  la * trE[ed2] * trE[ed1];
-                      la * trE[ed2] * trE[ed1];
     }
   }
+
+  // copy to triangular representation
+  int id = 0;
+  for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
+    for (int ed2 = ed1; ed2 < T::kNumEdges; ed2++) {
+      metric[21*idx + id++] = k[T::kNumEdges*ed1 + ed2];
+    }
+  }
+
+  if (id != T::kNumEdges*(T::kNumEdges+1)/2) {
+    mju_error("incorrect stiffness matrix size");
+  }
 }
 
 // convert from Flex connectivity to stencils

plugin/elasticity/membrane.cc

@@ -129,9 +129,6 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   // generate triangles from the vertices
   nt = CreateStencils<Stencil2D>(elements, edges, simplex, edgeidx);
 
-  // allocate metric induced by geometry
-  metric.assign(kNumEdges*kNumEdges*nt, 0);
-
   // loop over all triangles
   for (int t = 0; t < nt; t++) {
     int* v = elements[t].vertices;
@@ -160,7 +157,9 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
     }
 
     // compute metric tensor
-    MetricTensor<Stencil2D>(metric, t, mu, la, basis);
+    // TODO: do not write in a const mjModel
+    MetricTensor<Stencil2D>(m->flex_stiffness + 21 * m->flex_elemadr[f0], t, mu,
+                            la, basis);
   }
 
   // allocate array
@@ -196,7 +195,7 @@ void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
   mjtNum* xpos = d->flexvert_xpos + 3*flex_vertadr;
   mjtNum* qfrc = d->qfrc_passive;
 
-  ComputeForce<Stencil2D>(force, elements, metric, elongation, m, xpos);
+  ComputeForce<Stencil2D>(force, elements, elongation, m, f0, xpos);
 
   // insert into passive force
   AddFlexForce(qfrc, force, m, d, xpos, f0);

plugin/elasticity/membrane.h

@@ -52,7 +52,6 @@ class Membrane {
   std::vector<std::pair<int, int> > edges;  // edge to vertex map     (ne x 2)
 
   // precomputed quantities
-  std::vector<mjtNum> metric;               // geom-induced metric    (nt x 9)
   std::vector<mjtNum> prev;                 // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
   std::vector<mjtNum> force;                // force at all vertices  (nv x 3)

plugin/elasticity/solid.cc

@@ -137,9 +137,6 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   // generate tetrahedra from the vertices
   nt = CreateStencils<Stencil3D>(elements, edges, simplex, edgeidx);
 
-  // allocate arrays
-  metric.assign(kNumEdges*kNumEdges*nt, 0);
-
   // loop over all tetrahedra
   for (int t = 0; t < nt; t++) {
     int* v = elements[t].vertices;
@@ -167,7 +164,9 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
     mjtNum la = E*nu / ((1+nu)*(1-2*nu)) * volume;
 
     // compute metric tensor
-    MetricTensor<Stencil3D>(metric, t, mu, la, basis);
+    // TODO: do not write in a const mjModel
+    MetricTensor<Stencil3D>(m->flex_stiffness + 21 * m->flex_elemadr[f0], t, mu,
+                            la, basis);
   }
 
   // allocate array
@@ -203,7 +202,7 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) {
   mjtNum* xpos = d->flexvert_xpos + 3*flex_vertadr;
   mjtNum* qfrc = d->qfrc_passive;
 
-  ComputeForce<Stencil3D>(force, elements, metric, elongation, m, xpos);
+  ComputeForce<Stencil3D>(force, elements, elongation, m, f0, xpos);
 
   // insert into passive force
   AddFlexForce(qfrc, force, m, d, xpos, f0);

plugin/elasticity/solid.h

@@ -50,7 +50,6 @@ class Solid {
   std::vector<std::pair<int, int> > edges;  // edge to vertex map     (ne x 2)
 
   // precomputed quantities
-  std::vector<mjtNum> metric;               // geom-induced metric    (nt x 36)
   std::vector<mjtNum> prev;                 // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
   std::vector<mjtNum> force;                // force at all vertices  (nv x 3)

src/user/user_model.cc

@@ -2525,6 +2525,10 @@ void mjCModel::CopyObjects(mjModel* m) {
     m->flex_gap[i] = (mjtNum)pfl->gap;
     mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4);
 
+    // elasticity
+    // TODO: these are now written by plugins, they will be moved to mjCFlex
+    mjuu_zerovec(m->flex_stiffness + 21 * elem_adr, 21 * pfl->nelem);
+
     // set fields: mesh-like
     m->flex_dim[i] = pfl->dim;
     m->flex_vertadr[i] = vert_adr;

test/plugin/elasticity/elasticity_test.cc

@@ -148,6 +148,7 @@ TEST_F(PluginTest, ElasticEnergyMembrane) {
 
   mj_kinematics(m, d);
   mj_flex(m, d);
+  mjtNum* metric = m->flex_stiffness + 21 * m->flex_elemadr[0];
 
   // check that if the entire geometry is rescaled by a factor "scale", then
   // trace(strain^2) = 2*scale^2
@@ -156,15 +157,16 @@ TEST_F(PluginTest, ElasticEnergyMembrane) {
     for (int t = 0; t < membrane->nt; t++) {
       mjtNum energy = 0;
       mjtNum volume = 1./2.;
+      int idx = 0;
       for (int e1 = 0; e1 < 3; e1++) {
-        for (int e2 = 0; e2 < 3; e2++) {
+        for (int e2 = e1; e2 < 3; e2++) {
           int idx1 = membrane->elements[t].edges[e1] + m->flex_edgeadr[0];
           int idx2 = membrane->elements[t].edges[e2] + m->flex_edgeadr[0];
-          mjtNum elongation1 =
+          mjtNum elong1 =
               scale * m->flexedge_length0[idx1] * m->flexedge_length0[idx1];
-          mjtNum elongation2 =
+          mjtNum elong2 =
               scale * m->flexedge_length0[idx2] * m->flexedge_length0[idx2];
-          energy += membrane->metric[9*t+3*e2+e1] * elongation1 * elongation2;
+          energy += metric[21*t+idx++] * elong1 * elong2 * (e1 == e2 ? 1. : 2.);
         }
       }
       EXPECT_NEAR(
@@ -230,6 +232,7 @@ TEST_F(ElasticityTest, ElasticEnergySolid) {
 
   mj_kinematics(m, d);
   mj_flex(m, d);
+  mjtNum* metric = m->flex_stiffness + 21 * m->flex_elemadr[0];
 
   // check that if the entire geometry is rescaled by a factor "scale", then
   // trace(strain^2) = 3*scale^2
@@ -238,15 +241,16 @@ TEST_F(ElasticityTest, ElasticEnergySolid) {
     for (int t = 0; t < solid->nt; t++) {
       mjtNum energy = 0;
       mjtNum volume = 1./6.;
+      int idx = 0;
       for (int e1 = 0; e1 < 6; e1++) {
-        for (int e2 = 0; e2 < 6; e2++) {
+        for (int e2 = e1; e2 < 6; e2++) {
           int idx1 = solid->elements[t].edges[e1] + m->flex_edgeadr[0];
           int idx2 = solid->elements[t].edges[e2] + m->flex_edgeadr[0];
-          mjtNum elongation1 =
+          mjtNum elong1 =
               scale * m->flexedge_length0[idx1] * m->flexedge_length0[idx1];
-          mjtNum elongation2 =
+          mjtNum elong2 =
               scale * m->flexedge_length0[idx2] * m->flexedge_length0[idx2];
-          energy += solid->metric[36*t+6*e2+e1] * elongation1 * elongation2;
+          energy += metric[21*t+idx++] * elong1 * elong2 * (e1 == e2 ? 1. : 2.);
         }
       }
       EXPECT_NEAR(

unity/Runtime/Bindings/MjBindings.cs

@@ -5421,6 +5421,7 @@ public unsafe struct mjModel_ {
   public double* flexedge_length0;
   public double* flexedge_invweight0;
   public double* flex_radius;
+  public double* flex_stiffness;
   public double* flex_edgestiffness;
   public double* flex_edgedamping;
   public byte* flex_edgeequality;