Add edge ids to flex elements in mjModel.

PiperOrigin-RevId: 675130981
Change-Id: If2ccfb268142d235273b7a50e15740f6399903b2

plugin/elasticity/elasticity.cc

@@ -19,72 +19,11 @@
 #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()) {  // SHOULD NOT OCCUR
-        if (elements[t].edges[e] != edgeidx[T::kNumEdges*t+e]) {
-          mju_error("edge ordering is incoherent between flex and plugin");
-        }
-      }
-    }
-  }
-
-  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

@@ -65,20 +65,19 @@ 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>& elongationglob,
                          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];
 
-  if (elements.size() != m->flex_elemnum[flex]) {
-    mju_error("plugin stencil does not match flex stencil");
-  }
+  int dim = m->flex_dim[flex];
+  const int* elem = m->flex_elem + m->flex_elemdataadr[flex];
+  const int* edgeelem = m->flex_elemedge + m->flex_elemedgeadr[flex];
 
   // compute force element-by-element
   for (int t = 0; t < m->flex_elemnum[flex]; t++)  {
-    const int* v = elements[t].vertices;
+    const int* v = elem + (dim+1) * t;
 
     // compute length gradient with respect to dofs
     mjtNum gradient[T::kNumEdges][2][3];
@@ -87,7 +86,7 @@ inline void ComputeForce(std::vector<mjtNum>& qfrc_passive,
     // extract elongation of edges belonging to this element
     mjtNum elongation[T::kNumEdges];
     for (int e = 0; e < T::kNumEdges; e++) {
-      int idx = elements[t].edges[e];
+      int idx = edgeelem[t * T::kNumEdges + e];
       elongation[e] = elongationglob[idx];
     }
 

plugin/elasticity/membrane.cc

@@ -87,10 +87,7 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
         strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
     mjtNum damp =
             strtod(mj_getPluginConfig(m, instance, "damping"), 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, damp, face, edge);
+    return Membrane(m, d, instance, nu, E, thick, damp);
   } else {
     mju_warning("Invalid parameter specification in shell plugin");
     return std::nullopt;
@@ -99,9 +96,7 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
 
 // plugin constructor
 Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
-                   mjtNum E, mjtNum thick, mjtNum damp,
-                   const std::vector<int>& simplex,
-                   const std::vector<int>& edgeidx)
+                   mjtNum E, mjtNum thick, mjtNum damp)
     : f0(-1), damping(damp), thickness(thick) {
   // count plugin bodies
   nv = ne = 0;
@@ -126,12 +121,10 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   // vertex positions
   mjtNum* body_pos = m->flex_xvert0 + 3*m->flex_vertadr[f0];
 
-  // generate triangles from the vertices
-  nt = CreateStencils<Stencil2D>(elements, edges, simplex, edgeidx);
-
   // loop over all triangles
-  for (int t = 0; t < nt; t++) {
-    int* v = elements[t].vertices;
+  const int* elem = m->flex_elem + m->flex_elemdataadr[f0];
+  for (int t = 0; t < m->flex_elemnum[f0]; t++) {
+    const int* v = elem + (m->flex_dim[f0]+1) * t;
     for (int i = 0; i < kNumVerts; i++) {
       int bi = m->flex_vertbodyid[m->flex_vertadr[f0]+v[i]];
       if (bi && m->body_plugin[bi] != instance) {
@@ -163,7 +156,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   }
 
   // allocate array
-  ne = edges.size();
+  ne = m->flex_edgenum[f0];
   elongation.assign(ne, 0);
   force.assign(3*nv, 0);
 }
@@ -195,7 +188,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, elongation, m, f0, xpos);
+  ComputeForce<Stencil2D>(force, elongation, m, f0, xpos);
 
   // insert into passive force
   AddFlexForce(qfrc, force, m, d, xpos, f0);

plugin/elasticity/membrane.h

@@ -44,13 +44,8 @@ class Membrane {
   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> prev;                 // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
@@ -61,8 +56,7 @@ class Membrane {
 
  private:
   Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-           mjtNum thick, mjtNum damp, const std::vector<int>& simplex,
-           const std::vector<int>& edgeidx);
+           mjtNum thick, mjtNum damp);
 };
 
 }  // namespace mujoco::plugin::elasticity

plugin/elasticity/solid.cc

@@ -93,10 +93,7 @@ std::optional<Solid> Solid::Create(const mjModel* m, mjData* d, int instance) {
         mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
         mjtNum damp =
             strtod(mj_getPluginConfig(m, instance, "damping"), nullptr);
-        std::vector<int> face, edge;
-        String2Vector(mj_getPluginConfig(m, instance, "face"), face);
-        String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
-        return Solid(m, d, instance, nu, E, damp, face, edge);
+        return Solid(m, d, instance, nu, E, damp);
     } else {
         mju_warning("Invalid parameter specification in solid plugin");
         return std::nullopt;
@@ -105,8 +102,7 @@ std::optional<Solid> Solid::Create(const mjModel* m, mjData* d, int instance) {
 
 // plugin constructor
 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)
+             mjtNum damp)
     : f0(-1), damping(damp) {
   // count plugin bodies
   nv = ne = 0;
@@ -134,12 +130,10 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   // vertex positions
   mjtNum* body_pos = m->flex_xvert0 + 3*m->flex_vertadr[f0];
 
-  // generate tetrahedra from the vertices
-  nt = CreateStencils<Stencil3D>(elements, edges, simplex, edgeidx);
-
   // loop over all tetrahedra
-  for (int t = 0; t < nt; t++) {
-    int* v = elements[t].vertices;
+  const int* elem = m->flex_elem + m->flex_elemdataadr[f0];
+  for (int t = 0; t < m->flex_elemnum[f0]; t++) {
+    const int* v = elem + (m->flex_dim[f0]+1) * t;
     for (int i = 0; i < kNumVerts; i++) {
       int bi = m->flex_vertbodyid[m->flex_vertadr[f0]+v[i]];
       if (bi && m->body_plugin[bi] != instance) {
@@ -170,7 +164,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   }
 
   // allocate array
-  ne = edges.size();
+  ne = m->flex_edgenum[f0];
   elongation.assign(ne, 0);
   force.assign(3*nv, 0);
 }
@@ -202,7 +196,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, elongation, m, f0, xpos);
+  ComputeForce<Stencil3D>(force, elongation, m, f0, xpos);
 
   // insert into passive force
   AddFlexForce(qfrc, force, m, d, xpos, f0);

plugin/elasticity/solid.h

@@ -42,13 +42,8 @@ class Solid {
   int i0;  // index of first body
   int nc;  // number of cubes in the grid
   int nv;  // number of vertices (bodies) in the solid
-  int nt;  // number of volumetric elements (tetrahedra)
   int ne;  // number of edges in the solid
 
-  // connectivity info for mapping tetrahedra to edges and vertices
-  std::vector<Stencil3D> elements;          // 4 vertices and 6 edges (nt x 10)
-  std::vector<std::pair<int, int> > edges;  // edge to vertex map     (ne x 2)
-
   // precomputed quantities
   std::vector<mjtNum> prev;                 // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
@@ -58,8 +53,7 @@ class Solid {
 
  private:
   Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-        mjtNum damp, const std::vector<int>& simplex,
-        const std::vector<int>& edgeidx);
+        mjtNum damp);
 };
 
 }  // namespace mujoco::plugin::elasticity