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

PiperOrigin-RevId: 733258529
Change-Id: I903548ae60139e4858550194e1aedf0efa2f3033

Author: yuvaltassa

src/engine/engine_io.c

@@ -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,16 +1219,16 @@ 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);
 
   // make mapM2D
   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

src/engine/engine_print.c

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