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

PiperOrigin-RevId: 704287869
Change-Id: I15d1faa0b5c3ff54721f692ec6329726456ec318

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)

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)

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           ) \

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(

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),

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

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);
   }
 

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,
+static void printSparsity(const char* str, int nr, int nc, const int* rowadr, const int* diag,
-                          const int* rowsuper, const int* colind, FILE* fp) {
+                          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,
+    printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, NULL,
-                  NULL, d->ten_J_colind, fp);
+                  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,
+  printSparsity("C: reduced dof-dof matrix", m->nv, m->nv, d->C_rowadr, d->C_diag, d->C_rownnz,
-                d->C_colind, fp);
+                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,
+    printSparsity("J: constraint Jacobian", d->nefc, m->nv,
-                  d->efc_J_rowsuper, d->efc_J_colind, fp);
+                  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);

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";
   }

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;