Skip to content

Commit

Permalink
Improve printing of inertia-like matrices, allow mj_printData to take…
Browse files Browse the repository at this point in the history 
… `const mjData*`

PiperOrigin-RevId: 718176884
Change-Id: I057dabe6973e19db86e9ff9ce4430c6d17dfce2b
  • Loading branch information
@yuvaltassa @copybara-github
yuvaltassa authored and copybara-github committed Jan 22, 2025
1 parent 6a7a7a5 commit 67dd482
Show file tree
Hide file tree
Showing 5 changed files with 56 additions and 40 deletions.
4 changes: 2 additions & 2 deletions doc/includes/references.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3186,9 +3186,9 @@ int mjs_activatePlugin(mjSpec* s, const char* name);
int mjs_setDeepCopy(mjSpec* s, int deepcopy);
void mj_printFormattedModel(const mjModel* m, const char* filename, const char* float_format);
void mj_printModel(const mjModel* m, const char* filename);
void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filename,
const char* float_format);
void mj_printData(const mjModel* m, mjData* d, const char* filename);
void mj_printData(const mjModel* m, const mjData* d, const char* filename);
void mju_printMat(const mjtNum* mat, int nr, int nc);
void mju_printMatSparse(const mjtNum* mat, int nr,
const int* rownnz, const int* rowadr, const int* colind);
Expand Down
4 changes: 2 additions & 2 deletions include/mujoco/mujoco.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -257,11 +257,11 @@ MJAPI void mj_printModel(const mjModel* m, const char* filename);

// Print mjData to text file, specifying format.
// float_format must be a valid printf-style format string for a single float value
MJAPI void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
MJAPI void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filename,
const char* float_format);

// Print data to text file.
MJAPI void mj_printData(const mjModel* m, mjData* d, const char* filename);
MJAPI void mj_printData(const mjModel* m, const mjData* d, const char* filename);

