Add membrane plugin (2D Flex elastic stiffness).

PiperOrigin-RevId: 574109411
Change-Id: I4a701d8189cecf540cd200bd30cb4582f3c7dd43

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