Print supernode values in mj_printData

PiperOrigin-RevId: 702677257
Change-Id: Icecc5dd8604fcadc02bba5fc69c9fd5fd4fbd825

src/engine/engine_print.c

@@ -103,7 +103,7 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,
     fprintf(fp, "  ");
     for (int adr=rowadr[r]; adr < rowadr[r]+rownnz[r]; adr++) {
       fprintf(fp, "  ");
-      fprintf(fp, "%d: ", colind[adr]);
+      fprintf(fp, "%2d: ", colind[adr]);
       fprintf(fp, float_format, mat[adr]);
     }
     fprintf(fp, "\n");
@@ -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,
+static void printSparsity(const char* str, int nr, int nc, const int* rowadr, const int* rownnz,
-                     const int* rowadr, const int* rownnz, const int* colind, FILE* fp) {
+                          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;
@@ -136,7 +136,9 @@ static void printSparsity(const char* str, int nr, int nc,
         fprintf(fp, " ");
       }
     }
-    fprintf(fp, " |\n");
+    fprintf(fp, " |");
+    if (rowsuper && rowsuper[r] > 0) fprintf(fp, " %d", rowsuper[r]);
+    fprintf(fp, "\n");
     if (r < nr-1) fprintf(fp, " ");
   }
   for (int c=0; c < nc+2; c++) fprintf(fp, "-");
@@ -1055,7 +1057,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, d->flexedge_J_colind, fp);
+                  d->flexedge_J_rowadr, 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,
@@ -1067,8 +1069,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,
+    printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, d->ten_J_rownnz,
-             d->ten_J_rowadr, d->ten_J_rownnz, d->ten_J_colind, fp);
+                  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,
@@ -1085,7 +1087,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, d->moment_colind, fp);
+                d->moment_rowadr, 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);
@@ -1104,7 +1106,8 @@ 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, d->B_colind, fp);
+  printSparsity("B: body-dof matrix", m->nbody, m->nv, d->B_rowadr, d->B_rownnz, NULL,
+                d->B_colind, fp);
 
   // B_rownnz
   fprintf(fp, NAME_FORMAT, "B_rownnz");
@@ -1128,7 +1131,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, d->C_colind, fp);
+  printSparsity("C: reduced dof-dof matrix", m->nv, m->nv, d->C_rowadr, d->C_rownnz, NULL,
+                d->C_colind, fp);
 
   fprintf(fp, NAME_FORMAT, "C_rownnz");
   for (int i = 0; i < m->nv; i++) {
@@ -1158,7 +1162,7 @@ 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, d->D_colind, fp);
+  printSparsity("D: dof-dof matrix", m->nv, m->nv, d->D_rowadr, d->D_rownnz, NULL, d->D_colind, fp);
 
   // D_rownnz
   fprintf(fp, NAME_FORMAT, "D_rownnz");
@@ -1255,13 +1259,14 @@ 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,
+    printSparsity("J: constraint Jacobian", d->nefc, m->nv, d->efc_J_rowadr, d->efc_J_rownnz,
-             d->efc_J_rowadr, d->efc_J_rownnz, d->efc_J_colind, fp);
+                  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,
+                  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);
     printSparse("EFC_AR", d->efc_AR, d->nefc, d->efc_AR_rownnz,