Add membrane plugin (2D Flex elastic stiffness).

PiperOrigin-RevId: 574109411
Change-Id: I4a701d8189cecf540cd200bd30cb4582f3c7dd43

model/plugin/elasticity/trampoline_flex.xml

@@ -0,0 +1,62 @@
+<!-- Copyright 2021 DeepMind Technologies Limited
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<mujoco model="Trampoline">
+  <include file="scene.xml"/>
+
+  <extension>
+    <plugin plugin="mujoco.elasticity.shell"/>
+    <plugin plugin="mujoco.elasticity.membrane"/>
+  </extension>
+
+  <option timestep="0.001" solver="CG" tolerance="1e-6" integrator="implicitfast"/>
+
+  <size memory="10M"/>
+
+  <visual>
+    <map stiffness="100"/>
+  </visual>
+
+  <default>
+    <default class="wall">
+      <geom type="plane" size=".5 .5 .05"/>
+    </default>
+  </default>
+
+  <worldbody>
+    <body pos=".2 0 1.5">
+      <freejoint/>
+      <geom type="sphere" size=".1"/>
+    </body>
+    <flexcomp type="grid" count="16 16 1" spacing=".1 .1 .1" pos="0 0 1"
+              radius=".001" mass="10" name="plate" dim="2">
+      <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
+      <edge equality="false" damping="10"/>
+      <plugin plugin="mujoco.elasticity.membrane">
+        <config key="poisson" value="0"/>
+        <config key="thickness" value="1e-2"/>
+        <!--Units are in Pa (SI)-->
+        <config key="young" value="3e5"/>
+      </plugin>
+    </flexcomp>
+  </worldbody>
+
+  <equality>
+    <connect body1="plate_0" anchor="0 0 0"/>
+    <connect body1="plate_15" anchor="0 0 0"/>
+    <connect body1="plate_240" anchor="0 0 0"/>
+    <connect body1="plate_255" anchor="0 0 0"/>
+  </equality>
+</mujoco>

plugin/elasticity/CMakeLists.txt

