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

Compute diagonal indices in mj_sqrMatTDSparse

Browse Source 
PiperOrigin-RevId: 713714087
Change-Id: Icc8eae74e6e47ba1d12a6e9aa0d774ab006cfe38
This commit is contained in:
Yuval Tassa
2025-01-09 10:16:37 -08:00
committed by Copybara-Service
parent daba00ca9d
commit 00010f5848
6 changed files with 53 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+5 -8
View File
@@ -2174,23 +2174,20 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// construct supernodes
mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind);
// AR = JM2 * JM2'
// pre-count efc_AR_rownnz, efc_AR_rowadr
mju_sqrMatTDSparseCount(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT,
rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
// AR = JM2 * JM2'
int* diagind = mjSTACKALLOC(d, nefc, int);
mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc,
d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind,
rownnzT, rowadrT, colindT, NULL,
rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
rownnz, rowadr, colind, rowsuper, d, diagind);
// add R to diagonal of AR
for (int i=0; i < nefc; i++) {
for (int j=0; j < d->efc_AR_rownnz[i]; j++) {
if (i == d->efc_AR_colind[d->efc_AR_rowadr[i]+j]) {
d->efc_AR[d->efc_AR_rowadr[i]+j] += d->efc_R[i];
break;
}
}
d->efc_AR[diagind[i]] += d->efc_R[i];
}
}