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

Implement multi-cell finite element method for interpolated flexes.

Browse Source 
This change introduces a `flex_cellcount` field to `mjModel` to specify the number of cells in each dimension for interpolated flexes. The stiffness computation, passive force calculation, and Jacobian derivatives are updated to operate on a per-cell basis, significantly improving performance by localizing computations to the nodes within each cell.

PiperOrigin-RevId: 901216393
Change-Id: Ic23132e609de11e71bb7fef8d1f139daad2ec264
This commit is contained in:
Alessio Quaglino
2026-04-17 04:10:54 -07:00
committed by Copybara-Service
parent 8415dff307
commit 6c7ed66781
36 changed files with 1777 additions and 1039 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+170 -137
View File
@@ -432,21 +432,31 @@ static int mj_vertBodyWeight(const mjModel* m, const mjData* d, int f, int* v,
return 0;
}
// compute parametric coordinates of the vertex in [0, 1]^3
mjtNum coord[3] = {0, 0, 0};
for (int i = 0; i < nw; i++) {
mju_addToScl3(coord, m->flex_vert0 + 3*v[i], vweight[i]);
mju_addToScl3(coord, m->flex_vert0 + 3*v[i], vweight[i]);
}
int order = m->flex_interp[f];
int npc = (order+1)*(order+1)*(order+1); // number of nodes per cell
// cell lookup: get local coords and node indices
mjtNum local[3];
int nodeindices[27]; // max npc for quadratic: 3^3 = 27
mju_cellLookup(coord, m->flex_cellnum+3*f, order, local, nodeindices);
// evaluate basis functions for this cell's local nodes
int nstart = m->flex_nodeadr[f];
int nend = m->flex_nodeadr[f] + m->flex_nodenum[f];
int nb = 0;
for (int i = nstart; i < nend; i++) {
mjtNum w = mju_evalBasis(coord, i-nstart, m->flex_interp[f]);
for (int j = 0; j < npc; j++) {
mjtNum w = mju_evalBasis(local, j, order);
if (w < 1e-5) {
continue;
}
if (bweight) bweight[nb] = w;
body[nb++] = m->flex_nodebodyid[i];
body[nb++] = m->flex_nodebodyid[nstart + nodeindices[j]];
}
return nb;
@@ -871,6 +881,11 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
break;
}
int npc = (order+1)*(order+1)*(order+1);
int cx = m->flex_cellnum[3*f+0];
int cy = m->flex_cellnum[3*f+1];
int cz = m->flex_cellnum[3*f+2];
// allocate stack for node positions and Jacobians
mj_markStack(d);
mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
@@ -910,152 +925,164 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
}
}
// loop over Gauss points
// get reference positions from m->flex_node0 (Cartesian positions at qpos0)
// reference positions for all nodes
int nstart = m->flex_nodeadr[f];
mjtNum* refpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
for (int n = 0; n < nodenum; n++) {
mju_copy3(refpos + 3*n, m->flex_node0 + 3*(n + nstart));
}
// B-bar: precompute center-point values for volumetric constraint (trilinear only)
if (order == 1) {
mjtNum center[3] = {0.5, 0.5, 0.5};
mjtNum Fcur_c[9], Fref_c[9], Fref_inv_center[9], F_center[9];
// compute deformation gradient at center
mju_defGradient(Fcur_c, center, xpos, order);
mju_defGradient(Fref_c, center, refpos, order);
mat3_inverse(Fref_c, Fref_inv_center);
mju_mulMatMat3(F_center, Fcur_c, Fref_inv_center);
// compute C and E at center
mjtNum C_c[9], E_c[9];
mju_mulMatTMat3(C_c, F_center, F_center);
mju_scl(E_c, C_c, 0.5, 9);
E_c[0] -= 0.5;
E_c[4] -= 0.5;
E_c[8] -= 0.5;
// J = det(F) at center
mjtNum I1_center = E_c[0] + E_c[4] + E_c[8];
mjtNum J_center = mat3_det(F_center);
// compute shape function gradients at center (8 nodes for trilinear)
mjtNum grad_center[8][3];
shape_gradients(order, center, grad_center);
// add I1 and J-1 constraints at center (reduced integration for volumetric)
mjtNum* dSdx = mjSTACKALLOC(d, 3*nodenum, mjtNum);
for (int inv = 0; inv < 2; inv++) {
if (inv == 0) {
// I1 = tr(E), dI1/dE = I
cpos[0] = I1_center;
} else {
// J - 1 = det(F) - 1, dJ/dF = cofactor(F)
cpos[0] = J_center - 1.0;
}
volumetric_dSdx(inv, nodenum, grad_center, F_center, Fref_inv_center, dSdx);
strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
if (issparse) {
mjtNum* sparse_jac = mjSTACKALLOC(d, combined_nnz, mjtNum);
for (int k = 0; k < combined_nnz; k++) {
sparse_jac[k] = strain_jac[combined_chain[k]];
}
mj_addConstraint(m, d, sparse_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
combined_nnz, combined_chain);
} else {
mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
}
}
}
// add I1 and J-1 constraints at center (reduced integration for volumetric)
// per-cell arrays
mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* refpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* dSdx_local = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* dSdx = mjSTACKALLOC(d, 3*nodenum, mjtNum);
for (int g = 0; g < ngauss; g++) {
mjtNum* p = gauss[g];
int gindices[125]; // max npc = 125 for quadratic
// F = Fcur * Fref_inv
mjtNum Fcur[9], Fref[9], Fref_inv[9], F[9];
mju_defGradient(Fcur, p, xpos, order);
mju_defGradient(Fref, p, refpos, order);
mat3_inverse(Fref, Fref_inv);
mju_mulMatMat3(F, Fcur, Fref_inv);
// loop over cells
for (int ci = 0; ci < cx; ci++) {
for (int cj = 0; cj < cy; cj++) {
for (int ck = 0; ck < cz; ck++) {
// gather cell-local node positions
mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, refpos, xpos_c, NULL,
refpos_c, gindices, NULL);
// compute Green-Lagrange strain E = 0.5*(C - I)
mjtNum C[9], E[9];
mju_mulMatTMat3(C, F, F);
for (int j = 0; j < 9; j++) {
E[j] = 0.5 * C[j];
}
E[0] -= 0.5;
E[4] -= 0.5;
E[8] -= 0.5;
// B-bar: center-point volumetric constraints (trilinear)
if (order == 1) {
mjtNum center[3] = {0.5, 0.5, 0.5};
mjtNum Fcur_c[9], Fref_c[9], Fref_inv_c[9], F_c[9];
// compute 3 invariants of E
mjtNum I1 = E[0] + E[4] + E[8];
mjtNum trE2 = E[0]*E[0] + E[1]*E[3] + E[2]*E[6] +
E[3]*E[1] + E[4]*E[4] + E[5]*E[7] +
E[6]*E[2] + E[7]*E[5] + E[8]*E[8];
mjtNum I2 = 0.5 * (I1*I1 - trE2);
mjtNum I3 = mat3_det(E);
mju_defGradient(Fcur_c, center, xpos_c, order);
mju_defGradient(Fref_c, center, refpos_c, order);
mat3_inverse(Fref_c, Fref_inv_c);
mju_mulMatMat3(F_c, Fcur_c, Fref_inv_c);
// compute shape function gradients at this Gauss point
mjtNum grad[27][3];
shape_gradients(order, p, grad);
mjtNum C_c[9], E_c[9];
mju_mulMatTMat3(C_c, F_c, F_c);
mju_scl(E_c, C_c, 0.5, 9);
E_c[0] -= 0.5; E_c[4] -= 0.5; E_c[8] -= 0.5;
// trilinear: 3 constraints per Gauss point (I1, I2, I3 skipped - only shear)
// quadratic: 6 constraints per Gauss point
for (int s = 0; s < 6; s++) {
// skip I1, I2, I3 for trilinear (I1, J-1 at center; I2 is small for small strain)
if (order == 1 && (s == 0 || s == 1 || s == 2)) {
continue;
}
mjtNum I1_c = E_c[0] + E_c[4] + E_c[8];
mjtNum J_c = mat3_det(F_c);
mjtNum dSdE[9];
mju_zero(dSdE, 9);
mjtNum grad_c[8][3];
shape_gradients(order, center, grad_c);
if (s == 0) {
// I1 = tr(E), dI1/dE = I (only for quadratic)
cpos[0] = I1;
dSdE[0] = dSdE[4] = dSdE[8] = 1.0;
} else if (s == 1) {
// I2 = 0.5*(tr(E)^2 - tr(E^2)), dI2/dE = tr(E)*I - E
cpos[0] = I2;
dSdE[0] = I1 - E[0];
dSdE[4] = I1 - E[4];
dSdE[8] = I1 - E[8];
dSdE[1] = -E[1]; dSdE[3] = -E[3];
dSdE[2] = -E[2]; dSdE[6] = -E[6];
dSdE[5] = -E[5]; dSdE[7] = -E[7];
} else if (s == 2) {
// I3 = det(E), dI3/dE = cofactor(E)
cpos[0] = I3;
mat3_cofactor(E, dSdE);
} else {
// off-diagonal entries: s=3->E12, s=4->E13, s=5->E23
int offdiag_idx[3] = {1, 2, 5};
int ij = offdiag_idx[s - 3];
cpos[0] = E[ij];
dSdE[ij] = 1.0;
}
for (int inv = 0; inv < 2; inv++) {
cpos[0] = (inv == 0) ? I1_c : J_c - 1.0;
// compute dS/dx for all nodes
invariant_dSdx(nodenum, grad, F, Fref_inv, dSdE, dSdx);
strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
// compute local dSdx
volumetric_dSdx(inv, npc, grad_c, F_c, Fref_inv_c, dSdx_local);
// add constraint
if (issparse) {
mjtNum* sparse_jac = mjSTACKALLOC(d, combined_nnz, mjtNum);
for (int k = 0; k < combined_nnz; k++) {
sparse_jac[k] = strain_jac[combined_chain[k]];
// scatter to global dSdx
mju_zero(dSdx, 3*nodenum);
for (int n = 0; n < npc; n++) {
mju_addTo3(dSdx + 3*gindices[n], dSdx_local + 3*n);
}
strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
if (issparse) {
mj_markStack(d);
mjtNum* sj = mjSTACKALLOC(d, combined_nnz, mjtNum);
for (int k = 0; k < combined_nnz; k++) {
sj[k] = strain_jac[combined_chain[k]];
}
mj_addConstraint(m, d, sj, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
combined_nnz, combined_chain);
mj_freeStack(d);
} else {
mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
}
}
}
// Gauss integration per cell
for (int g = 0; g < ngauss; g++) {
mjtNum* p = gauss[g];
// F = Fcur * Fref_inv
mjtNum Fcur[9], Fref[9], Fref_inv[9], F[9];
mju_defGradient(Fcur, p, xpos_c, order);
mju_defGradient(Fref, p, refpos_c, order);
mat3_inverse(Fref, Fref_inv);
mju_mulMatMat3(F, Fcur, Fref_inv);
// Green-Lagrange strain E = 0.5*(C - I)
mjtNum C[9], E[9];
mju_mulMatTMat3(C, F, F);
for (int j = 0; j < 9; j++) {
E[j] = 0.5 * C[j];
}
E[0] -= 0.5; E[4] -= 0.5; E[8] -= 0.5;
// 3 invariants of E
mjtNum I1 = E[0] + E[4] + E[8];
mjtNum trE2 = E[0]*E[0] + E[1]*E[3] + E[2]*E[6]
+ E[3]*E[1] + E[4]*E[4] + E[5]*E[7]
+ E[6]*E[2] + E[7]*E[5] + E[8]*E[8];
mjtNum I2 = 0.5 * (I1*I1 - trE2);
mjtNum I3 = mat3_det(E);
// shape function gradients at Gauss point
mjtNum grad[27][3];
shape_gradients(order, p, grad);
for (int s = 0; s < 6; s++) {
// skip I1,I2,I3 for trilinear (B-bar handles vol)
if (order == 1 && (s == 0 || s == 1 || s == 2)) {
continue;
}
mjtNum dSdE[9];
mju_zero(dSdE, 9);
if (s == 0) {
cpos[0] = I1;
dSdE[0] = dSdE[4] = dSdE[8] = 1.0;
} else if (s == 1) {
cpos[0] = I2;
dSdE[0] = I1-E[0]; dSdE[4] = I1-E[4];
dSdE[8] = I1-E[8];
dSdE[1] = -E[1]; dSdE[3] = -E[3];
dSdE[2] = -E[2]; dSdE[6] = -E[6];
dSdE[5] = -E[5]; dSdE[7] = -E[7];
} else if (s == 2) {
cpos[0] = I3;
mat3_cofactor(E, dSdE);
} else {
int offdiag_idx[3] = {1, 2, 5};
int ij = offdiag_idx[s - 3];
cpos[0] = E[ij];
dSdE[ij] = 1.0;
}
// compute local dS/dx for cell nodes
invariant_dSdx(npc, grad, F, Fref_inv, dSdE,
dSdx_local);
// scatter to global dSdx
mju_zero(dSdx, 3*nodenum);
for (int n = 0; n < npc; n++) {
mju_addTo3(dSdx + 3*gindices[n], dSdx_local + 3*n);
}
strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
if (issparse) {
mj_markStack(d);
mjtNum* sj = mjSTACKALLOC(d, combined_nnz, mjtNum);
for (int k = 0; k < combined_nnz; k++) {
sj[k] = strain_jac[combined_chain[k]];
}
mj_addConstraint(m, d, sj, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
combined_nnz, combined_chain);
mj_freeStack(d);
} else {
mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
}
}
}
mj_addConstraint(m, d, sparse_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
combined_nnz, combined_chain);
} else {
mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
}
}
}
@@ -1899,10 +1926,13 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
int nstart = m->flex_nodeadr[flex_id];
int order = m->flex_interp[flex_id];
// compute constraint count: trilinear (2 + 3*8 = 26), quadratic (6*27 = 162)
// compute constraint count per cell, then multiply by ncells
int nquad = order + 1;
int ngauss = nquad * nquad * nquad;
int nconstraint = (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss);
int ncells = m->flex_cellnum[3*flex_id+0]
* m->flex_cellnum[3*flex_id+1]
* m->flex_cellnum[3*flex_id+2];
int nconstraint = ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss));
mjtNum avg_invweight = 0;
for (int n = 0; n < nodenum; n++) {
@@ -2510,7 +2540,10 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
}
int nquad = order + 1; // 2 for order=1, 3 for order=2
int ngauss = nquad * nquad * nquad; // 8 or 27
size = (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss); // 26 or 162
int ncells = m->flex_cellnum[3*id[0]+0]
* m->flex_cellnum[3*id[0]+1]
* m->flex_cellnum[3*id[0]+2];
size = ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss));
if (nnz) {
// Count unique DOFs across all node bodies (matching instantiation)
src/engine/engine_core_smooth.c
+22 -5
View File
@@ -580,9 +580,11 @@ void mj_flex(const mjModel* m, mjData* d) {
}
}
// trilinear interpolation
// trilinear/quadratic interpolation
else {
mjtNum nodexpos[3*mjMAXFLEXNODES];
int nodenum = nend - nstart;
mj_markStack(d);
mjtNum* nodexpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
if (m->flex_centered[f]) {
for (int i=nstart; i < nend; i++) {
mji_copy3(nodexpos + 3*(i-nstart), d->xpos + 3*m->flex_nodebodyid[i]);
@@ -596,14 +598,26 @@ void mj_flex(const mjModel* m, mjData* d) {
}
int order = m->flex_interp[f];
if (nend - nstart != (order + 1) * (order + 1) * (order + 1)) {
int cx = m->flex_cellnum[3*f+0];
int cy = m->flex_cellnum[3*f+1];
int cz = m->flex_cellnum[3*f+2];
int nx_g = cx * order + 1;
int ny_g = cy * order + 1;
int nz_g = cz * order + 1;
if (nend - nstart != nx_g * ny_g * nz_g) {
mjERROR("flex_interp_order mismatch");
}
for (int i=vstart; i < vend; i++) {
mju_zero3(d->flexvert_xpos+3*i);
mju_interpolate3D(d->flexvert_xpos+3*i, m->flex_vert0 + 3*i, nodexpos, order);
// cell lookup: get local coords and node indices
mjtNum local[3];
int nodeindices[27]; // max npc for quadratic: 3^3 = 27
mju_cellLookup(m->flex_vert0 + 3*i, m->flex_cellnum+3*f, order, local, nodeindices);
mju_interpolate3D(d->flexvert_xpos+3*i, local, nodexpos, order, nodeindices);
}
mj_freeStack(d);
}
}
@@ -2617,7 +2631,10 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
if (order && nodenum) {
int nquad = order + 1;
int ngauss = nquad * nquad * nquad;
i += (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss);
int ncells = m->flex_cellnum[3*k+0]
* m->flex_cellnum[3*k+1]
* m->flex_cellnum[3*k+2];
i += ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss));
}
break;
}
src/engine/engine_core_util.c
+28
View File
@@ -987,6 +987,34 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
//-------------------------- miscellaneous utilities -----------------------------------------------
// gather global node positions and velocities
void mju_flexGatherState(const mjModel* m, mjData* d, int f, mjtNum* xpos, mjtNum* vel) {
int nodenum = m->flex_nodenum[f];
int nstart = m->flex_nodeadr[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
// compute positions
if (m->flex_centered[f]) {
for (int i=0; i < nodenum; i++) {
mju_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
if (vel) {
mju_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]);
}
}
} else {
mjtNum screw[6];
for (int i=0; i < nodenum; i++) {
mju_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart));
mju_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
if (vel) {
mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0);
mju_copy3(vel + 3*i, screw + 3);
}
}
}
}
// extract 6D force:torque for one contact, in contact frame
void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) {
mjContact* con;
src/engine/engine_core_util.h
+3
View File
@@ -129,6 +129,9 @@ MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
//-------------------------- miscellaneous ---------------------------------------------------------
// gather global node positions and velocities
MJAPI void mju_flexGatherState(const mjModel* m, mjData* d, int f, mjtNum* xpos, mjtNum* vel);
// extract 6D force:torque for one contact, in contact frame
MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
src/engine/engine_derivative.c
+113 -94
View File
@@ -915,127 +915,146 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
continue;
}
int order = m->flex_interp[f];
int npc = (order+1)*(order+1)*(order+1);
int cx = m->flex_cellnum[3*f+0];
int cy = m->flex_cellnum[3*f+1];
int cz = m->flex_cellnum[3*f+2];
int nodenum = m->flex_nodenum[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
// standard stack allocation
mj_markStack(d);
mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* K_rot = mjSTACKALLOC(d, 9*nodenum*nodenum, mjtNum);
// sparse Jacobian allocations
int dim = 3 * nodenum;
int* rownnz = mjSTACKALLOC(d, dim, int);
int* rowadr = mjSTACKALLOC(d, dim, int);
mjtNum* J_val = mjSTACKALLOC(d, dim*nv, mjtNum);
int* J_colind = mjSTACKALLOC(d, dim*nv, int);
// per-cell arrays
int dim_c = 3 * npc;
mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* K_rot_cell = mjSTACKALLOC(d, dim_c*dim_c, mjtNum);
// sparse Jacobian for one cell
int* J_rownnz = mjSTACKALLOC(d, dim_c, int);
int* J_rowadr = mjSTACKALLOC(d, dim_c, int);
mjtNum* J_val = mjSTACKALLOC(d, dim_c*nv, mjtNum);
int* J_colind = mjSTACKALLOC(d, dim_c*nv, int);
// temp allocations for chain
int* chain_colind = mjSTACKALLOC(d, nv, int);
mjtNum* blk_jac = mjSTACKALLOC(d, 3*nv, mjtNum);
// compute positions, rotation and Jacobian
mjtNum quat[4] = {1, 0, 0, 0};
mj_flexInterpState(m, d, f, xpos, NULL, quat);
// gather raw node positions (unrotated)
mju_flexGatherState(m, d, f, xpos, NULL);
// compute generalized stiffness in global frame: K_rot = R * K * R^T
mjtNum R[9];
mju_quat2Mat(R, quat); // R = R_global2local
mjtNum RT[9];
mju_transpose(RT, R, 3, 3); // RT = R_local2global
// loop over cells
int cell_idx = 0;
for (int ci = 0; ci < cx; ci++) {
for (int cj = 0; cj < cy; cj++) {
for (int ck = 0; ck < cz; ck++) {
// gather cell-local node positions
int gindices[125]; // max npc = 125 for quadratic
mjtNum quat[4];
mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, NULL,
xpos_c, NULL, NULL, gindices, quat);
// blockwise rotation: K_rot(i,j) = scale * RT * K_local(i,j) * R
// note: k stores -K, so K_rot = scale * (-K_phys)
for (int i=0; i < nodenum; i++) {
for (int j=0; j < nodenum; j++) {
mjtNum blk[9], tmp[9];
// R = R_global2local, RT = R_local2global
mjtNum R[9], RT[9];
mju_quat2Mat(R, quat);
mju_transpose(RT, R, 3, 3);
// get K_local(i,j)
int adr = (3*i)*(3*nodenum) + 3*j;
for (int r=0; r < 3; r++) {
for (int c=0; c < 3; c++) {
blk[3*r+c] = k[adr + r*(3*nodenum) + c];
// get cell stiffness
mjtNum* k_cell = k + cell_idx * 3*npc * 3*npc;
// compute K_rot_cell = RT * K_cell * R (block-wise)
mju_zero(K_rot_cell, dim_c*dim_c);
for (int a = 0; a < npc; a++) {
for (int b = 0; b < npc; b++) {
mjtNum blk[9], tmp[9];
// get K_cell(a,b) 3x3 block
int adr_cell = (3*a)*(3*npc) + 3*b;
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
blk[3*r+c] = k_cell[adr_cell + r*(3*npc) + c];
}
}
// tmp = K * R
mju_mulMatMat3(tmp, blk, R);
// blk = RT * tmp = RT * K * R
mju_mulMatMat3(blk, RT, tmp);
// store in K_rot_cell at (a, b)
int adr_out = (3*a)*dim_c + 3*b;
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
K_rot_cell[adr_out + r*dim_c + c] = scale * blk[3*r+c];
}
}
}
}
}
// tmp = K * R
mju_mulMatMat3(tmp, blk, R);
// construct sparse Jacobian for this cell's nodes
int current_adr = 0;
for (int n = 0; n < npc; n++) {
int bid = bodyid[gindices[n]];
int chain_nnz = mj_bodyChain(m, bid, chain_colind);
mj_jacSparse(m, d, blk_jac, NULL, xpos+3*gindices[n], bid,
chain_nnz, chain_colind, /*flg_skipcommon=*/0);
// blk = RT * tmp = RT * K * R
mju_mulMatMat3(blk, RT, tmp);
for (int r = 0; r < 3; r++) {
int row_idx = 3*n + r;
J_rownnz[row_idx] = chain_nnz;
J_rowadr[row_idx] = current_adr;
// store scaled into K_rot
for (int r=0; r < 3; r++) {
for (int c=0; c < 3; c++) {
K_rot[adr + r*(3*nodenum) + c] = scale * blk[3*r+c];
}
}
}
}
// construct sparse Jacobian J_val
int current_adr = 0;
for (int i=0; i < nodenum; i++) {
// get chain for this node
int chain_nnz = mj_bodyChain(m, bodyid[i], chain_colind);
// compute sparse Jacobian for this node (3 rows)
mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind,
/*flg_skipcommon=*/0);
// copy to sparse structure
for (int r=0; r<3; r++) {
int row_idx = 3*i + r;
rownnz[row_idx] = chain_nnz;
rowadr[row_idx] = current_adr;
for (int idx=0; idx<chain_nnz; idx++) {
for (int idx = 0; idx < chain_nnz; idx++) {
J_colind[current_adr] = chain_colind[idx];
J_val[current_adr] = blk_jac[r*chain_nnz + idx];
current_adr++;
}
}
}
// perform operation
if (op == mjFLEXOP_VEC) {
// res += J^T * K_rot * J * vec
addJTBJ_mulSparse(m, d, res, vec, rownnz, rowadr, J_colind, J_val, K_rot, dim);
} else if (op == mjFLEXOP_ADDH) {
// H += -J^T * K_rot * J
// H is dense ndof x ndof
// reuse stack for J_reduced (but now we extract from sparse J)
mjtNum* J_reduced = mjSTACKALLOC(d, dim*ndof, mjtNum);
mju_zero(J_reduced, dim*ndof);
// extract columns of J into J_reduced
for (int i=0; i<dim; i++) {
int nnz = rownnz[i];
int adr = rowadr[i];
for (int idx=0; idx<nnz; idx++) {
int global_col = J_colind[adr + idx];
int local_idx = global2local[global_col];
if (local_idx >= 0) {
J_reduced[i*ndof + local_idx] = J_val[adr + idx];
}
}
}
}
// H -= J_reduced^T * K_rot * J_reduced
// K_rot * J_reduced (dim x ndof)
mjtNum* KJ = mjSTACKALLOC(d, dim*ndof, mjtNum);
mju_mulMatMat(KJ, K_rot, J_reduced, dim, dim, ndof);
// apply operation with cell's K_rot and J
if (op == mjFLEXOP_VEC) {
addJTBJ_mulSparse(m, d, res, vec, J_rownnz, J_rowadr, J_colind,
J_val, K_rot_cell, dim_c);
} else if (op == mjFLEXOP_ADDH) {
mj_markStack(d);
// H -= J_cell^T * K_rot_cell * J_cell
mjtNum* J_reduced = mjSTACKALLOC(d, dim_c*ndof, mjtNum);
mju_zero(J_reduced, dim_c*ndof);
// H[i, j] -= sum_k J_reduced[k, i] * KJ[k, j]
for (int i=0; i<ndof; i++) {
for (int j=0; j<ndof; j++) {
mjtNum val = 0;
for (int dim_idx=0; dim_idx<dim; dim_idx++) {
val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j];
for (int i = 0; i < dim_c; i++) {
int nnz = J_rownnz[i];
int adr = J_rowadr[i];
for (int idx = 0; idx < nnz; idx++) {
int global_col = J_colind[adr + idx];
int local_idx = global2local[global_col];
if (local_idx >= 0) {
J_reduced[i*ndof + local_idx] = J_val[adr + idx];
}
}
}
// KJ = K_rot_cell * J_reduced (dim_c x ndof)
mjtNum* KJ = mjSTACKALLOC(d, dim_c*ndof, mjtNum);
mju_mulMatMat(KJ, K_rot_cell, J_reduced, dim_c, dim_c, ndof);
// H[i,j] -= J_reduced[k,i] * KJ[k,j]
for (int i = 0; i < ndof; i++) {
for (int j = 0; j < ndof; j++) {
mjtNum val = 0;
for (int dim_idx = 0; dim_idx < dim_c; dim_idx++) {
val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j];
}
res[i*ndof + j] -= val;
}
}
mj_freeStack(d);
}
// res is H
res[i*ndof + j] -= val;
cell_idx++;
}
}
}
src/engine/engine_passive.c
+95 -80
View File
@@ -58,61 +58,7 @@ static void inline GradSquaredLengths(mjtNum gradient[6][2][3],
}
}
// compute interpolated flex state: xpos, vel, quat
// f: flex index
// xpos: (output) 3*nodenum
// vel: (output) 3*nodenum, can be NULL
// quat: (output) 4, rotation from global to local
void mj_flexInterpState(const mjModel* m, mjData* d, int f,
mjtNum* xpos, mjtNum* vel, mjtNum* quat) {
int nodenum = m->flex_nodenum[f];
int nstart = m->flex_nodeadr[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
mjtNum com[3] = {0};
// compute positions
if (m->flex_centered[f]) {
for (int i=0; i < nodenum; i++) {
mji_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
if (vel) {
mji_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]);
}
}
} else {
mjtNum screw[6];
for (int i=0; i < nodenum; i++) {
mji_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart));
mji_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
if (vel) {
mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0);
mji_copy3(vel + 3*i, screw + 3);
}
}
}
// compute center of mass
for (int i = 0; i < nodenum; i++) {
mji_addToScl3(com, xpos+3*i, 1.0/nodenum);
}
// compute the Jacobian at the center of mass
mjtNum mat[9] = {0};
mjtNum p[3] = {.5, .5, .5};
mju_defGradient(mat, p, xpos, m->flex_interp[f]);
// find rotation
mju_mat2Rot(quat, mat);
mju_negQuat(quat, quat);
// rotate vertices to quat and add reference center of mass
for (int i = 0; i < nodenum; i++) {
mju_rotVecQuat(xpos+3*i, xpos+3*i, quat);
mji_addTo3(xpos+3*i, p);
if (vel) {
mju_rotVecQuat(vel+3*i, vel+3*i, quat);
}
}
}
// spring and damper forces
static void mj_springdamper(const mjModel* m, mjData* d) {
@@ -284,42 +230,111 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
}
if (m->flex_interp[f]) {
int order = m->flex_interp[f];
int npc = (order+1)*(order+1)*(order+1); // nodes per cell
int cx = m->flex_cellnum[3*f+0];
int cy = m->flex_cellnum[3*f+1];
int cz = m->flex_cellnum[3*f+2];
int ny_g = cy * order + 1;
int nz_g = cz * order + 1;
mj_markStack(d);
mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* displ = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* vel = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* frc = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* dmp = mjSTACKALLOC(d, 3*nodenum, mjtNum);
// allocate global arrays
mjtNum* xpos_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* vel_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* frc_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* dmp_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* xpos0 = m->flex_node0 + 3*m->flex_nodeadr[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
mjtNum quat[4] = {1, 0, 0, 0};
mj_flexInterpState(m, d, f, xpos, vel, quat);
// gather global node positions and velocities (unrotated)
mju_flexGatherState(m, d, f, xpos_g, vel_g);
// compute displacement
for (int i = 0; i < nodenum; i++) {
mji_addScl3(displ+3*i, xpos+3*i, xpos0+3*i, -1);
// zero global force accumulators
mju_zero(frc_g, 3*nodenum);
mju_zero(dmp_g, 3*nodenum);
// per-cell arrays
mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* vel_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* xpos0_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* displ_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* frc_c = mjSTACKALLOC(d, 3*npc, mjtNum);
mjtNum* dmp_c = mjSTACKALLOC(d, 3*npc, mjtNum);
// loop over cells
int cell_idx = 0;
for (int ci = 0; ci < cx; ci++) {
for (int cj = 0; cj < cy; cj++) {
for (int ck = 0; ck < cz; ck++) {
// gather cell-local node data
mjtNum quat[4];
mjtNum p[3] = {.5, .5, .5};
mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos_g, vel_g, xpos0,
xpos_c, vel_c, xpos0_c, NULL, quat);
// rotate to corotational frame
for (int n = 0; n < npc; n++) {
mju_rotVecQuat(xpos_c+3*n, xpos_c+3*n, quat);
mji_addTo3(xpos_c+3*n, p);
mju_rotVecQuat(vel_c+3*n, vel_c+3*n, quat);
}
// compute displacement
for (int n = 0; n < npc; n++) {
mji_addScl3(displ_c+3*n, xpos_c+3*n, xpos0_c+3*n, -1);
}
// get cell stiffness matrix
mjtNum* k_cell = k + cell_idx * 3*npc * 3*npc;
// compute force in corotational frame
if (enbl_spring) {
mju_mulMatVec(frc_c, k_cell, displ_c, 3*npc, 3*npc);
}
if (enbl_damper) {
mju_mulMatVec(dmp_c, k_cell, vel_c, 3*npc, 3*npc);
}
// rotate back to global frame and scatter
mju_negQuat(quat, quat);
int local = 0;
for (int li = 0; li <= order; li++) {
for (int lj = 0; lj <= order; lj++) {
for (int lk = 0; lk <= order; lk++) {
int gi = ci*order + li;
int gj = cj*order + lj;
int gk = ck*order + lk;
int gidx = gi*ny_g*nz_g + gj*nz_g + gk;
mjtNum qfrc[3], qdmp[3];
mji_rotVecQuat(qfrc, frc_c+3*local, quat);
mji_rotVecQuat(qdmp, dmp_c+3*local, quat);
if (enbl_spring) {
mji_addTo3(frc_g + 3*gidx, qfrc);
}
if (enbl_damper) {
mji_addTo3(dmp_g + 3*gidx, qdmp);
}
local++;
}
}
}
cell_idx++;
}
}
}
// compute force in the stretch frame
if (enbl_spring) mju_mulMatVec(frc, k, displ, 3*nodenum, 3*nodenum);
// compute damping force in stretch frame
if (enbl_damper) mju_mulMatVec(dmp, k, vel, 3*nodenum, 3*nodenum);
// rotate forces to global frame and add to qfrc
mju_negQuat(quat, quat);
// apply accumulated forces to bodies
for (int i = 0; i < nodenum; i++) {
mjtNum qfrc[3], qdmp[3];
mji_rotVecQuat(qfrc, frc+3*i, quat);
mji_rotVecQuat(qdmp, dmp+3*i, quat);
mju_scl3(qdmp, qdmp, m->flex_damping[f]);
mju_scl3(dmp_g+3*i, dmp_g+3*i, m->flex_damping[f]);
if (m->flex_centered[f]) {
if (enbl_spring) mji_addTo3(d->qfrc_spring+m->body_dofadr[bodyid[i]], qfrc);
if (enbl_damper) mji_addTo3(d->qfrc_damper+m->body_dofadr[bodyid[i]], qdmp);
if (enbl_spring) mji_addTo3(d->qfrc_spring + m->body_dofadr[bodyid[i]], frc_g+3*i);
if (enbl_damper) mji_addTo3(d->qfrc_damper + m->body_dofadr[bodyid[i]], dmp_g+3*i);
} else {
if (enbl_spring) mj_applyFT(m, d, qfrc, 0, xpos+3*i, bodyid[i], d->qfrc_spring);
if (enbl_damper) mj_applyFT(m, d, qdmp, 0, xpos+3*i, bodyid[i], d->qfrc_damper);
if (enbl_spring) mj_applyFT(m, d, frc_g+3*i, 0, xpos_g+3*i, bodyid[i], d->qfrc_spring);
if (enbl_damper) mj_applyFT(m, d, dmp_g+3*i, 0, xpos_g+3*i, bodyid[i], d->qfrc_damper);
}
}
src/engine/engine_passive.h
+1 -3
View File
@@ -28,9 +28,7 @@ extern "C" {
// all passive forces
MJAPI void mj_passive(const mjModel* m, mjData* d);
// compute interpolated flex state: xpos, vel, quat
MJAPI void mj_flexInterpState(const mjModel* m, mjData* d, int f,
mjtNum* xpos, mjtNum* vel, mjtNum* quat);
//------------------------- fluid models -----------------------------------------------------------
src/engine/engine_util_misc.c
+107 -2
View File
@@ -612,11 +612,116 @@ mjtNum mju_evalBasis(const mjtNum x[3], int i, int order) {
}
}
// map global parametric coord to cell-local coord and build node indices
// coord: [0,1]^3 parametric coordinates
// cellnum: cell counts (cx, cy, cz)
// order: interpolation order (1=trilinear, 2=triquadratic)
// local: output local parametric coordinates within cell [0,1]^3
// nodeindices: output array of global node indices for the cell (size (order+1)^3, may be NULL)
// returns: number of nodes per cell (order+1)^3
int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNum local[3],
int* nodeindices) {
int cx = cellnum[0], cy = cellnum[1], cz = cellnum[2];
// find containing cell
int ci = (int)mju_floor(coord[0] * cx);
int cj = (int)mju_floor(coord[1] * cy);
int ck = (int)mju_floor(coord[2] * cz);
ci = mjMIN(ci, cx - 1); ci = mjMAX(ci, 0);
cj = mjMIN(cj, cy - 1); cj = mjMAX(cj, 0);
ck = mjMIN(ck, cz - 1); ck = mjMAX(ck, 0);
// local parametric coordinates within cell
local[0] = mju_clip(coord[0] * cx - ci, 0, 1);
local[1] = mju_clip(coord[1] * cy - cj, 0, 1);
local[2] = mju_clip(coord[2] * cz - ck, 0, 1);
// build node indices for this cell
if (nodeindices) {
int ny_g = cy * order + 1;
int nz_g = cz * order + 1;
int ni = 0;
for (int li = 0; li <= order; li++) {
for (int lj = 0; lj <= order; lj++) {
for (int lk = 0; lk <= order; lk++) {
int gi = ci*order + li;
int gj = cj*order + lj;
int gk = ck*order + lk;
nodeindices[ni++] = gi*ny_g*nz_g + gj*nz_g + gk;
}
}
}
}
int npc = (order + 1) * (order + 1) * (order + 1);
return npc;
}
// interpolate a function at x with given interpolation coefficients and order n
void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order) {
void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order,
const int* nodeindices) {
int npoint = (order + 1) * (order + 1) * (order + 1);
for (int j=0; j < npoint; j++) {
mju_addToScl3(res, coeff+3*j, mju_evalBasis(x, j, order));
int idx = nodeindices ? nodeindices[j] : j;
mju_addToScl3(res, coeff+3*idx, mju_evalBasis(x, j, order));
}
}
static void flexInterpRotation(int order, const mjtNum* xpos_c,
const mjtNum local[3], mjtNum* quat) {
mjtNum mat[9] = {0};
if (order > 0) {
mju_defGradient(mat, local, xpos_c, order);
} else {
// order 0: fallback to identity matrix
mat[0] = 1;
mat[4] = 1;
mat[8] = 1;
}
// find rotation
quat[0] = 1;
quat[1] = 0;
quat[2] = 0;
quat[3] = 0;
mju_mat2Rot(quat, mat);
mju_negQuat(quat, quat);
}
// gather cell-local quantities and optionally compute rotation
void mju_flexGatherCellState(int order, int cy, int cz, int ci, int cj, int ck,
const mjtNum* xpos_g, const mjtNum* vel_g, const mjtNum* xpos0_g,
mjtNum* xpos_c, mjtNum* vel_c, mjtNum* xpos0_c,
int* nodeindices, mjtNum* quat) {
int ny_g = cy * order + 1;
int nz_g = cz * order + 1;
int local = 0;
for (int li = 0; li <= order; li++) {
for (int lj = 0; lj <= order; lj++) {
for (int lk = 0; lk <= order; lk++) {
int gi = ci*order + li;
int gj = cj*order + lj;
int gk = ck*order + lk;
int gidx = gi*ny_g*nz_g + gj*nz_g + gk;
if (xpos_c && xpos_g) mju_copy3(xpos_c + 3*local, xpos_g + 3*gidx);
if (vel_c && vel_g) mju_copy3(vel_c + 3*local, vel_g + 3*gidx);
if (xpos0_c && xpos0_g) mju_copy3(xpos0_c + 3*local, xpos0_g + 3*gidx);
if (nodeindices) nodeindices[local] = gidx;
local++;
}
}
}
if (quat && xpos_c) {
mjtNum p[3] = {.5, .5, .5};
flexInterpRotation(order, xpos_c, p, quat);
}
}
src/engine/engine_util_misc.h
+13 -1
View File
@@ -89,8 +89,20 @@ MJAPI void mju_defGradient(mjtNum res[9], const mjtNum p[3], const mjtNum* dof,
// evaluate the basis function at x for the i-th node
MJAPI mjtNum mju_evalBasis(const mjtNum x[3], int i, int order);
// map global parametric coord to cell-local coord and build node indices
MJAPI int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNum local[3],
int* nodeindices);
// interpolate a function at x with given interpolation coefficients and order n
MJAPI void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order);
MJAPI void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order,
const int* nodeindices);
// gather cell-local quantities and optionally compute rotation
MJAPI void mju_flexGatherCellState(int order, int cy, int cz, int ci, int cj, int ck,
const mjtNum* xpos_g, const mjtNum* vel_g,
const mjtNum* xpos0_g, mjtNum* xpos_c, mjtNum* vel_c,
mjtNum* xpos0_c, int* nodeindices, mjtNum* quat);
// ----------------------------- Base64 ------------------------------------------------------------
src/engine/engine_vis_interact.c
+16 -10
View File
@@ -863,24 +863,30 @@ int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
flexdist = newdist;
if (m->flex_interp[i]) {
mjtNum* coord = m->flex_vert0 + 3*(m->flex_vertadr[i] + vertid);
int order = m->flex_interp[i];
int npc = (order+1)*(order+1)*(order+1);
// cell lookup: get local coords and node indices
mjtNum loc[3];
int nodeindices[27]; // max npc for quadratic: 3^3 = 27
mju_cellLookup(coord, m->flex_cellnum+3*i, order, loc, nodeindices);
// find node with largest weight in this cell
int nodeid = -1;
int nstart = m->flex_nodeadr[i];
int nend = nstart + m->flex_nodenum[i];
mjtNum w = 0;
for (int j = nstart; j < nend; j++) {
if (mju_evalBasis(coord, j-nstart, m->flex_interp[i]) > w) {
w = mju_evalBasis(coord, j-nstart, m->flex_interp[i]);
nodeid = j;
for (int j = 0; j < npc; j++) {
mjtNum ww = mju_evalBasis(loc, j, order);
if (ww > w) {
w = ww;
nodeid = nodeindices[j];
}
}
if (nodeid < 0) {
mjERROR("flex %d: node closest to vertex %d not found", i, vertid);
}
flexbodyid = m->flex_nodebodyid[m->flex_nodeadr[i] + nodeid];
flexbodyid = m->flex_nodebodyid[nstart + nodeid];
if (m->flex_centered[i]) {
mju_copy3(flexpnt, d->xpos + 3*flexbodyid);
} else {
mju_mulMatVec3(flexpnt, d->xmat + 9*flexbodyid, m->flex_node + 3*nodeid);
mju_mulMatVec3(flexpnt, d->xmat + 9*flexbodyid, m->flex_node + 3*(nstart + nodeid));
mju_addTo3(flexpnt, d->xpos + 3*flexbodyid);
}
} else {
src/engine/engine_vis_visualize.c
+20 -13
View File
@@ -1434,10 +1434,9 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
}
// control points box
mjtNum xpos[mjMAXFLEXNODES];
mjtNum* xpos = mjSTACKALLOC(d, 3*m->flex_nodenum[f], mjtNum);
int nstart = m->flex_nodeadr[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
int nnode = m->flex_interp[f]+1;
if (m->flex_centered[f]) {
for (int i=0; i < m->flex_nodenum[f]; i++) {
mju_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
@@ -1448,15 +1447,23 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
mju_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
}
}
for (int i=0; i < nnode; i++) {
for (int j=0; j < nnode; j++) {
for (int k=0; k < nnode; k++) {
int nn = nnode*nnode;
int offset = 3*(nn*(i+0) + nnode*(j+0) + k);
int offset1 = 3*(nn*(i+1) + nnode*(j+0) + k);
int offset2 = 3*(nn*(i+0) + nnode*(j+1) + k);
int offset3 = 3*(nn*(i+0) + nnode*(j+0) + (k+1));
if (i < nnode-1) {
int cx = m->flex_cellnum[3*f+0];
int cy = m->flex_cellnum[3*f+1];
int cz = m->flex_cellnum[3*f+2];
int order = m->flex_interp[f];
int NX = cx * order + 1;
int NY = cy * order + 1;
int NZ = cz * order + 1;
for (int i=0; i < NX; i++) {
for (int j=0; j < NY; j++) {
for (int k=0; k < NZ; k++) {
int offset = 3*(i*NY*NZ + j*NZ + k);
int offset1 = 3*((i+1)*NY*NZ + j*NZ + k);
int offset2 = 3*(i*NY*NZ + (j+1)*NZ + k);
int offset3 = 3*(i*NY*NZ + j*NZ + (k+1));
if (i < NX-1) {
mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN);
if (!thisgeom) {
return;
@@ -1465,7 +1472,7 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
mjv_connector(thisgeom, mjGEOM_LINE, 3, xpos+offset, xpos+offset1);
releaseGeom(&thisgeom, scn);
}
if (j < nnode-1) {
if (j < NY-1) {
mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN);
if (!thisgeom) {
return;
@@ -1474,7 +1481,7 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
mjv_connector(thisgeom, mjGEOM_LINE, 3, xpos+offset, xpos+offset2);
releaseGeom(&thisgeom, scn);
}
if (k < nnode-1) {
if (k < NZ-1) {
mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN);
if (!thisgeom) {
return;
src/user/user_flexcomp.cc
+77 -22
View File
@@ -77,6 +77,7 @@ bool IsValidElementOrNodeHeader22(const std::string& line) {
mjCFlexcomp::mjCFlexcomp(void) {
type = mjFCOMPTYPE_GRID;
count[0] = count[1] = count[2] = 10;
cellcount[0] = cellcount[1] = cellcount[2] = -1;
mjuu_setvec(spacing, 0.02, 0.02, 0.02);
mjuu_setvec(scale, 1, 1, 1);
mass = 1;
@@ -269,10 +270,19 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
// construct pinned array
int nnode = 0;
if (doftype == mjFCOMPDOF_TRILINEAR) {
nnode = 8;
} else if (doftype == mjFCOMPDOF_QUADRATIC) {
nnode = 27;
if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) {
int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
// multi-cell count for mesh/direct/gmsh, else single cell
int cx = 1, cy = 1, cz = 1;
if (type == mjFCOMPTYPE_MESH || type == mjFCOMPTYPE_DIRECT ||
type == mjFCOMPTYPE_GMSH) {
if (cellcount[0] >= 0) {
cx = cellcount[0];
cy = cellcount[1];
cz = cellcount[2];
}
}
nnode = (cx*order+1) * (cy*order+1) * (cz*order+1);
}
pinned = vector<bool>(std::max(npnt, nnode), rigid);
@@ -562,36 +572,81 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
}
}
// create nodal mesh for trilinear interpolation
// create nodal mesh for trilinear/quadratic interpolation
if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) {
int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
flex->SetOrder(order);
std::vector<double> node(3*(order+1)*(order+1)*(order+1), 0);
flex->spec.order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
if (cellcount[0] >= 0) {
flex->spec.cellcount[0] = cellcount[0];
flex->spec.cellcount[1] = cellcount[1];
flex->spec.cellcount[2] = cellcount[2];
}
// total number of nodes with shared boundaries
int nx = flex->spec.cellcount[0] * flex->spec.order + 1;
int ny = flex->spec.cellcount[1] * flex->spec.order + 1;
int nz = flex->spec.cellcount[2] * flex->spec.order + 1;
int nnode = nx * ny * nz;
std::vector<double> node(3 * nnode, 0);
int idx = 0;
double step = 1.0 / (double)order;
// Simpson's rule weights for quadratic mass distribution
double massP2[3] = {1. / 6., 2. / 3., 1. / 6.};
for (int i=0; i <= order; i++) {
for (int j=0; j <= order; j++) {
for (int k=0; k <= order; k++) {
// compute per-node mass for trilinear:
// mass / nnode (uniform), or use Simpson for quadratic
double node_mass_uniform = mass / nnode;
for (int gi = 0; gi < nx; gi++) {
for (int gj = 0; gj < ny; gj++) {
for (int gk = 0; gk < nz; gk++) {
// parametric position in [0, 1]^3
double s = (double)gi / (flex->spec.cellcount[0] * flex->spec.order);
double t = (double)gj / (flex->spec.cellcount[1] * flex->spec.order);
double u = (double)gk / (flex->spec.cellcount[2] * flex->spec.order);
// physical position
double px = minmax[0] + s * (minmax[3] - minmax[0]);
double py = minmax[1] + t * (minmax[4] - minmax[1]);
double pz = minmax[2] + u * (minmax[5] - minmax[2]);
if (pinned[idx]) {
node[3*idx+0] = minmax[0] + i * step * (minmax[3] - minmax[0]);
node[3*idx+1] = minmax[1] + j * step * (minmax[4] - minmax[1]);
node[3*idx+2] = minmax[2] + k * step * (minmax[5] - minmax[2]);
mjs_appendString(pf->nodebody, mjs_getName(body->element)->c_str());
node[3*idx+0] = px;
node[3*idx+1] = py;
node[3*idx+2] = pz;
mjs_appendString(pf->nodebody,
mjs_getName(body->element)->c_str());
idx++;
continue;
}
mjsBody* pb = mjs_addBody(body, 0);
pb->pos[0] = minmax[0] + i * step * (minmax[3] - minmax[0]);
pb->pos[1] = minmax[1] + j * step * (minmax[4] - minmax[1]);
pb->pos[2] = minmax[2] + k * step * (minmax[5] - minmax[2]);
pb->pos[0] = px;
pb->pos[1] = py;
pb->pos[2] = pz;
mjuu_zerovec(pb->ipos, 3);
// mass distribution
if (doftype == mjFCOMPDOF_TRILINEAR) {
pb->mass = mass / 8;
pb->mass = node_mass_uniform;
} else {
pb->mass = mass * massP2[i] * massP2[j] * massP2[k];
// local index within the cell for mass computation
int li = gi % flex->spec.order;
int lj = gj % flex->spec.order;
int lk = gk % flex->spec.order;
// boundary nodes: average mass contribution
int ncells_i = (gi > 0 && gi < nx-1 && li == 0) ? 2 : 1;
int ncells_j = (gj > 0 && gj < ny-1 && lj == 0) ? 2 : 1;
int ncells_k = (gk > 0 && gk < nz-1 && lk == 0) ? 2 : 1;
// use Simpson weights scaled by cell count
double wi = massP2[li == 0 ? 0 : li];
double wj = massP2[lj == 0 ? 0 : lj];
double wk = massP2[lk == 0 ? 0 : lk];
pb->mass = mass * wi * wj * wk * ncells_i * ncells_j * ncells_k
/ (flex->spec.cellcount[0] * flex->spec.cellcount[1] * flex->spec.cellcount[2]);
}
pb->inertia[0] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
pb->inertia[1] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
pb->inertia[2] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
@@ -607,7 +662,7 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
// construct node name, add to nodebody
char txt[100];
mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), i, j, k);
mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), gi, gj, gk);
mjs_setName(pb->element, txt);
mjs_appendString(pf->nodebody, mjs_getName(pb->element)->c_str());
src/user/user_flexcomp.h
+1
View File
@@ -78,6 +78,7 @@ class mjCFlexcomp {
std::string name; // flex name
mjtFcompType type; // flexcomp type
int count[3]; // grid count in each dimension
int cellcount[3]; // number of cells for interpolation
double spacing[3]; // spacing between grid elements
double scale[3]; // scaling for mesh and direct
double origin[3]; // origin for generating a 3D mesh from a convex 2D mesh
src/user/user_init.c
+3
View File
@@ -224,6 +224,9 @@ void mjs_defaultFlex(mjsFlex* flex) {
// set other defaults
flex->dim = 2;
flex->radius = 0.005;
flex->cellcount[0] = 1;
flex->cellcount[1] = 1;
flex->cellcount[2] = 1;
flex->internal = 0;
flex->selfcollide = mjFLEXSELF_AUTO;
flex->activelayers = 1;
src/user/user_mesh.cc
+59 -10
View File
@@ -3961,7 +3961,10 @@ std::string mjCFlex::ComputeStiffnessCacheKey() const {
combine(std::hash<double>{}(young));
combine(std::hash<double>{}(poisson));
combine(std::hash<int>{}(order_));
combine(std::hash<int>{}(spec.order));
combine(std::hash<int>{}(spec.cellcount[0]));
combine(std::hash<int>{}(spec.cellcount[1]));
combine(std::hash<int>{}(spec.cellcount[2]));
// compute bounding box from vertex positions
if (!vert_.empty()) {
@@ -4086,10 +4089,20 @@ void mjCFlex::Compile(const mjVFS* vfs) {
// set nnode
nnode = static_cast<int>(nodebody_.size());
if (nnode && !order_) {
order_ = std::pow(nnode, 1.0 / 3) - 1;
if (nnode != std::pow(order_ + 1, 3)) {
throw mjCError(this, "number of nodes must be %d^3 but it is %d", nullptr, order_, nnode);
if (nnode && !spec.order) {
throw mjCError(this, "Interpolation order must be explicitly specified (dof is missing)");
}
// check node compatibility with count and dof
if (spec.order > 0) {
int expected_nodes = (spec.cellcount[0] * spec.order + 1) *
(spec.cellcount[1] * spec.order + 1) *
(spec.cellcount[2] * spec.order + 1);
if (nnode != expected_nodes) {
std::string msg = "number of nodes (" + std::to_string(nnode) +
") does not match cellcount and dof expected (" +
std::to_string(expected_nodes) + ")";
throw mjCError(this, msg.c_str());
}
}
@@ -4329,12 +4342,48 @@ void mjCFlex::Compile(const mjVFS* vfs) {
}
if (!stiffness_cached && young > 0 && interpolated) {
int n = pow(order_ + 1, 3);
int ndof = 3 * n;
if (stiffness.size() < ndof * ndof) {
stiffness.resize(ndof * ndof, 0);
int npc = pow(spec.order + 1, 3); // nodes per cell
int ndof_cell = 3 * npc;
int cx = spec.cellcount[0], cy = spec.cellcount[1], cz = spec.cellcount[2];
int ncells = cx * cy * cz;
int ny_global = cy * spec.order + 1;
int nz_global = cz * spec.order + 1;
// total stiffness = ncells * ndof_cell^2
stiffness.resize(ncells * ndof_cell * ndof_cell, 0);
// compute stiffness per cell
for (int ci = 0; ci < cx; ci++) {
for (int cj = 0; cj < cy; cj++) {
for (int ck = 0; ck < cz; ck++) {
int cell_idx = ci * cy * cz + cj * cz + ck;
// gather cell's local node positions
std::vector<double> cell_pos(3 * npc);
int local = 0;
for (int li = 0; li <= spec.order; li++) {
for (int lj = 0; lj <= spec.order; lj++) {
for (int lk = 0; lk <= spec.order; lk++) {
int gi = ci * spec.order + li;
int gj = cj * spec.order + lj;
int gk = ck * spec.order + lk;
int global = gi * ny_global * nz_global + gj * nz_global + gk;
mjuu_copyvec(cell_pos.data() + 3*local, nodexpos.data() + 3*global, 3);
local++;
}
}
}
// compute per-cell stiffness
std::vector<double> K_cell(ndof_cell * ndof_cell, 0);
ComputeLinearStiffness(K_cell, cell_pos.data(), young, poisson, spec.order);
// copy into global stiffness array
mjuu_copyvec(stiffness.data() + cell_idx * ndof_cell * ndof_cell,
K_cell.data(), ndof_cell * ndof_cell);
}
}
}
ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_);
}
// create bounding volume hierarchy
src/user/user_model.cc
+30 -8
View File
@@ -2190,8 +2190,13 @@ void mjCModel::SetSizes() {
nflexshelldata += (int)flexes_[i]->shell.size();
nflexevpair += (int)flexes_[i]->evpair.size()/2;
nflextexcoord += (flexes_[i]->HasTexcoord() ? flexes_[i]->get_texcoord().size()/2 : 0);
if (flexes_[i]->order_ != 0) {
extra_stiffness_size += (3 * flexes_[i]->nnode) * (3 * flexes_[i]->nnode);
if (flexes_[i]->spec.order != 0) {
int npc = (int)pow(flexes_[i]->spec.order + 1, 3);
int ndof_cell = 3 * npc;
int ncells = flexes_[i]->spec.cellcount[0] *
flexes_[i]->spec.cellcount[1] *
flexes_[i]->spec.cellcount[2];
extra_stiffness_size += ncells * ndof_cell * ndof_cell;
}
if (flexes_[i]->interpolated || flexes_[i]->rigid) {
continue;
@@ -3467,18 +3472,30 @@ void mjCModel::CopyObjects(mjModel* m) {
mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4);
// elasticity
if (pfl->order_ == 0) {
if (pfl->spec.order == 0) {
m->flex_stiffnessadr[i] = 21 * elem_adr;
} else {
m->flex_stiffnessadr[i] = current_extra_stiffness_adr;
current_extra_stiffness_adr += (3 * pfl->nnode) * (3 * pfl->nnode);
int npc = (int)pow(pfl->spec.order + 1, 3);
int ndof_cell = 3 * npc;
int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2];
current_extra_stiffness_adr += ncells * ndof_cell * ndof_cell;
}
if (!pfl->stiffness.empty()) {
mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i], pfl->stiffness.data(), pfl->stiffness.size());
mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i],
pfl->stiffness.data(), pfl->stiffness.size());
} else {
int size = (pfl->order_ == 0) ? 21 * pfl->nelem : (3 * pfl->nnode) * (3 * pfl->nnode);
mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], size);
int stiff_size;
if (pfl->spec.order == 0) {
stiff_size = 21 * pfl->nelem;
} else {
int npc = (int)pow(pfl->spec.order + 1, 3);
int ndof_cell = 3 * npc;
int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2];
stiff_size = ncells * ndof_cell * ndof_cell;
}
mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], stiff_size);
}
if (!pfl->bending.empty()) {
mjuu_copyvec(m->flex_bending + 17 * edge_adr, pfl->bending.data(), pfl->bending.size());
@@ -3613,7 +3630,12 @@ void mjCModel::CopyObjects(mjModel* m) {
}
// set interpolation type, only two types for now
m->flex_interp[i] = pfl->order_;
m->flex_interp[i] = pfl->spec.order;
// set cell count for multi-cell finite cell method
m->flex_cellnum[3*i+0] = pfl->spec.cellcount[0];
m->flex_cellnum[3*i+1] = pfl->spec.cellcount[1];
m->flex_cellnum[3*i+2] = pfl->spec.cellcount[2];
// convert edge pairs to int array, set edge rigid
for (int k=0; k < pfl->nedge; k++) {
src/user/user_objects.h
+2 -2
View File
@@ -1042,7 +1042,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
static constexpr int kNumEdges[3] = {1, 3, 6}; // number of edges per element indexed by dim
void SetOrder(int order) { order_ = order; } // set interpolation order
private:
void Compile(const mjVFS* vfs); // compiler
@@ -1052,7 +1052,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
std::vector<double> vert0_; // vertex positions in [0, 1]^d in the bounding box
std::vector<double> node0_; // node Cartesian positions
int order_ = 0; // interpolation order
// stiffness caching
std::string ComputeStiffnessCacheKey() const;
src/xml/xml_native_reader.cc
+14 -3
View File
@@ -315,7 +315,7 @@ std::vector<const char*> MJCF[nMJCF] = {
{">"},
{">"},
{"flexcomp", "*", "name", "type", "group", "dim", "dof",
"count", "spacing", "radius", "rigid", "mass", "inertiabox",
"count", "cellcount", "spacing", "radius", "rigid", "mass", "inertiabox",
"scale", "file", "point", "element", "texcoord", "material", "rgba",
"flatskin", "pos", "quat", "axisangle", "xyaxes", "zaxis", "euler", "origin"},
{"<"},
@@ -334,8 +334,8 @@ std::vector<const char*> MJCF[nMJCF] = {
{"deformable", "*"},
{"<"},
{"flex", "*", "name", "group", "dim", "radius", "material",
"rgba", "flatskin", "body", "vertex", "element", "texcoord", "elemtexcoord", "node"},
{"flex", "*", "name", "group", "dim", "radius", "material", "rgba", "flatskin", "body",
"vertex", "element", "texcoord", "elemtexcoord", "node", "cellcount", "dof"},
{"<"},
{"contact", "?", "contype", "conaffinity", "condim", "priority",
"friction", "solmix", "solref", "solimp", "margin", "gap",
@@ -1501,6 +1501,16 @@ void mjXReader::OneFlex(XMLElement* elem, mjsFlex* flex) {
ReadAttrInt(elem, "dim", &flex->dim);
ReadAttrInt(elem, "group", &flex->group);
flex->cellcount[0] = 1;
flex->cellcount[1] = 1;
flex->cellcount[2] = 1;
ReadAttr(elem, "cellcount", 3, flex->cellcount, text);
flex->order = 0;
if (MapValue(elem, "dof", &n, fdof_map, mjNFCOMPDOFS)) {
flex->order = (n == mjFCOMPDOF_QUADRATIC) ? 2 : (n == mjFCOMPDOF_TRILINEAR ? 1 : 0);
}
// read data vectors
if (ReadAttrTxt(elem, "body", text, true)) {
mjs_setStringVec(flex->vertbody, text.c_str());
@@ -2794,6 +2804,7 @@ void mjXReader::OneFlexcomp(XMLElement* elem, mjsBody* body, const mjVFS* vfs) {
fcomp.type = (mjtFcompType)n;
}
ReadAttr(elem, "count", 3, fcomp.count, text);
ReadAttr(elem, "cellcount", 3, fcomp.cellcount, text);
ReadAttr(elem, "spacing", 3, fcomp.spacing, text);
ReadAttr(elem, "scale", 3, fcomp.scale, text);
ReadAttr(elem, "mass", 1, &fcomp.mass, text);
src/xml/xml_native_writer.cc
+7
View File
@@ -141,6 +141,13 @@ void mjXWriter::OneFlex(XMLElement* elem, const mjCFlex* flex) {
WriteAttrKey(elem, "flatskin", bool_map, 2, flex->flatskin, defflex.flatskin);
WriteAttrInt(elem, "dim", flex->dim, defflex.dim);
WriteAttrInt(elem, "group", flex->group, defflex.group);
WriteAttr(elem, "cellcount", 3, flex->spec.cellcount, defflex.spec.cellcount);
if (flex->spec.order != defflex.spec.order) {
string dof_str = "full";
if (flex->spec.order == 1) dof_str = "trilinear";
else if (flex->spec.order == 2) dof_str = "quadratic";
WriteAttrTxt(elem, "dof", dof_str);
}
// data vectors
if (!flex->get_vertbody().empty()) {