@@ -21,6 +21,8 @@ set(MUJOCO_ELASTICITY_SRCS
     cable.h
     elasticity.cc
     elasticity.h
+    membrane.cc
+    membrane.h
     register.cc
     shell.cc
     shell.h

plugin/elasticity/elasticity.cc

@@ -14,15 +14,75 @@
 
 #include "elasticity.h"
 #include <algorithm>
+#include <cassert>
 #include <cctype>
 #include <cstdlib>
 #include <sstream>
 #include <string>
+#include <utility>
 #include <vector>
+#include <unordered_map>
 #include <mujoco/mujoco.h>
 
 namespace mujoco::plugin::elasticity {
 
+template <typename T>
+int CreateStencils(std::vector<T>& elements,
+                   std::vector<std::pair<int, int>>& edges,
+                   const std::vector<int>& simplex,
+                   const std::vector<int>& edgeidx) {
+  int ne = 0;
+  int nt = simplex.size() / T::kNumVerts;
+  elements.resize(nt);
+  for (int t = 0; t < nt; t++) {
+    for (int v = 0; v < T::kNumVerts; v++) {
+      elements[t].vertices[v] = simplex[T::kNumVerts*t+v];
+    }
+  }
+
+  // map from edge vertices to their index in `edges` vector
+  std::unordered_map<std::pair<int, int>, int, PairHash> edge_indices;
+
+  // loop over all tetrahedra
+  for (int t = 0; t < nt; t++) {
+    int* v = elements[t].vertices;
+
+    // compute edges to vertices map for fast computations
+    for (int e = 0; e < T::kNumEdges; e++) {
+      auto pair = std::pair(
+        std::min(v[T::edge[e][0]], v[T::edge[e][1]]),
+        std::max(v[T::edge[e][0]], v[T::edge[e][1]])
+      );
+
+      // if edge is already present in the vector only store its index
+      auto [it, inserted] = edge_indices.insert({pair, ne});
+
+      if (inserted) {
+        edges.push_back(pair);
+        elements[t].edges[e] = ne++;
+      } else {
+        elements[t].edges[e] = it->second;
+      }
+
+      if (!edgeidx.empty()) {
+        assert(elements[t].edges[e] == edgeidx[T::kNumEdges*t+e]);
+      }
+    }
+  }
+
+  return nt;
+}
+
+template int CreateStencils<Stencil2D>(std::vector<Stencil2D>& elements,
+                                       std::vector<std::pair<int, int>>& edges,
+                                       const std::vector<int>& simplex,
+                                       const std::vector<int>& edgeidx);
+
+template int CreateStencils<Stencil3D>(std::vector<Stencil3D>& elements,
+                                       std::vector<std::pair<int, int>>& edges,
+                                       const std::vector<int>& simplex,
+                                       const std::vector<int>& edgeidx);
+
 void String2Vector(const std::string& txt, std::vector<int>& vec) {
   std::stringstream strm(txt);
   vec.clear();

plugin/elasticity/elasticity.h

@@ -15,8 +15,10 @@
 #ifndef MUJOCO_PLUGIN_ELASTICITY_ELASTICITY_H_
 #define MUJOCO_PLUGIN_ELASTICITY_ELASTICITY_H_
 
-#include <sstream>
+#include <cstddef>
+#include <functional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <mujoco/mujoco.h>
@@ -46,6 +48,78 @@ inline void UpdateSquaredLengths(std::vector<mjtNum>& len,
   }
 }
 
+struct Stencil2D {
+  static constexpr int kNumEdges = 3;
+  static constexpr int kNumVerts = 3;
+  static constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}};
+  int vertices[kNumVerts];
+  int edges[kNumEdges];
+};
+
+struct Stencil3D {
+  static constexpr int kNumEdges = 6;
+  static constexpr int kNumVerts = 4;
+  static constexpr int edge[kNumEdges][2] = {{0, 1}, {1, 2}, {2, 0},
+                                             {2, 3}, {0, 3}, {1, 3}};
+  int vertices[kNumVerts];
+  int edges[kNumEdges];
+};
+
+// gradients of edge lengths with respect to vertex positions
+template <typename T>
+void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3],
+                               const mjtNum* x,
+                               const int v[T::kNumVerts]) {
+  for (int e = 0; e < T::kNumEdges; e++) {
+    for (int d = 0; d < 3; d++) {
+      gradient[e][0][d] = x[3*v[T::edge[e][0]]+d] - x[3*v[T::edge[e][1]]+d];
+      gradient[e][1][d] = x[3*v[T::edge[e][1]]+d] - x[3*v[T::edge[e][0]]+d];
+    }
+  }
+}
+
+// compute metric tensor of edge lengths inner product
+template <typename T>
+void inline MetricTensor(std::vector<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};
+
+  // compute first invariant i.e. trace(strain)
+  for (int e = 0; e < T::kNumEdges; e++) {
+    for (int i = 0; i < 3; i++) {
+      trE[e] += basis[e][4*i];
+    }
+  }
+
+  // compute second invariant i.e. trace(strain^2)
+  for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
+    for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
+      for (int i = 0; i < 3; i++) {
+        for (int j = 0; j < 3; j++) {
+          trEE[T::kNumEdges*ed1+ed2] += basis[ed1][3*i+j] * basis[ed2][3*j+i];
+        }
+      }
+    }
+  }
+
+  // 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;
+      metric[index] = mu * trEE[T::kNumEdges * ed1 + ed2] +
+                      la * trE[ed2] * trE[ed1];
+    }
+  }
+}
+
+// convert from Flex connectivity to stencils
+template <typename T>
+int CreateStencils(std::vector<T>& elements,
+                   std::vector<std::pair<int, int>>& edges,
+                   const std::vector<int>& simplex,
+                   const std::vector<int>& edgeidx);
+
 // copied from mjXUtil
 void String2Vector(const std::string& txt, std::vector<int>& vec);
 

plugin/elasticity/membrane.cc

