chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add flex elasticity stiffness matrix to mjModel.

Browse Source 
PiperOrigin-RevId: 675101645
Change-Id: Ic169c108b9eece3657ea4ec211df357abe615b7e
This commit is contained in:
Alessio Quaglino
2024-09-16 04:59:55 -07:00
committed by Copybara-Service
parent c77babe046
commit 2da15cf137
12 changed files with 71 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+1
View File
@@ -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)
include/mujoco/mjmodel.h
+1
View File
@@ -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)
include/mujoco/mjxmacro.h
+1
View File
@@ -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 ) \
introspect/structs.py
+7
View File
@@ -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(
plugin/elasticity/elasticity.h
+36 -9
View File
@@ -66,13 +66,18 @@ 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>& metric,
const std::vector<mjtNum>& 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<mjtNum>& 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 <typename T>
void inline MetricTensor(std::vector<mjtNum>& 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<mjtNum>& 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
plugin/elasticity/membrane.cc
+4 -5
View File
@@ -129,9 +129,6 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
// 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;
@@ -160,7 +157,9 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
}
// compute metric tensor
MetricTensor<Stencil2D>(metric, t, mu, la, basis);
// TODO: do not write in a const mjModel
MetricTensor<Stencil2D>(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<Stencil2D>(force, elements, metric, elongation, m, xpos);
ComputeForce<Stencil2D>(force, elements, elongation, m, f0, xpos);
// insert into passive force
AddFlexForce(qfrc, force, m, d, xpos, f0);
plugin/elasticity/membrane.h
-1
View File
@@ -52,7 +52,6 @@ class Membrane {
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)
std::vector<mjtNum> prev; // previous-step lengths (ne x 1)
std::vector<mjtNum> elongation; // edge elongation (ne x 1)
std::vector<mjtNum> force; // force at all vertices (nv x 3)
plugin/elasticity/solid.cc
+4 -5
View File
@@ -137,9 +137,6 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
// generate tetrahedra from the vertices
nt = CreateStencils<Stencil3D>(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<Stencil3D>(metric, t, mu, la, basis);
// TODO: do not write in a const mjModel
MetricTensor<Stencil3D>(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<Stencil3D>(force, elements, metric, elongation, m, xpos);
ComputeForce<Stencil3D>(force, elements, elongation, m, f0, xpos);
// insert into passive force
AddFlexForce(qfrc, force, m, d, xpos, f0);
plugin/elasticity/solid.h
-1
View File
@@ -50,7 +50,6 @@ class Solid {
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 36)
std::vector<mjtNum> prev; // previous-step lengths (ne x 1)
std::vector<mjtNum> elongation; // edge elongation (ne x 1)
std::vector<mjtNum> force; // force at all vertices (nv x 3)
src/user/user_model.cc
+4
View File
@@ -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;
test/plugin/elasticity/elasticity_test.cc
+12 -8
View File
@@ -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(
unity/Runtime/Bindings/MjBindings.cs
+1
View File
@@ -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;