Store index of diagonal element for C and D in mjData.

PiperOrigin-RevId: 704287869
Change-Id: I15d1faa0b5c3ff54721f692ec6329726456ec318
This commit is contained in:
Yuval Tassa
2024-12-09 08:00:29 -08:00
committed by Copybara-Service
parent b8344c191f
commit 3a14ada308
10 changed files with 67 additions and 18 deletions
+17 -3
View File
@@ -921,7 +921,8 @@ int mj_sizeModel(const mjModel* m) {
// construct sparse representation of dof-dof matrix
static void makeDofDofSparse(const mjModel* m, mjData* d,
int* rownnz, int* rowadr, int* colind, int reduced) {
int* rownnz, int* rowadr, int* diag, int* colind,
int reduced) {
int nv = m->nv;
// no dofs, nothing to do
@@ -986,6 +987,19 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
mjERROR("sum of rownnz different from expected");
}
// find diagonal indices
for (int i = 0; i < nv; i++) {
int adr = rowadr[i];
int j = 0;
while (colind[adr + j] < i && j < rownnz[i]) {
j++;
}
if (colind[adr + j] != i) {
mjERROR("diagonal index not found");
}
diag[i] = j;
}
mj_freeStack(d);
}
@@ -1915,14 +1929,14 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
// construct sparse matrix representations
if (m->body_dofadr) {
// make D
makeDofDofSparse(m, d, d->D_rownnz, d->D_rowadr, d->D_colind, /*reduced=*/0);
makeDofDofSparse(m, d, d->D_rownnz, d->D_rowadr, d->D_diag, d->D_colind, /*reduced=*/0);
// make B, check D and B
makeBSparse(m, d);
checkDBSparse(m, d);
// make C
makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, d->C_colind, /*reduced=*/1);
makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind, /*reduced=*/1);
makeDmap(m, d);
}
+19 -14
View File
@@ -114,8 +114,8 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,
// print sparse matrix structure
static void printSparsity(const char* str, int nr, int nc, const int* rowadr, const int* rownnz,
const int* rowsuper, const int* colind, FILE* fp) {
static void printSparsity(const char* str, int nr, int nc, const int* rowadr, const int* diag,
const int* rownnz, const int* rowsuper, const int* colind, FILE* fp) {
// if no rows / columns, or too many columns to be visually useful, return
if (!nr || !nc || nc > 300) {
return;
@@ -130,7 +130,11 @@ static void printSparsity(const char* str, int nr, int nc, const int* rowadr, co
int nnz = 0;
for (int c=0; c < nc; c++) {
if (nnz < rownnz[r] && colind[adr + nnz] == c) {
fprintf(fp, "x");
if (diag && diag[r] == nnz) {
fprintf(fp, "D");
} else {
fprintf(fp, "x");
}
nnz++;
} else {
fprintf(fp, " ");
@@ -1057,7 +1061,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
printArray("FLEXEDGE_J", m->nflexedge, m->nv, d->flexedge_J, fp, float_format);
} else {
printSparsity("FLEXEDGE_J: flex edge connectivity", m->nflexedge, m->nv,
d->flexedge_J_rowadr, d->flexedge_J_rownnz, NULL, d->flexedge_J_colind, fp);
d->flexedge_J_rowadr, NULL, d->flexedge_J_rownnz, NULL, d->flexedge_J_colind, fp);
printArrayInt("FLEXEDGE_J_ROWNNZ", m->nflexedge, 1, d->flexedge_J_rownnz, fp);
printArrayInt("FLEXEDGE_J_ROWADR", m->nflexedge, 1, d->flexedge_J_rowadr, fp);
printSparse("FLEXEDGE_J", d->flexedge_J, m->nflexedge, d->flexedge_J_rownnz,
@@ -1069,8 +1073,8 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
if (!mj_isSparse(m)) {
printArray("TEN_MOMENT", m->ntendon, m->nv, d->ten_J, fp, float_format);
} else {
printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, d->ten_J_rownnz,
NULL, d->ten_J_colind, fp);
printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, NULL,
d->ten_J_rownnz, NULL, d->ten_J_colind, fp);
printArrayInt("TEN_J_ROWNNZ", m->ntendon, 1, d->ten_J_rownnz, fp);
printArrayInt("TEN_J_ROWADR", m->ntendon, 1, d->ten_J_rowadr, fp);
printSparse("TEN_J", d->ten_J, m->ntendon, d->ten_J_rownnz,
@@ -1087,7 +1091,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
printArray("ACTUATOR_LENGTH", m->nu, 1, d->actuator_length, fp, float_format);
printSparsity("actuator_moment", m->nu, m->nv,
d->moment_rowadr, d->moment_rownnz, NULL, d->moment_colind, fp);
d->moment_rowadr, NULL, d->moment_rownnz, NULL, d->moment_colind, fp);
printSparse("ACTUATOR_MOMENT", d->actuator_moment, m->nu, d->moment_rownnz,
d->moment_rowadr, d->moment_colind, fp, float_format);
printArray("CRB", m->nbody, 10, d->crb, fp, float_format);
@@ -1106,7 +1110,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format);
// B sparse structure
printSparsity("B: body-dof matrix", m->nbody, m->nv, d->B_rowadr, d->B_rownnz, NULL,
printSparsity("B: body-dof matrix", m->nbody, m->nv, d->B_rowadr, NULL, d->B_rownnz, NULL,
d->B_colind, fp);
// B_rownnz
@@ -1131,8 +1135,8 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
fprintf(fp, "\n\n");
// C sparse structure
printSparsity("C: reduced dof-dof matrix", m->nv, m->nv, d->C_rowadr, d->C_rownnz, NULL,
d->C_colind, fp);
printSparsity("C: reduced dof-dof matrix", m->nv, m->nv, d->C_rowadr, d->C_diag, d->C_rownnz,
NULL, d->C_colind, fp);
fprintf(fp, NAME_FORMAT, "C_rownnz");
for (int i = 0; i < m->nv; i++) {
@@ -1162,7 +1166,8 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
fprintf(fp, "\n\n");
// D sparse structure
printSparsity("D: dof-dof matrix", m->nv, m->nv, d->D_rowadr, d->D_rownnz, NULL, d->D_colind, fp);
printSparsity("D: dof-dof matrix", m->nv, m->nv,
d->D_rowadr, d->D_diag, d->D_rownnz, NULL, d->D_colind, fp);
// D_rownnz
fprintf(fp, NAME_FORMAT, "D_rownnz");
@@ -1259,13 +1264,13 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
printArray("EFC_J", d->nefc, m->nv, d->efc_J, fp, float_format);
printArray("EFC_AR", d->nefc, d->nefc, d->efc_AR, fp, float_format);
} else {
printSparsity("J: constraint Jacobian", d->nefc, m->nv, d->efc_J_rowadr, d->efc_J_rownnz,
d->efc_J_rowsuper, d->efc_J_colind, fp);
printSparsity("J: constraint Jacobian", d->nefc, m->nv,
d->efc_J_rowadr, NULL, d->efc_J_rownnz, d->efc_J_rowsuper, d->efc_J_colind, fp);
printArrayInt("EFC_J_ROWNNZ", d->nefc, 1, d->efc_J_rownnz, fp);
printArrayInt("EFC_J_ROWADR", d->nefc, 1, d->efc_J_rowadr, fp);
printSparse("EFC_J", d->efc_J, d->nefc, d->efc_J_rownnz,
d->efc_J_rowadr, d->efc_J_colind, fp, float_format);
printSparsity("JT: constraint Jacobian transposed", m->nv, d->nefc, d->efc_JT_rowadr,
printSparsity("JT: constraint Jacobian transposed", m->nv, d->nefc, d->efc_JT_rowadr, NULL,
d->efc_JT_rownnz, d->efc_JT_rowsuper, d->efc_JT_colind, fp);
printArrayInt("EFC_AR_ROWNNZ", d->nefc, 1, d->efc_AR_rownnz, fp);
printArrayInt("EFC_AR_ROWADR", d->nefc, 1, d->efc_AR_rowadr, fp);