@@ -0,0 +1,237 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cstdint>
+#include <cstdlib>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include <mujoco/mjplugin.h>
+#include <mujoco/mjtnum.h>
+#include <mujoco/mujoco.h>
+#include "elasticity.h"
+#include "membrane.h"
+
+
+namespace mujoco::plugin::elasticity {
+namespace {
+
+// local tetrahedron numbering
+constexpr int kNumEdges = Stencil2D::kNumEdges;
+constexpr int kNumVerts = Stencil2D::kNumVerts;
+
+// area of a triangle
+mjtNum ComputeVolume(const mjtNum* x, const int v[kNumVerts]) {
+  mjtNum normal[3];
+  mjtNum edge1[3];
+  mjtNum edge2[3];
+
+  mju_sub3(edge1, x+3*v[1], x+3*v[0]);
+  mju_sub3(edge2, x+3*v[2], x+3*v[0]);
+  mju_cross(normal, edge1, edge2);
+
+  return mju_norm3(normal) / 2;
+}
+
+// compute local basis
+void ComputeBasis(mjtNum basis[9], const mjtNum* x, const int v[kNumVerts],
+                  const int faceL[2], const int faceR[2], mjtNum area) {
+  mjtNum basisL[3], basisR[3];
+  mjtNum edgesL[3], edgesR[3];
+  mjtNum normal[3];
+
+  mju_sub3(edgesL, x+3*v[faceL[0]], x+3*v[faceL[1]]);
+  mju_sub3(edgesR, x+3*v[faceR[1]], x+3*v[faceR[0]]);
+
+  mju_cross(normal, edgesR, edgesL);
+  mju_normalize3(normal);
+  mju_cross(basisL, normal, edgesL);
+  mju_cross(basisR, edgesR, normal);
+
+  // we use as basis the symmetrized tensor products of the edge normals of the
+  // other two edges; this is shown in Weischedel "A discrete geometric view on
+  // shear-deformable shell models" in the remark at the end of section 4.1;
+  // equivalent to linear finite elements but in a coordinate-free formulation.
+
+  for (int i = 0; i < 3; i++) {
+    for (int j = 0; j < 3; j++) {
+      basis[3*i+j] = ( basisL[i]*basisR[j] +
+                       basisR[i]*basisL[j] ) / (8*area*area);
+    }
+  }
+}
+
+}  // namespace
+
+// factory function
+std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
+                                         int instance) {
+  if (CheckAttr("face", m, instance) && CheckAttr("poisson", m, instance) &&
+      CheckAttr("young", m, instance) && CheckAttr("thickness", m, instance)) {
+    mjtNum nu = strtod(mj_getPluginConfig(m, instance, "poisson"), nullptr);
+    mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
+    mjtNum thick =
+        strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
+    std::vector<int> face, edge;
+    String2Vector(mj_getPluginConfig(m, instance, "face"), face);
+    String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
+    return Membrane(m, d, instance, nu, E, thick, face, edge);
+  } else {
+    mju_warning("Invalid parameter specification in shell plugin");
+    return std::nullopt;
+  }
+}
+
+// plugin constructor
+Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
+                   mjtNum E, mjtNum thick, const std::vector<int>& simplex,
+                   const std::vector<int>& edgeidx)
+    : thickness(thick) {
+  // count plugin bodies
+  nv = ne = 0;
+  for (int i = 1; i < m->nbody; i++) {
+    if (m->body_plugin[i] == instance) {
+      if (!nv++) {
+        i0 = i;
+      }
+    }
+  }
+
+  // count flexes
+  for (int i = 0; i < m->nflex; i++) {
+    if (m->flex_vertbodyid[m->flex_vertadr[i]] == i0) {
+      f0 = i;
+      break;
+    }
+  }
+
+  // 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;
+    for (int i = 0; i < kNumVerts; i++) {
+      if (m->body_plugin[i0+v[i]] != instance) {
+        mju_error("This body does not have the requested plugin instance");
+      }
+    }
+
+    // triangles area
+    mjtNum volume = ComputeVolume(m->body_pos+3*i0, v);
+
+    // material parameters
+    mjtNum mu = E / (2*(1+nu)) * mju_abs(volume) / 4 * thickness;
+    mjtNum la = E*nu / ((1+nu)*(1-2*nu)) * mju_abs(volume) / 4 * thickness;
+
+    // local geometric quantities
+    mjtNum basis[kNumEdges][9] = {{0}, {0}, {0}};
+
+    // compute edge basis
+    for (int e = 0; e < kNumEdges; e++) {
+      ComputeBasis(basis[e], m->body_pos+3*i0, v,
+                   Stencil2D::edge[Stencil2D::edge[e][0]],
+                   Stencil2D::edge[Stencil2D::edge[e][1]], volume);
+    }
+
+    // compute metric tensor
+    MetricTensor<Stencil2D>(metric, t, mu, la, basis);
+  }
+}
+
+void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
+  for (int t = 0; t < nt; t++)  {
+    int* v = elements[t].vertices;
+
+    // compute length gradient with respect to dofs
+    mjtNum gradient[kNumEdges][2][3];
+    GradSquaredLengths<Stencil2D>(gradient, d->xpos+3*i0, v);
+
+    // compute elongation
+    mjtNum elongation[kNumEdges];
+    for (int e = 0; e < kNumEdges; e++) {
+      int idx = elements[t].edges[e] + m->flex_edgeadr[f0];
+      mjtNum deformed = d->flexedge_length[idx]*d->flexedge_length[idx];
+      mjtNum reference = m->flexedge_length0[idx]*m->flexedge_length0[idx];
+      elongation[e] = deformed - reference;
+    }
+
+    // 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[kNumVerts*3] = {0};
+    int offset = kNumEdges*kNumEdges;
+    for (int ed1 = 0; ed1 < kNumEdges; ed1++) {
+      for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
+        for (int i = 0; i < 2; i++) {
+          for (int x = 0; x < 3; x++) {
+            force[3 * Stencil2D::edge[ed2][i] + x] +=
+                elongation[ed1] * gradient[ed2][i][x] *
+                metric[offset * t + kNumEdges * ed1 + ed2];
+          }
+        }
+      }
+    }
+
+    // insert into global force
+    for (int i = 0; i < kNumVerts; i++) {
+      for (int x = 0; x < 3; x++) {
+        d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x];
+      }
+    }
+  }
+}
+
+
+
+void Membrane::RegisterPlugin() {
+  mjpPlugin plugin;
+  mjp_defaultPlugin(&plugin);
+
+  plugin.name = "mujoco.elasticity.membrane";
+  plugin.capabilityflags |= mjPLUGIN_PASSIVE;
+
+  const char* attributes[] = {"face", "edge", "young", "poisson", "thickness"};
+  plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
+  plugin.attributes = attributes;
+  plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
+
+  plugin.init = +[](const mjModel* m, mjData* d, int instance) {
+    auto elasticity_or_null = Membrane::Create(m, d, instance);
+    if (!elasticity_or_null.has_value()) {
+      return -1;
+    }
+    d->plugin_data[instance] = reinterpret_cast<uintptr_t>(
+        new Membrane(std::move(*elasticity_or_null)));
+    return 0;
+  };
+  plugin.destroy = +[](mjData* d, int instance) {
+    delete reinterpret_cast<Membrane*>(d->plugin_data[instance]);
+    d->plugin_data[instance] = 0;
+  };
+  plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) {
+    auto* elasticity = reinterpret_cast<Membrane*>(d->plugin_data[instance]);
+    elasticity->Compute(m, d, instance);
+  };
+
+  mjp_registerPlugin(&plugin);
+}
+
+}  // namespace mujoco::plugin::elasticity

