Alessio Quaglino
2da15cf137
PiperOrigin-RevId: 675101645 Change-Id: Ic169c108b9eece3657ea4ec211df357abe615b7e
218 lines
6.9 KiB
C++
218 lines
6.9 KiB
C++
// Copyright 2022 DeepMind Technologies Limited
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#ifndef MUJOCO_PLUGIN_ELASTICITY_ELASTICITY_H_
|
|
#define MUJOCO_PLUGIN_ELASTICITY_ELASTICITY_H_
|
|
|
|
#include <cstddef>
|
|
#include <functional>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include <mujoco/mujoco.h>
|
|
|
|
namespace mujoco::plugin::elasticity {
|
|
|
|
struct PairHash
|
|
{
|
|
template <class T1, class T2>
|
|
std::size_t operator() (const std::pair<T1, T2>& pair) const {
|
|
return std::hash<T1>()(pair.first) ^ std::hash<T2>()(pair.second);
|
|
}
|
|
};
|
|
|
|
struct Stencil2D {
|
|
static constexpr int kNumEdges = 3;
|
|
static constexpr int kNumVerts = 3;
|
|
static constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}};
|
|
int vertices[kNumVerts];
|
|
int edges[kNumEdges];
|
|
};
|
|
|
|
struct Stencil3D {
|
|
static constexpr int kNumEdges = 6;
|
|
static constexpr int kNumVerts = 4;
|
|
static constexpr int edge[kNumEdges][2] = {{0, 1}, {1, 2}, {2, 0},
|
|
{2, 3}, {0, 3}, {1, 3}};
|
|
int vertices[kNumVerts];
|
|
int edges[kNumEdges];
|
|
};
|
|
|
|
// gradients of edge lengths with respect to vertex positions
|
|
template <typename T>
|
|
void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3],
|
|
const mjtNum* x,
|
|
const int v[T::kNumVerts]) {
|
|
for (int e = 0; e < T::kNumEdges; e++) {
|
|
for (int d = 0; d < 3; d++) {
|
|
gradient[e][0][d] = x[3*v[T::edge[e][0]]+d] - x[3*v[T::edge[e][1]]+d];
|
|
gradient[e][1][d] = x[3*v[T::edge[e][1]]+d] - x[3*v[T::edge[e][0]]+d];
|
|
}
|
|
}
|
|
}
|
|
|
|
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");
|
|
}
|
|
|
|
// 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
|
|
mjtNum gradient[T::kNumEdges][2][3];
|
|
GradSquaredLengths<T>(gradient, xpos, v);
|
|
|
|
// 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];
|
|
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};
|
|
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[T::kNumEdges * ed1 + ed2];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// insert into global force
|
|
for (int i = 0; i < T::kNumVerts; i++) {
|
|
for (int x = 0; x < 3; x++) {
|
|
qfrc_passive[3*v[i]+x] += force[3*i+x];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// add flex force to degrees of freedom
|
|
inline void AddFlexForce(mjtNum* qfrc,
|
|
const std::vector<mjtNum>& force,
|
|
const mjModel* m, mjData* d,
|
|
const mjtNum* xpos,
|
|
int f0) {
|
|
int* bodyid = m->flex_vertbodyid + m->flex_vertadr[f0];
|
|
|
|
for (int v = 0; v < m->flex_vertnum[f0]; v++) {
|
|
int bid = bodyid[v];
|
|
if (m->body_simple[bid] != 2) {
|
|
// this should only occur for pinned flex vertices
|
|
mj_applyFT(m, d, force.data() + 3*v, 0, xpos + 3*v, bid, qfrc);
|
|
} else {
|
|
int body_dofnum = m->body_dofnum[bid];
|
|
int body_dofadr = m->body_dofadr[bid];
|
|
for (int x = 0; x < body_dofnum; x++) {
|
|
qfrc[body_dofadr+x] += force[3*v+x];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// compute metric tensor of edge lengths inner product
|
|
template <typename T>
|
|
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++) {
|
|
for (int i = 0; i < 3; i++) {
|
|
trE[e] += basis[e][4*i];
|
|
}
|
|
}
|
|
|
|
// compute second invariant i.e. trace(strain^2)
|
|
for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
|
|
for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
|
|
for (int i = 0; i < 3; i++) {
|
|
for (int j = 0; j < 3; j++) {
|
|
trEE[T::kNumEdges*ed1+ed2] += basis[ed1][3*i+j] * basis[ed2][3*j+i];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// assembly of strain metric tensor
|
|
for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
|
|
for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
|
|
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
|
|
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);
|
|
|
|
// copied from mjXUtil
|
|
void String2Vector(const std::string& txt, std::vector<int>& vec);
|
|
|
|
// reads numeric attributes
|
|
bool CheckAttr(const char* name, const mjModel* m, int instance);
|
|
|
|
} // namespace mujoco::plugin::elasticity
|
|
|
|
#endif // MUJOCO_PLUGIN_ELASTICITY_ELASTICITY_H_
|