diff --git a/doc/includes/references.h b/doc/includes/references.h index 6b9610a9..948d88f9 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -314,10 +314,12 @@ struct mjData_ { int* B_colind; // body-dof: column indices of non-zeros (nB x 1) int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1) int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1) + int* C_diag; // reduced dof-dof: index of diagonal element (nv x 1) int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1) int* mapM2C; // index mapping from M to C (nC x 1) int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1) int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1) + int* D_diag; // dof-dof: index of diagonal element (nv x 1) int* D_colind; // dof-dof: column indices of non-zeros (nD x 1) int* mapM2D; // index mapping from M to D (nD x 1) int* mapD2M; // index mapping from D to M (nM x 1) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 7ef96b12..d7185142 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -342,10 +342,12 @@ struct mjData_ { int* B_colind; // body-dof: column indices of non-zeros (nB x 1) int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1) int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1) + int* C_diag; // reduced dof-dof: index of diagonal element (nv x 1) int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1) int* mapM2C; // index mapping from M to C (nC x 1) int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1) int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1) + int* D_diag; // dof-dof: index of diagonal element (nv x 1) int* D_colind; // dof-dof: column indices of non-zeros (nD x 1) int* mapM2D; // index mapping from M to D (nD x 1) int* mapD2M; // index mapping from D to M (nM x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 95a06b2c..a54c4ff5 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -655,10 +655,12 @@ X ( int, B_colind, nB, 1 ) \ X ( int, C_rownnz, nv, 1 ) \ X ( int, C_rowadr, nv, 1 ) \ + X ( int, C_diag, nv, 1 ) \ X ( int, C_colind, nC, 1 ) \ X ( int, mapM2C, nC, 1 ) \ X ( int, D_rownnz, nv, 1 ) \ X ( int, D_rowadr, nv, 1 ) \ + X ( int, D_diag, nv, 1 ) \ X ( int, D_colind, nD, 1 ) \ X ( int, mapM2D, nD, 1 ) \ X ( int, mapD2M, nM, 1 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index c1e00a86..194e410d 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -5396,6 +5396,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='reduced dof-dof: address of each row in C_colind (nv x 1)', # pylint: disable=line-too-long ), + StructFieldDecl( + name='C_diag', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='reduced dof-dof: index of diagonal element', + array_extent=('nv',), + ), StructFieldDecl( name='C_colind', type=PointerType( @@ -5428,6 +5436,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='dof-dof: address of each row in D_colind', array_extent=('nv',), ), + StructFieldDecl( + name='D_diag', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='dof-dof: index of diagonal element', + array_extent=('nv',), + ), StructFieldDecl( name='D_colind', type=PointerType( diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index d09ae016..10182e76 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -332,10 +332,12 @@ def make_data( 'B_colind': (m.nB, jp.int32), 'C_rownnz': (m.nv, jp.int32), 'C_rowadr': (m.nv, jp.int32), + 'C_diag': (m.nv, jp.int32), 'C_colind': (m.nC, jp.int32), 'mapM2C': (m.nC, jp.int32), 'D_rownnz': (m.nv, jp.int32), 'D_rowadr': (m.nv, jp.int32), + 'D_diag': (m.nv, jp.int32), 'D_colind': (m.nD, jp.int32), 'mapM2D': (m.nD, jp.int32), 'mapD2M': (m.nM, jp.int32), diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 35fbdb19..83fe9f4a 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -1273,10 +1273,12 @@ class Data(PyTreeNode): B_colind: body-dof: column indices of non-zeros (nB,) C_rownnz: reduced dof-dof: non-zeros in each row (nv,) C_rowadr: reduced dof-dof: address of each row in C_colind (nv,) + C_diag: reduced dof-dof: index of diagonal element (nv,) C_colind: reduced dof-dof: column indices of non-zeros (nC,) mapM2C: index mapping from M to C (nC,) D_rownnz: dof-dof: non-zeros in each row (nv,) D_rowadr: dof-dof: address of each row in D_colind (nv,) + D_diag: dof-dof: index of diagonal element (nv,) D_colind: dof-dof: column indices of non-zeros (nD,) mapM2D: index mapping from M to D (nD,) mapD2M: index mapping from D to M (nM,) @@ -1402,10 +1404,12 @@ class Data(PyTreeNode): B_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name C_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name C_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + C_diag: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name C_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name mapM2C: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name D_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name D_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + D_diag: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name D_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name mapM2D: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name mapD2M: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index ec24d4cf..59704140 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -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); } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 24bd02bb..10ad5668 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -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); diff --git a/test/fixture.h b/test/fixture.h index 4d65623d..69224c15 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -118,7 +118,7 @@ inline void PrintMatrix(const mjtNum* mat, int nrow, int ncol, int p = 5) { std::cerr << "\n"; for (int r = 0; r < nrow; r++) { for (int c = 0; c < ncol; c++) { - std::cerr << std::fixed << std::setw(4 + p) << mat[c + r*ncol] << " "; + std::cerr << std::fixed << std::setw(3 + p) << mat[c + r*ncol] << " "; } std::cerr << "\n"; } diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index de038bef..8dec655e 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4940,10 +4940,12 @@ public unsafe struct mjData_ { public int* B_colind; public int* C_rownnz; public int* C_rowadr; + public int* C_diag; public int* C_colind; public int* mapM2C; public int* D_rownnz; public int* D_rowadr; + public int* D_diag; public int* D_colind; public int* mapM2D; public int* mapD2M;