plugin/elasticity/membrane.h

@@ -0,0 +1,67 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_PLUGIN_ELASTICITY_MEMBRANE_H_
+#define MUJOCO_PLUGIN_ELASTICITY_MEMBRANE_H_
+
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
+#include "elasticity.h"
+
+
+namespace mujoco::plugin::elasticity {
+
+class Membrane {
+ public:
+  // Returns a new Membrane instance or nullopt on failure.
+  static std::optional<Membrane> Create(const mjModel* m, mjData* d,
+                                     int instance);
+  Membrane(Membrane&&) = default;
+
+  Membrane& operator=(Membrane&& other) = default;
+
+  void Compute(const mjModel* m, mjData* d, int instance);
+
+  static void RegisterPlugin();
+
+  int f0;  // index of corresponding flex
+  int i0;  // index of first body
+  int nc;  // number of quads in the grid
+  int nv;  // number of vertices (bodies) in the Membrane
+  int nt;  // number of area elements (triangles)
+  int ne;  // number of edges in the Membrane
+
+  // connectivity info for mapping tetrahedra to edges and vertices
+  std::vector<Stencil2D> elements;          // triangles              (nt x 6)
+  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)
+
+  mjtNum thickness;
+
+ private:
+  Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
+           mjtNum thick, const std::vector<int>& simplex,
+           const std::vector<int>& edgeidx);
+};
+
+}  // namespace mujoco::plugin::elasticity
+
+#endif  // MUJOCO_PLUGIN_ELASTICITY_MEMBRANE_H_

