From 2da15cf137667fd76cd06fa48fd7c8d43ae40603 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Mon, 16 Sep 2024 04:59:55 -0700 Subject: [PATCH] Add flex elasticity stiffness matrix to mjModel. PiperOrigin-RevId: 675101645 Change-Id: Ic169c108b9eece3657ea4ec211df357abe615b7e --- doc/includes/references.h | 1 + include/mujoco/mjmodel.h | 1 + include/mujoco/mjxmacro.h | 1 + introspect/structs.py | 7 ++++ plugin/elasticity/elasticity.h | 45 ++++++++++++++++++----- plugin/elasticity/membrane.cc | 9 ++--- plugin/elasticity/membrane.h | 1 - plugin/elasticity/solid.cc | 9 ++--- plugin/elasticity/solid.h | 1 - src/user/user_model.cc | 4 ++ test/plugin/elasticity/elasticity_test.cc | 20 ++++++---- unity/Runtime/Bindings/MjBindings.cs | 1 + 12 files changed, 71 insertions(+), 29 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index f6407f6d..840faa3e 100644 --- a/doc/includes/references.h +++ b/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) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 349dffe0..d5e08bb5 100644 --- a/include/mujoco/mjmodel.h +++ b/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) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 56b17584..d32729c2 100644 --- a/include/mujoco/mjxmacro.h +++ b/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 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index 444b8398..d5216fd1 100644 --- a/introspect/structs.py +++ b/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( diff --git a/plugin/elasticity/elasticity.h b/plugin/elasticity/elasticity.h index d12b5a82..155c4279 100644 --- a/plugin/elasticity/elasticity.h +++ b/plugin/elasticity/elasticity.h @@ -66,13 +66,18 @@ void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3], template inline void ComputeForce(std::vector& qfrc_passive, const std::vector& elements, - const std::vector& metric, const std::vector& 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& 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 -void inline MetricTensor(std::vector& 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& 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; - metric[index] = mu * trEE[T::kNumEdges * ed1 + ed2] + - la * trE[ed2] * trE[ed1]; + k[T::kNumEdges*ed1 + ed2] = mu * trEE[T::kNumEdges * ed1 + ed2] + + 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 diff --git a/plugin/elasticity/membrane.cc b/plugin/elasticity/membrane.cc index da7c17d5..5c8148ad 100644 --- a/plugin/elasticity/membrane.cc +++ b/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(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(metric, t, mu, la, basis); + // TODO: do not write in a const mjModel + MetricTensor(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(force, elements, metric, elongation, m, xpos); + ComputeForce(force, elements, elongation, m, f0, xpos); // insert into passive force AddFlexForce(qfrc, force, m, d, xpos, f0); diff --git a/plugin/elasticity/membrane.h b/plugin/elasticity/membrane.h index 815c0a5b..b06819a4 100644 --- a/plugin/elasticity/membrane.h +++ b/plugin/elasticity/membrane.h @@ -52,7 +52,6 @@ class Membrane { std::vector > edges; // edge to vertex map (ne x 2) // precomputed quantities - std::vector metric; // geom-induced metric (nt x 9) std::vector prev; // previous-step lengths (ne x 1) std::vector elongation; // edge elongation (ne x 1) std::vector force; // force at all vertices (nv x 3) diff --git a/plugin/elasticity/solid.cc b/plugin/elasticity/solid.cc index d9fce036..537ab714 100644 --- a/plugin/elasticity/solid.cc +++ b/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(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(metric, t, mu, la, basis); + // TODO: do not write in a const mjModel + MetricTensor(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(force, elements, metric, elongation, m, xpos); + ComputeForce(force, elements, elongation, m, f0, xpos); // insert into passive force AddFlexForce(qfrc, force, m, d, xpos, f0); diff --git a/plugin/elasticity/solid.h b/plugin/elasticity/solid.h index 7b282b42..d68f29cd 100644 --- a/plugin/elasticity/solid.h +++ b/plugin/elasticity/solid.h @@ -50,7 +50,6 @@ class Solid { std::vector > edges; // edge to vertex map (ne x 2) // precomputed quantities - std::vector metric; // geom-induced metric (nt x 36) std::vector prev; // previous-step lengths (ne x 1) std::vector elongation; // edge elongation (ne x 1) std::vector force; // force at all vertices (nv x 3) diff --git a/src/user/user_model.cc b/src/user/user_model.cc index b1049ff3..09da35e3 100644 --- a/src/user/user_model.cc +++ b/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; diff --git a/test/plugin/elasticity/elasticity_test.cc b/test/plugin/elasticity/elasticity_test.cc index 2cdd410c..bc87203a 100644 --- a/test/plugin/elasticity/elasticity_test.cc +++ b/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( diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 70eb8412..bd1a56e6 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/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;