Fill-in qM CSR matrix structure on mjData (currently unused).

PiperOrigin-RevId: 733258529
Change-Id: I903548ae60139e4858550194e1aedf0efa2f3033
This commit is contained in:
Yuval Tassa
2025-03-04 03:17:56 -08:00
committed by Copybara-Service
parent c72c83aa3b
commit 7d9df8c9a9
2 changed files with 75 additions and 23 deletions
+44 -23
View File
@@ -928,7 +928,7 @@ 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* diag, int* colind,
int reduced) {
int reduced, int upper) {
int nv = m->nv;
// no dofs, nothing to do
@@ -947,13 +947,13 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
rownnz[i]++;
// process below diagonal unless reduced and dof is simple
if (!reduced || !m->dof_simplenum[i]) {
if (!(reduced && m->dof_simplenum[i])) {
while ((j = m->dof_parentid[j]) >= 0) {
// both reduced and non-reduced have lower triangle
rownnz[i]++;
// only non-reduced has upper triangle
if (!reduced) rownnz[j]++;
// add upper triangle if requested
if (upper) rownnz[j]++;
}
}
}
@@ -972,14 +972,14 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
colind[rowadr[i] + remaining[i]] = i;
// process below diagonal unless reduced and dof is simple
if (!reduced || !m->dof_simplenum[i]) {
if (!(reduced && m->dof_simplenum[i])) {
int j = i;
while ((j = m->dof_parentid[j]) >= 0) {
remaining[i]--;
colind[rowadr[i] + remaining[i]] = j;
// only non-reduced has upper triangle
if (!reduced) {
// add upper triangle if requested
if (upper) {
remaining[j]--;
colind[rowadr[j] + remaining[j]] = i;
}
@@ -995,7 +995,8 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
}
// check total nnz; SHOULD NOT OCCUR
if (rowadr[nv - 1] + rownnz[nv - 1] != (reduced ? m->nC : m->nD)) {
int expected_nnz = upper ? m->nD : (reduced ? m->nC : m->nM);
if (rowadr[nv - 1] + rownnz[nv - 1] != expected_nnz) {
mjERROR("sum of rownnz different from expected");
}
@@ -1131,18 +1132,23 @@ static void checkDBSparse(const mjModel* m, mjData* d) {
// integer valued dst[D or C] = src[M], handle different sparsity representations
// integer valued dst[D or C or M] = src[M (legacy)], handle different sparsity representations
static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
int reduced) {
int reduced, int upper) {
int nv = m->nv;
const int* rownnz;
const int* rowadr;
if (reduced) {
if (reduced && !upper) {
rownnz = d->C_rownnz;
rowadr = d->C_rowadr;
} else {
} else if (!reduced && !upper) {
rownnz = d->M_rownnz;
rowadr = d->M_rowadr;
} else if (!reduced && upper) {
rownnz = d->D_rownnz;
rowadr = d->D_rowadr;
} else {
mjERROR("unsupported sparsity structure (reduced + upper)");
}
mj_markStack(d);
@@ -1160,14 +1166,14 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
adr++;
// process below diagonal unless reduced and dof is simple
if (!reduced || !m->dof_simplenum[i]) {
if (!(reduced && m->dof_simplenum[i])) {
int j = i;
while ((j = m->dof_parentid[j]) >= 0) {
remaining[i]--;
dst[rowadr[i] + remaining[i]] = src[adr];
// only non-reduced has upper triangle
if (!reduced) {
// add upper triangle if requested
if (upper) {
remaining[j]--;
dst[rowadr[j] + remaining[j]] = src[adr];
}
@@ -1213,8 +1219,8 @@ static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int
// construct index mappings between M <-> D and M -> C
static void makeDofDofmap(const mjModel* m, mjData* d) {
// construct index mappings between M <-> D, M -> C, M (legacy) -> M (CSR)
static void makeDofDofmaps(const mjModel* m, mjData* d) {
int nM = m->nM, nC = m->nC, nD = m->nD;
mj_markStack(d);
@@ -1222,7 +1228,7 @@ static void makeDofDofmap(const mjModel* m, mjData* d) {
int* M = mjSTACKALLOC(d, nM, int);
for (int i=0; i < nM; i++) M[i] = i;
for (int i=0; i < nD; i++) d->mapM2D[i] = -1;
copyM2Sparse(m, d, d->mapM2D, M, /*reduced=*/0);
copyM2Sparse(m, d, d->mapM2D, M, /*reduced=*/0, /*upper=*/1);
// check that all indices are filled in
for (int i=0; i < nD; i++) {
@@ -1246,7 +1252,7 @@ static void makeDofDofmap(const mjModel* m, mjData* d) {
// make mapM2C
for (int i=0; i < nC; i++) d->mapM2C[i] = -1;
copyM2Sparse(m, d, d->mapM2C, M, /*reduced=*/1);
copyM2Sparse(m, d, d->mapM2C, M, /*reduced=*/1, /*upper=*/0);
// check that all indices are filled in
for (int i=0; i < nC; i++) {
@@ -1255,6 +1261,17 @@ static void makeDofDofmap(const mjModel* m, mjData* d) {
}
}
// make mapM2M
for (int i=0; i < nM; i++) d->mapM2M[i] = -1;
copyM2Sparse(m, d, d->mapM2M, M, /*reduced=*/0, /*upper=*/0);
// check that all indices are filled in
for (int i=0; i < nM; i++) {
if (d->mapM2M[i] < 0) {
mjERROR("unassigned index in mapM2M");
}
}
mj_freeStack(d);
}
@@ -1965,15 +1982,19 @@ 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_diag, d->D_colind, /*reduced=*/0);
makeDofDofSparse(m, d, d->D_rownnz, d->D_rowadr, d->D_diag, d->D_colind,
/*reduced=*/0, /*upper=*/1);
// make B, check D and B
makeBSparse(m, d);
checkDBSparse(m, d);
// make C
makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, NULL, d->C_colind, /*reduced=*/1);
makeDofDofmap(m, d);
// make M, C
makeDofDofSparse(m, d, d->M_rownnz, d->M_rowadr, NULL, d->M_colind, /*reduced=*/0, /*upper=*/0);
makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, NULL, d->C_colind, /*reduced=*/1, /*upper=*/0);
// make index mappings: mapM2D, mapD2M, mapM2C, mapM2M
makeDofDofmaps(m, d);
}
// restore pluginstate and plugindata
+31
View File
@@ -1161,6 +1161,37 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
}
fprintf(fp, "\n\n");
// M sparse structure
mj_printSparsity("M: inertia matrix", m->nv, m->nv, d->M_rowadr, NULL, d->M_rownnz,
NULL, d->M_colind, fp);
fprintf(fp, NAME_FORMAT, "M_rownnz");
for (int i = 0; i < m->nv; i++) {
fprintf(fp, " %d", d->M_rownnz[i]);
}
fprintf(fp, "\n\n");
// M_rowadr
fprintf(fp, NAME_FORMAT, "M_rowadr");
for (int i = 0; i < m->nv; i++) {
fprintf(fp, " %d", d->M_rowadr[i]);
}
fprintf(fp, "\n\n");
// M_colind
fprintf(fp, NAME_FORMAT, "M_colind");
for (int i = 0; i < m->nM; i++) {
fprintf(fp, " %d", d->M_colind[i]);
}
fprintf(fp, "\n\n");
// mapM2M
fprintf(fp, NAME_FORMAT, "mapM2M");
for (int i = 0; i < m->nM; i++) {
fprintf(fp, " %d", d->mapM2M[i]);
}
fprintf(fp, "\n\n");
// C sparse structure
mj_printSparsity("C: reduced dof-dof matrix", m->nv, m->nv, d->C_rowadr, NULL, d->C_rownnz,
NULL, d->C_colind, fp);