plugin/elasticity/register.cc

@@ -15,12 +15,14 @@
 #include <mujoco/mjplugin.h>
 #include "cable.h"
 #include "shell.h"
+#include "membrane.h"
 #include "solid.h"
 
 namespace mujoco::plugin::elasticity {
 
 mjPLUGIN_LIB_INIT {
   Cable::RegisterPlugin();
+  Membrane::RegisterPlugin();
   Shell::RegisterPlugin();
   Solid::RegisterPlugin();
 }

plugin/elasticity/shell.h

@@ -21,17 +21,11 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
+#include "elasticity.h"
 
 
 namespace mujoco::plugin::elasticity {
 
-struct Stencil2D {
-  static constexpr int kNumEdges = 3;
-  static constexpr int kNumVerts = 3;
-  int vertices[kNumVerts];
-  int edges[kNumEdges];
-};
-
 struct StencilFlap {
   static constexpr int kNumVerts = 4;
   int vertices[kNumVerts];

plugin/elasticity/solid.cc

@@ -17,7 +17,6 @@
 #include <cstdint>
 #include <cstdlib>
 #include <optional>
-#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -34,8 +33,6 @@ namespace {
 // local tetrahedron numbering
 constexpr int kNumEdges = Stencil3D::kNumEdges;
 constexpr int kNumVerts = Stencil3D::kNumVerts;
-constexpr int edge[kNumEdges][2] = {{0, 1}, {1, 2}, {2, 0},
-                                    {2, 3}, {0, 3}, {1, 3}};
 constexpr int face[kNumVerts][3] = {{2, 1, 0}, {0, 1, 3}, {1, 2, 3}, {2, 0, 3}};
 constexpr int e2f[kNumEdges][2] = {{2, 3}, {1, 3}, {2, 1},
                                    {1, 0}, {0, 2}, {0, 3}};
@@ -83,19 +80,6 @@ void ComputeBasis(mjtNum basis[9], const mjtNum* x, const int v[kNumVerts],
   }
 }
 
-// gradients of edge lengths with respect to vertex positions
-void GradSquaredLengths(mjtNum gradient[kNumEdges][2][3],
-                        const mjtNum* x,
-                        const int v[kNumVerts],
-                        const int edge[kNumEdges][2]) {
-  for (int e = 0; e < kNumEdges; e++) {
-    for (int d = 0; d < 3; d++) {
-      gradient[e][0][d] = x[3*v[edge[e][0]]+d] - x[3*v[edge[e][1]]+d];
-      gradient[e][1][d] = x[3*v[edge[e][1]]+d] - x[3*v[edge[e][0]]+d];
-    }
-  }
-}
-
 }  // namespace
 
 // factory function
@@ -118,49 +102,6 @@ std::optional<Solid> Solid::Create(const mjModel* m, mjData* d, int instance) {
     }
 }
 
-// create map from tetrahedra to vertices and edges and from edges to vertices
-void Solid::CreateStencils(const std::vector<int>& simplex,
-                           const std::vector<int>& edgeidx) {
-  // populate stencil
-  nt = simplex.size() / kNumVerts;
-  elements.resize(nt);
-  for (int t = 0; t < nt; t++) {
-    for (int v = 0; v < kNumVerts; v++) {
-      elements[t].vertices[v] = simplex[kNumVerts*t+v];
-    }
-  }
-
-  // map from edge vertices to their index in `edges` vector
-  std::unordered_map<std::pair<int, int>, int, PairHash> edge_indices;
-
-  // loop over all tetrahedra
-  for (int t = 0; t < nt; t++) {
-    int* v = elements[t].vertices;
-
-    // compute edges to vertices map for fast computations
-    for (int e = 0; e < kNumEdges; e++) {
-      auto pair = std::pair(
-        std::min(v[edge[e][0]], v[edge[e][1]]),
-        std::max(v[edge[e][0]], v[edge[e][1]])
-      );
-
-      // if edge is already present in the vector only store its index
-      auto [it, inserted] = edge_indices.insert({pair, ne});
-
-      if (inserted) {
-        edges.push_back(pair);
-        elements[t].edges[e] = ne++;
-      } else {
-        elements[t].edges[e] = it->second;
-      }
-
-      if (!edgeidx.empty()) {
-        assert(elements[t].edges[e] == edgeidx[kNumEdges*t+e]);
-      }
-    }
-  }
-}
-
 // plugin constructor
 Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
              mjtNum damp, const std::vector<int>& simplex,
