From 34e8ff1aad4a6311be6455d703a3f497d4cfa077 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 14 May 2025 02:49:46 -0700 Subject: [PATCH] Add flex edge flap connectivity to mjModel. PiperOrigin-RevId: 758593152 Change-Id: I79836df961972c9eae8037e5fcbfc3d0122645de --- doc/includes/references.h | 1 + include/mujoco/mjmodel.h | 1 + include/mujoco/mjxmacro.h | 1 + plugin/elasticity/shell.cc | 105 ++++++---------------- plugin/elasticity/shell.h | 13 +-- python/mujoco/introspect/structs.py | 8 ++ src/user/user_mesh.cc | 68 ++++++++++++-- src/user/user_model.cc | 7 ++ src/user/user_objects.h | 6 ++ test/plugin/elasticity/elasticity_test.cc | 6 +- unity/Runtime/Bindings/MjBindings.cs | 1 + 11 files changed, 122 insertions(+), 95 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index eada1572..ffada3b8 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -1229,6 +1229,7 @@ struct mjModel_ { int* flex_nodebodyid; // node body ids (nflexnode x 1) int* flex_vertbodyid; // vertex body ids (nflexvert x 1) int* flex_edge; // edge vertex ids (2 per edge) (nflexedge x 2) + int* flex_edgeflap; // adjacent vertex ids (dim=2 only) (nflexedge x 2) int* flex_elem; // element vertex ids (dim+1 per elem) (nflexelemdata x 1) int* flex_elemtexcoord; // element texture coordinates (dim+1) (nflexelemdata x 1) int* flex_elemedge; // element edge ids (nflexelemedge x 1) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index ceaa76d1..fd63b272 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -896,6 +896,7 @@ struct mjModel_ { int* flex_nodebodyid; // node body ids (nflexnode x 1) int* flex_vertbodyid; // vertex body ids (nflexvert x 1) int* flex_edge; // edge vertex ids (2 per edge) (nflexedge x 2) + int* flex_edgeflap; // adjacent vertex ids (dim=2 only) (nflexedge x 2) int* flex_elem; // element vertex ids (dim+1 per elem) (nflexelemdata x 1) int* flex_elemtexcoord; // element texture coordinates (dim+1) (nflexelemdata x 1) int* flex_elemedge; // element edge ids (nflexelemedge x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index aa55b3d7..f46b4ff5 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -349,6 +349,7 @@ XMJV( int, flex_nodebodyid, nflexnode, 1 ) \ X ( int, flex_vertbodyid, nflexvert, 1 ) \ X ( int, flex_edge, nflexedge, 2 ) \ + X ( int, flex_edgeflap, nflexedge, 2 ) \ XMJV( int, flex_elem, nflexelemdata, 1 ) \ XMJV( int, flex_elemtexcoord, nflexelemdata, 1 ) \ X ( int, flex_elemedge, nflexelemedge, 1 ) \ diff --git a/plugin/elasticity/shell.cc b/plugin/elasticity/shell.cc index 30b02051..772b06eb 100644 --- a/plugin/elasticity/shell.cc +++ b/plugin/elasticity/shell.cc @@ -12,14 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include #include #include #include #include #include -#include -#include #include #include @@ -34,9 +31,7 @@ namespace mujoco::plugin::elasticity { namespace { // local tetrahedron numbering -constexpr int kNumEdges = Stencil2D::kNumEdges; constexpr int kNumVerts = Stencil2D::kNumVerts; -constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}}; // cotangent between two edges mjtNum cot(mjtNum* x, int v0, int v1, int v2) { @@ -68,81 +63,26 @@ mjtNum ComputeVolume(const mjtNum* x, const int v[kNumVerts]) { // factory function std::optional Shell::Create(const mjModel* m, mjData* d, int instance) { - if (CheckAttr("face", m, instance) && - CheckAttr("edge", m, instance) && - CheckAttr("poisson", m, instance) && + if (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 face, edge; - String2Vector(mj_getPluginConfig(m, instance, "face"), face); - String2Vector(mj_getPluginConfig(m, instance, "edge"), edge); - return Shell(m, d, instance, nu, E, thick, face, edge); + return Shell(m, d, instance, nu, E, thick); } else { mju_warning("Invalid parameter specification in shell plugin"); return std::nullopt; } } -// create map from triangles to vertices and edges and from edges to vertices -void Shell::CreateStencils(const std::vector& simplex, - const std::vector& 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, int, PairHash> edge_indices; - - // loop over all triangles - 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) { - StencilFlap flap; - flap.vertices[0] = v[edge[e][0]]; - flap.vertices[1] = v[edge[e][1]]; - flap.vertices[2] = v[(edge[e][1]+1) % 3]; - flap.vertices[3] = -1; - flaps.push_back(flap); - elements[t].edges[e] = ne++; - } else { - elements[t].edges[e] = it->second; - flaps[it->second].vertices[3] = v[(edge[e][1]+1) % 3]; - } - - if (!edgeidx.empty()) { - assert(elements[t].edges[e] == edgeidx[kNumEdges*t+e]); - } - } - } -} - // plugin constructor Shell::Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, - mjtNum thick, const std::vector& face, - const std::vector& edgeidx) - : thickness(thick) { + mjtNum thick) + : f0(-1), thickness(thick) { // count plugin bodies - nv = ne = 0; + nv = 0; for (int i = 1; i < m->nbody; i++) { if (m->body_plugin[i] == instance) { if (!nv++) { @@ -151,15 +91,25 @@ Shell::Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, } } - // generate triangles from the vertices - CreateStencils(face, edgeidx); + // count flexes + for (int i = 0; i < m->nflex; i++) { + for (int j = 0; j < m->flex_vertnum[i]; j++) { + if (m->flex_vertbodyid[m->flex_vertadr[i]+j] == i0) { + f0 = i; + nv = m->flex_vertnum[f0]; + if (m->flex_dim[i] != 2) { // SHOULD NOT OCCUR + mju_error("mujoco.elasticity.shell requires a 2D mesh"); + } + } + } + } // material parameters mjtNum mu = E / (2*(1+nu)); // loop over all triangles - for (int t = 0; t < nt; t++) { - int* v = elements[t].vertices; + for (int t = 0; t < m->flex_elemnum[f0]; t++) { + int* v = m->flex_elem + 3*(t+m->flex_elemadr[f0]); 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"); @@ -169,14 +119,16 @@ Shell::Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, // allocate array position.assign(nv*3, 0); - bending.assign(ne*16, 0); + bending.assign(m->flex_edgenum[f0]*16, 0); // store previous positions mju_copy(position.data(), m->body_pos+3*i0, 3*nv); // assemble bending Hessian - for (int e = 0; e < ne; e++) { - int* v = flaps[e].vertices; + for (int e = 0; e < m->flex_edgenum[f0]; e++) { + int* edge = m->flex_edge + 2*(e+m->flex_edgeadr[f0]); + int* flap = m->flex_edgeflap + 2*(e+m->flex_edgeadr[f0]); + int v[4] = {edge[0], edge[1], flap[0], flap[1]}; int vadj[3] = {v[1], v[0], v[3]}; if (v[3]== -1) { @@ -205,8 +157,10 @@ Shell::Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, } void Shell::Compute(const mjModel* m, mjData* d, int instance) { - for (int e = 0; e < ne; e++) { - int* v = flaps[e].vertices; + for (int e = 0; e < m->flex_edgenum[f0]; e++) { + int* edge = m->flex_edge + 2*(e+m->flex_edgeadr[f0]); + int* flap = m->flex_edgeflap + 2*(e+m->flex_edgeadr[f0]); + int v[4] = {edge[0], edge[1], flap[0], flap[1]}; mjtNum force[12] = {0}; if (v[3] == -1) { // skip boundary edges @@ -241,8 +195,7 @@ void Shell::RegisterPlugin() { plugin.name = "mujoco.elasticity.shell"; plugin.capabilityflags |= mjPLUGIN_PASSIVE; - const char* attributes[] = {"face", "edge", "young", - "poisson", "thickness", "damping"}; + const char* attributes[] = {"young", "poisson", "thickness", "damping"}; plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]); plugin.attributes = attributes; plugin.nstate = +[](const mjModel* m, int instance) { return 0; }; diff --git a/plugin/elasticity/shell.h b/plugin/elasticity/shell.h index f92fffbb..e8496e70 100644 --- a/plugin/elasticity/shell.h +++ b/plugin/elasticity/shell.h @@ -45,14 +45,9 @@ class Shell { static void RegisterPlugin(); int i0; // index of first body + int f0; // index of corresponding flex int nc; // number of quads in the grid int nv; // number of vertices (bodies) in the Shell - int nt; // number of area elements (triangles) - int ne; // number of edges in the Shell - - // connectivity info for mapping tetrahedra to edges and vertices - std::vector elements; // triangles (nt x 6) - std::vector flaps; // adjacent triangles (ne x 4) // precomputed quantities std::vector position; // previous-step positions (nv x 3) @@ -62,11 +57,7 @@ class Shell { private: Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, - mjtNum thick, const std::vector& face, - const std::vector& edgeidx); - - void CreateStencils(const std::vector& simplex, - const std::vector& edgeidx); + mjtNum thick); }; } // namespace mujoco::plugin::elasticity diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index e9e6f801..4efca1ee 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -2619,6 +2619,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='edge vertex ids (2 per edge)', array_extent=('nflexedge', 2), ), + StructFieldDecl( + name='flex_edgeflap', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='adjacent vertex ids (dim=2 only)', + array_extent=('nflexedge', 2), + ), StructFieldDecl( name='flex_elem', type=PointerType( diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index 93c10004..b44ecff5 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -3018,6 +3018,63 @@ void inline ComputeStiffness(std::vector& stiffness, MetricTensor(stiffness.data(), t, mu, la, basis); } +// local tetrahedron numbering +constexpr int kNumEdges = Stencil2D::kNumEdges; +constexpr int kNumVerts = Stencil2D::kNumVerts; +constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}}; + +// create map from triangles to vertices and edges and from edges to vertices +static void CreateFlapStencil(std::vector& flaps, + const std::vector& simplex, + const std::vector& edgeidx) { + // populate stencil + int ne = 0; + int nt = simplex.size() / kNumVerts; + std::vector elements(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, int, PairHash> edge_indices; + + // loop over all triangles + 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) { + StencilFlap flap; + flap.vertices[0] = v[edge[e][0]]; + flap.vertices[1] = v[edge[e][1]]; + flap.vertices[2] = v[(edge[e][1] + 1) % 3]; + flap.vertices[3] = -1; + flaps.push_back(flap); + elements[t].edges[e] = ne++; + } else { + elements[t].edges[e] = it->second; + flaps[it->second].vertices[3] = v[(edge[e][1] + 1) % 3]; + } + + // double check that the edge indices are consistent + if (!edgeidx.empty()) { + if (elements[t].edges[e] != edgeidx[kNumEdges * t + e]) { + mju_error("edge indices do not match in CreateFlapStencil"); + } + } + } + } +} + //----------------------------- linear elasticity -------------------------------------------------- // Gauss Legendre quadrature points in 1 dimension on the interval [a, b] @@ -3543,10 +3600,6 @@ void mjCFlex::Compile(const mjVFS* vfs) { } // add plugins - std::string userface, useredge; - userface = VectorToString(elem_); - useredge = VectorToString(edgeidx_); - for (const auto& vbodyid : vertbodyid) { if (vbodyid < 0) { continue; @@ -3557,11 +3610,14 @@ void mjCFlex::Compile(const mjVFS* vfs) { if (damping > 0) { plugin_instance->config_attribs["damping"] = std::to_string(damping); } - plugin_instance->config_attribs["face"] = userface; - plugin_instance->config_attribs["edge"] = useredge; } } + // create flap stencil + if (dim == 2) { + CreateFlapStencil(flaps, elem_, edgeidx_); + } + // create shell fragments and element-vertex collision pairs CreateShellPair(); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index ff78e190..93449339 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -3158,6 +3158,13 @@ void mjCModel::CopyObjects(mjModel* m) { for (int k=0; k < pfl->nedge; k++) { m->flex_edge[2*(edge_adr+k)] = pfl->edge[k].first; m->flex_edge[2*(edge_adr+k)+1] = pfl->edge[k].second; + if (pfl->dim == 2) { + m->flex_edgeflap[2*(edge_adr+k)+0] = pfl->flaps[k].vertices[2]; + m->flex_edgeflap[2*(edge_adr+k)+1] = pfl->flaps[k].vertices[3]; + } else { + m->flex_edgeflap[2*(edge_adr+k)+0] = -1; + m->flex_edgeflap[2*(edge_adr+k)+1] = -1; + } if (pfl->rigid) { m->flexedge_rigid[edge_adr+k] = 1; diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 9cdae417..585d6abb 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -830,6 +830,11 @@ class mjCLight : public mjCLight_, private mjsLight { //------------------------- class mjCFlex ---------------------------------------------------------- // Describes a flex +struct StencilFlap { + static constexpr int kNumVerts = 4; + int vertices[kNumVerts]; +}; + class mjCFlex_ : public mjCBase { protected: int nvert; // number of vertices @@ -846,6 +851,7 @@ class mjCFlex_ : public mjCBase { std::vector shell; // shell fragment vertex ids (dim per fragment) std::vector elemlayer; // element layer (distance from border) std::vector evpair; // element-vertex pairs + std::vector flaps; // adjacent triangles std::vector vertxpos; // global vertex positions mjCBoundingVolumeHierarchy tree; // bounding volume hierarchy std::vector elemaabb_; // element bounding volume diff --git a/test/plugin/elasticity/elasticity_test.cc b/test/plugin/elasticity/elasticity_test.cc index f38559b5..16273900 100644 --- a/test/plugin/elasticity/elasticity_test.cc +++ b/test/plugin/elasticity/elasticity_test.cc @@ -83,8 +83,10 @@ TEST_F(ElasticityTest, ElasticEnergyShell) { // check that a plane is in the kernel of the energy for (mjtNum scale = 1; scale < 4; scale++) { - for (int e = 0; e < shell->ne; e++) { - int* v = shell->flaps[e].vertices; + for (int e = 0; e < m->flex_edgenum[0]; e++) { + int* edge = m->flex_edge + 2*(m->flex_edgeadr[0] + e); + int* flap = m->flex_edgeflap + 2*(m->flex_edgeadr[0] + e); + int v[4] = {edge[0], edge[1], flap[0], flap[1]}; if (v[3]== -1) { continue; } diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index d26592a9..d743436c 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5480,6 +5480,7 @@ public unsafe struct mjModel_ { public int* flex_nodebodyid; public int* flex_vertbodyid; public int* flex_edge; + public int* flex_edgeflap; public int* flex_elem; public int* flex_elemtexcoord; public int* flex_elemedge;