Refactor flex strain constraints to be per-cell.

Each mjEQ_FLEXSTRAIN equality now represents a single cell within a flex. This allows for more efficient sparse Jacobian computation by only considering the degrees of freedom of the nodes within each specific cell. This change gives a speedup of about 10x on a 3x3x3 model.

PiperOrigin-RevId: 902164069
Change-Id: I78eedf1d5cf39b8989fe9863c22d164922fc0efb
This commit is contained in:
Alessio Quaglino
2026-04-19 07:20:49 -07:00
committed by Copybara-Service
parent a9e61966b5
commit 3d45a33190
7 changed files with 289 additions and 256 deletions
+25 -10
View File
@@ -680,20 +680,35 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
mjs_setDouble(pf->vert, point.data(), point.size());
}
// create edge equality constraint
// create equality constraints
if (equality) {
mjsEquality* pe = mjs_addEquality(&model->spec, &def.spec);
mjs_setDefault(pe->element, &model->Default()->spec);
// equality 1=edge(mjEQ_FLEX), 2=vert(mjEQ_FLEXVERT), 3=strain(mjEQ_FLEXSTRAIN)
if (equality == 1) {
pe->type = mjEQ_FLEX;
} else if (equality == 2) {
pe->type = mjEQ_FLEXVERT;
if (equality == 1 || equality == 2) {
mjsEquality* pe = mjs_addEquality(&model->spec, &def.spec);
mjs_setDefault(pe->element, &model->Default()->spec);
pe->type = (equality == 1) ? mjEQ_FLEX : mjEQ_FLEXVERT;
pe->active = true;
mjs_setString(pe->name1, name.c_str());
} else if (equality == 3) {
pe->type = mjEQ_FLEXSTRAIN;
// create one strain constraint per cell, storing cell index in eq_data
int cell_cx = flex->spec.cellcount[0];
int cell_cy = flex->spec.cellcount[1];
int cell_cz = flex->spec.cellcount[2];
for (int ci = 0; ci < cell_cx; ci++) {
for (int cj = 0; cj < cell_cy; cj++) {
for (int ck = 0; ck < cell_cz; ck++) {
mjsEquality* pe = mjs_addEquality(&model->spec, &def.spec);
mjs_setDefault(pe->element, &model->Default()->spec);
pe->type = mjEQ_FLEXSTRAIN;
pe->active = true;
mjs_setString(pe->name1, name.c_str());
pe->data[0] = ci;
pe->data[1] = cj;
pe->data[2] = ck;
}
}
}
}
pe->active = true;
mjs_setString(pe->name1, name.c_str());
}
return true;