@@ -185,7 +126,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   }
 
   // generate tetrahedra from the vertices
-  CreateStencils(simplex, edgeidx);
+  nt = CreateStencils<Stencil3D>(elements, edges, simplex, edgeidx);
 
   // allocate arrays
   metric.assign(kNumEdges*kNumEdges*nt, 0);
@@ -204,8 +145,6 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
 
     // local geometric quantities
     mjtNum basis[kNumEdges][9] = {{0}, {0}, {0}, {0}, {0}, {0}};
-    mjtNum trT[kNumEdges] = {0};
-    mjtNum trTT[kNumEdges*kNumEdges] = {0};
 
     // compute edge basis
     for (int e = 0; e < kNumEdges; e++) {
@@ -213,38 +152,16 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
                    face[e2f[e][0]], face[e2f[e][1]], volume);
     }
 
-    // compute first invariant i.e. trace(strain)
-    for (int e = 0; e < kNumEdges; e++) {
-      for (int i = 0; i < 3; i++) {
-        trT[e] += basis[e][4*i];
-      }
-    }
-
-    // compute second invariant i.e. trace(strain^2)
-    for (int ed1 = 0; ed1 < kNumEdges; ed1++) {
-      for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
-        for (int i = 0; i < 3; i++) {
-          for (int j = 0; j < 3; j++) {
-            trTT[kNumEdges*ed1+ed2] += basis[ed1][3*i+j] * basis[ed2][3*j+i];
-          }
-        }
-      }
-    }
-
     // material parameters
     mjtNum mu = E / (2*(1+nu)) * volume;
     mjtNum la = E*nu / ((1+nu)*(1-2*nu)) * volume;
 
-    // assembly of strain metric tensor
+    // compute metric tensor
-    for (int ed1 = 0; ed1 < kNumEdges; ed1++) {
+    MetricTensor<Stencil3D>(metric, t, mu, la, basis);
-      for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
-        int index = kNumEdges*kNumEdges*t + kNumEdges*ed1 + ed2;
-        metric[index] = mu * trTT[kNumEdges*ed1+ed2] + la * trT[ed2]*trT[ed1];
-      }
-    }
   }
 
   // allocate array
+  ne = edges.size();
   reference.assign(ne, 0);
   deformed.assign(ne, 0);
   previous.assign(ne, 0);
