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;