// Print matrix to screen.
MJAPI void mju_printMat(const mjtNum* mat, int nr, int nc);
Expand Down
4 changes: 2 additions & 2 deletions introspect/functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1123,7 +1123,7 @@
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData'),
inner_type=ValueType(name='mjData', is_const=True),
),
),
FunctionParameterDecl(
Expand Down Expand Up @@ -1155,7 +1155,7 @@
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData'),
inner_type=ValueType(name='mjData', is_const=True),
),
),
FunctionParameterDecl(
Expand Down
80 changes: 48 additions & 32 deletions src/engine/engine_print.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,43 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,



// print sparse inertia-like matrix
static void printInertia(const char* str, const mjtNum* mat, const mjModel* m,
FILE* fp, const char* float_format) {
int nv = m->nv;
// if no data, or too many rows to be visually useful, return
if (!mat || !nv || nv > 300) {
return;
}

// get length of string produced by float_format
char test[100];
int len = snprintf(test, sizeof(test), float_format, 0.0);

fprintf(fp, "%s\n", str);

for (int i=0; i < nv; i++) {
fprintf(fp, " ");
int adr = (i == nv-1) ? m->nM - 1 : m->dof_Madr[i+1] - 1;
for (int k=0; k <= i; k++) {
int j = i;
while (j != k && j >= 0) {
j = m->dof_parentid[j];
}
if (j == k) {
fprintf(fp, " ");
fprintf(fp, float_format, mat[adr--]);
} else {
for (int d=0; d < len+1; d++) fprintf(fp, " ");
}
}
fprintf(fp, "\n");
}
fprintf(fp, "\n");
}



// print sparse matrix structure
void mj_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) {
Expand Down Expand Up @@ -859,16 +896,13 @@ void mj_printModel(const mjModel* m, const char* filename) {

// print mjModel to text file, specifying format. float_format must be a
// valid printf-style format string for a single float value
void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filename,
const char* float_format) {
// stack in use, SHOULD NOT OCCUR
if (d->pstack) {
mjERROR("attempting to print mjData when stack is in use");
}

mjtNum *M = NULL;
mj_markStack(d);

// check format string
if (!validateFloatFormat(float_format)) {
mju_warning("WARNING: Received invalid float_format. Using default instead.");
Expand All @@ -886,15 +920,9 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
// check for nullptr
if (!fp) {
mju_warning("Could not open file '%s' for writing mjModel", filename);
mj_freeStack(d);
return;
}

// allocate full inertia if it's small
if (m->nv <= 200) {
M = mjSTACKALLOC(d, m->nv*m->nv, mjtNum);
}

#ifdef MEMORY_SANITIZER
// If memory sanitizer is active, d->buffer will be marked as poisoned, even
// though it's really initialized to 0. This catches unintentionally
Expand Down Expand Up @@ -990,7 +1018,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
fprintf(fp, " solver_niter = %d\n", d->solver_niter[island]);
fprintf(fp, " solver_nnz = %d\n", d->solver_nnz[island]);
for (int i=0; i < niter_stat; i++) {
mjSolverStat* stat = d->solver + island*mjNSOLVER + i;
const mjSolverStat* stat = d->solver + island*mjNSOLVER + i;
fprintf(fp, " %d: improvement = ", i);
fprintf(fp, float_format, stat->improvement);
fprintf(fp, " gradient = ");
Expand Down Expand Up @@ -1097,16 +1125,9 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
d->moment_rowadr, d->moment_colind, fp, float_format);
printArray("CRB", m->nbody, 10, d->crb, fp, float_format);

if (M) {
// construct and print full M matrix
mj_fullM(m, M, d->qM);
printArray("QM", m->nv, m->nv, M, fp, float_format);

// construct and print full LD matrix
mj_fullM(m, M, d->qLD);
printArray("QLD", m->nv, m->nv, M, fp, float_format);
}
printInertia("QM", d->qM, m, fp, float_format);

printInertia("QLD", d->qLD, m, fp, float_format);
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);

// B sparse structure
Expand Down Expand Up @@ -1204,16 +1225,13 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
}
fprintf(fp, "\n\n");

if (M) {
// print qDeriv
mju_sparse2dense(M, d->qDeriv, m->nv, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind);
printArray("QDERIV", m->nv, m->nv, M, fp, float_format);
// print qDeriv
printSparse("QDERIV", d->qDeriv, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind,
fp, float_format);

// print qLU
mju_sparse2dense(M, d->qLU, m->nv, m->nv, d->D_rownnz, d->D_rowadr,
d->D_colind);
printArray("QLU", m->nv, m->nv, M, fp, float_format);
}
// print qLU
printSparse("QLU", d->qLU, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind,
fp, float_format);

// contact
fprintf(fp, "CONTACT\n");
Expand Down Expand Up @@ -1401,8 +1419,6 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
if (filename) {
fclose(fp);
}

mj_freeStack(d);
}


Expand All @@ -1412,6 +1428,6 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,


// print mjData to text file
void mj_printData(const mjModel* m, mjData* d, const char* filename) {
void mj_printData(const mjModel* m, const mjData* d, const char* filename) {
mj_printFormattedData(m, d, filename, FLOAT_FORMAT);
}
4 changes: 2 additions & 2 deletions src/engine/engine_print.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,11 +36,11 @@ MJAPI void mj_printModel(const mjModel* m, const char* filename);

// print mjData to text file, specifying format
// float_format must be a valid printf-style format string for a single float value
MJAPI void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
MJAPI void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filename,
const char* float_format);

// print data to text file
MJAPI void mj_printData(const mjModel* m, mjData* d, const char* filename);
MJAPI void mj_printData(const mjModel* m, const mjData* d, const char* filename);

// print sparse matrix structure
MJAPI void mj_printSparsity(const char* str, int nr, int nc, const int* rowadr, const int* diag,
Expand Down

0 comments on commit 67dd482

Please sign in to comment.