@@ -266,7 +183,7 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) {
 
     // compute length gradient with respect to dofs
     mjtNum gradient[kNumEdges][2][3];
-    GradSquaredLengths(gradient, d->xpos+3*i0, v, edge);
+    GradSquaredLengths<Stencil3D>(gradient, d->xpos+3*i0, v);
 
     // we add generalized Rayleigh damping as decribed in Section 5.2 of
     // Kharevych et al., "Geometric, Variational Integrators for Computer
@@ -299,7 +216,7 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) {
       for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
         for (int i = 0; i < 2; i++) {
           for (int x = 0; x < 3; x++) {
-            force[3 * edge[ed2][i] + x] +=
+            force[3 * Stencil3D::edge[ed2][i] + x] +=
                 elongation[ed1] * gradient[ed2][i][x] *
                 metric[offset * t + kNumEdges * ed1 + ed2];
           }

plugin/elasticity/solid.h

@@ -21,17 +21,11 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
+#include "elasticity.h"
 
 
 namespace mujoco::plugin::elasticity {
 
-struct Stencil3D {
-  static constexpr int kNumEdges = 6;
-  static constexpr int kNumVerts = 4;
-  int vertices[kNumVerts];
-  int edges[kNumEdges];
-};
-
 class Solid {
  public:
   // Returns a new Solid instance or nullopt on failure.
@@ -67,9 +61,6 @@ class Solid {
   Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
         mjtNum damp, const std::vector<int>& simplex,
         const std::vector<int>& edgeidx);
-
-  void CreateStencils(const std::vector<int>& simplex,
-                      const std::vector<int>& edgeidx);
 };
 
 }  // namespace mujoco::plugin::elasticity

src/user/user_flexcomp.cc

@@ -512,6 +512,7 @@ bool mjCFlexcomp::MakeGrid(char* error, int error_sz) {
 
   // 2D
   else if (dim==2) {
+    int quad2tri[2][3] = {{0, 1, 2}, {0, 2, 3}};
     for (int ix=0; ix<count[0]; ix++) {
       for (int iy=0; iy<count[1]; iy++) {
         // add point
@@ -527,13 +528,17 @@ bool mjCFlexcomp::MakeGrid(char* error, int error_sz) {
 
         // add elements
         if (ix<count[0]-1 && iy<count[1]-1) {
-          element.push_back(GridID(ix, iy));
+          int vert[4] = {
-          element.push_back(GridID(ix+1, iy));
+            count[2]*count[1]*(ix+0) + count[2]*(iy+0),
-          element.push_back(GridID(ix+1, iy+1));
+            count[2]*count[1]*(ix+1) + count[2]*(iy+0),
-
+            count[2]*count[1]*(ix+1) + count[2]*(iy+1),
-          element.push_back(GridID(ix, iy));
+            count[2]*count[1]*(ix+0) + count[2]*(iy+1),
-          element.push_back(GridID(ix+1, iy+1));
+          };
-          element.push_back(GridID(ix, iy+1));
+          for (int s = 0; s < 2; s++) {
+            for (int v = 0; v < 3; v++) {
+              element.push_back(vert[quad2tri[s][v]]);
+            }
+          }
         }
       }
     }

test/engine/engine_plugin_test.cc

@@ -32,7 +32,7 @@ namespace {
 using ::testing::HasSubstr;
 using ::testing::NotNull;
 
-constexpr int kNumTruePlugins = 10;
+constexpr int kNumTruePlugins = 11;
 constexpr int kNumFakePlugins = 30;
 constexpr int kNumTestPlugins = 3;
 

test/plugin/elasticity/elasticity_test.cc

@@ -22,6 +22,7 @@
 #include <gtest/gtest.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
+#include "plugin/elasticity/membrane.h"
 #include "plugin/elasticity/shell.h"
 #include "plugin/elasticity/solid.h"
 
@@ -84,6 +85,65 @@ TEST_F(ElasticityTest, ElasticEnergyShell) {
   mj_deleteModel(m);
 }
 
+// -------------------------------- membrane -----------------------------------
+TEST_F(PluginTest, ElasticEnergyMembrane) {
+  static constexpr char cantilever_xml[] = R"(
+  <mujoco>
+  <extension>
+    <plugin plugin="mujoco.elasticity.membrane"/>
+  </extension>
+
+  <worldbody>
+    <flexcomp type="grid" count="8 8 1" spacing="1 1 1"
+              radius=".025" name="test" dim="2">
+      <plugin plugin="mujoco.elasticity.membrane">
+        <config key="poisson" value="0"/>
+        <config key="young" value="2"/>
+        <config key="thickness" value="1"/>
+      </plugin>
+      <edge equality="false"/>
+    </flexcomp>
+  </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjModel* m = LoadModelFromString(cantilever_xml, error, sizeof(error));
+  ASSERT_THAT(m, testing::NotNull()) << error;
+  mjData* d = mj_makeData(m);
+  auto* membrane =
+      reinterpret_cast<plugin::elasticity::Membrane*>(d->plugin_data[0]);
+
+  mj_kinematics(m, d);
+  mj_flex(m, d);
+
+  // check that if the entire geometry is rescaled by a factor "scale", then
+  // trace(strain^2) = 2*scale^2
+
+  for (mjtNum scale = 1; scale < 4; scale++) {
+    for (int t = 0; t < membrane->nt; t++) {
+      mjtNum energy = 0;
+      mjtNum volume = 1./2.;
+      for (int e1 = 0; e1 < 3; e1++) {
+        for (int e2 = 0; 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 =
+              scale * m->flexedge_length0[idx1] * m->flexedge_length0[idx1];
+          mjtNum elongation2 =
+              scale * m->flexedge_length0[idx2] * m->flexedge_length0[idx2];
+          energy += membrane->metric[9*t+3*e2+e1] * elongation1 * elongation2;
+        }
+      }
+      EXPECT_NEAR(
+        4*energy/volume, 2*scale*scale, std::numeric_limits<float>::epsilon());
+    }
+  }
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 // -------------------------------- solid -----------------------------------
 TEST_F(ElasticityTest, ElasticEnergySolid) {
   static constexpr char cantilever_xml[] = R"(