Precompute mappings between M <-> D sparse representations.

PiperOrigin-RevId: 670898353
Change-Id: I4b15477b1b6eb4a64cf7d4a0bb5af8fcb95d4e45

doc/includes/references.h

@@ -309,6 +309,8 @@ struct mjData_ {
   int*    D_rownnz;          // non-zeros in each row                            (nv x 1)
   int*    D_rowadr;          // address of each row in D_colind                  (nv x 1)
   int*    D_colind;          // 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)
   int*    B_rownnz;          // non-zeros in each row                            (nbody x 1)
   int*    B_rowadr;          // address of each row in B_colind                  (nbody x 1)
   int*    B_colind;          // column indices of non-zeros                      (nB x 1)

include/mujoco/mjdata.h

@@ -337,6 +337,8 @@ struct mjData_ {
   int*    D_rownnz;          // non-zeros in each row                            (nv x 1)
   int*    D_rowadr;          // address of each row in D_colind                  (nv x 1)
   int*    D_colind;          // 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)
   int*    B_rownnz;          // non-zeros in each row                            (nbody x 1)
   int*    B_rowadr;          // address of each row in B_colind                  (nbody x 1)
   int*    B_colind;          // column indices of non-zeros                      (nB x 1)

include/mujoco/mjxmacro.h

@@ -643,6 +643,8 @@
     X   ( int,       D_rownnz,          nv,          1           ) \
     X   ( int,       D_rowadr,          nv,          1           ) \
     X   ( int,       D_colind,          nD,          1           ) \
+    X   ( int,       mapM2D,            nD,          1           ) \
+    X   ( int,       mapD2M,            nM,          1           ) \
     X   ( int,       B_rownnz,          nbody,       1           ) \
     X   ( int,       B_rowadr,          nbody,       1           ) \
     X   ( int,       B_colind,          nB,          1           ) \

introspect/structs.py

@@ -4866,6 +4866,20 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='column indices of non-zeros                      (nD x 1)',  # pylint: disable=line-too-long
              ),
+             StructFieldDecl(
+                 name='mapM2D',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='index mapping from M to D                        (nD x 1)',  # pylint: disable=line-too-long
+             ),
+             StructFieldDecl(
+                 name='mapD2M',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='index mapping from D to M                        (nM x 1)',  # pylint: disable=line-too-long
+             ),
              StructFieldDecl(
                  name='B_rownnz',
                  type=PointerType(

src/engine/engine_forward.c

@@ -792,7 +792,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
       mjtNum* MhB = mj_stackAllocNum(d, nM);
 
       // MhB = M + h*diag(B)
-      mju_copy(MhB, d->qM, m->nM);
+      mju_copy(MhB, d->qM, nM);
       for (int i=0; i < nv; i++) {
         MhB[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i];
       }
@@ -933,7 +933,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
 // fully implicit in velocity, possibly skipping factorization
 void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
   TM_START;
-  int nv = m->nv;
+  int nv = m->nv, nM = m->nM, nD = m->nD;
 
   mj_markStack(d);
   mjtNum* qfrc = mj_stackAllocNum(d, nv);
@@ -949,7 +949,9 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mjd_smooth_vel(m, d, /* flg_bias = */ 1);
 
       // set qLU = qM
-      mj_copyM2DSparse(m, d, d->qLU, d->qM);
+      for (int i=0; i < nD; i++) {
+        d->qLU[i] = d->qM[d->mapM2D[i]];
+      }
 
       // set qLU = qM - dt*qDeriv
       mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
@@ -970,18 +972,20 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mjd_smooth_vel(m, d, /* flg_bias = */ 0);
 
       // modified mass matrix MhB = qDeriv[Lower]
-      mjtNum* MhB = mj_stackAllocNum(d, m->nM);
+      mjtNum* MhB = mj_stackAllocNum(d, nM);
-      mj_copyD2MSparse(m, d, MhB, d->qDeriv);
+      for (int i=0; i < nM; i++) {
+        MhB[i] = d->qDeriv[d->mapD2M[i]];
+      }
 
       // set MhB = M - dt*qDeriv
-      mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, m->nM);
+      mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
 
       // factorize
       mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL);
     }
 
     // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
-    mju_copy(qacc, qfrc, m->nv);
+    mju_copy(qacc, qfrc, nv);
     mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv);
   } else {
     mjERROR("integrator must be implicit or implicitfast");

src/engine/engine_inverse.c

@@ -71,7 +71,7 @@ void mj_invVelocity(const mjModel* m, mjData* d) {
 
 // convert discrete-time qacc to continuous-time qacc
 static void mj_discreteAcc(const mjModel* m, mjData* d) {
-  int nv = m->nv, dof_damping;
+  int nv = m->nv, nM = m->nM, nD = m->nD, dof_damping;
   mjtNum *qacc = d->qacc;
 
   mj_markStack(d);
@@ -114,7 +114,9 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
     mjd_smooth_vel(m, d, /* flg_bias = */ 1);
 
     // set qLU = qM
-    mj_copyM2DSparse(m, d, d->qLU, d->qM);
+    for (int i=0; i < nD; i++) {
+      d->qLU[i] = d->qM[d->mapM2D[i]];
+    }
 
     // set qLU = qM - dt*qDeriv
     mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
@@ -134,7 +136,9 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
 
     // set M = M - dt*qDeriv (reduced to M nonzeros)
     mjtNum* qDerivReduced = mj_stackAllocNum(d, m->nM);
-    mj_copyD2MSparse(m, d, qDerivReduced, d->qDeriv);
+    for (int i=0; i < nM; i++) {
+      qDerivReduced[i] = d->qDeriv[d->mapD2M[i]];
+    }
     mju_addToScl(d->qM, qDerivReduced, -m->opt.timestep, m->nM);
 
     // set qfrc = (M - dt*qDeriv) * qacc

src/engine/engine_io.c

@@ -1059,6 +1059,107 @@ static void checkDBSparse(const mjModel* m, mjData* d) {
 
 
 
+// integer valued dst[D] = src[M], handle different sparsity representations
+static void copyM2DSparse(const mjModel* m, mjData* d, int* dst, const int* src) {
+  int nv = m->nv;
+  mj_markStack(d);
+
+  // init remaining
+  int* remaining = mj_stackAllocInt(d, nv);
+  mju_copyInt(remaining, d->D_rownnz, nv);
+
+  // copy data
+  for (int i = nv - 1; i >= 0; i--) {
+    // init at diagonal
+    int adr = m->dof_Madr[i];
+    remaining[i]--;
+    dst[d->D_rowadr[i] + remaining[i]] = src[adr];
+    adr++;
+
+    // process below diagonal
+    int j = i;
+    while ((j = m->dof_parentid[j]) >= 0) {
+      remaining[i]--;
+      dst[d->D_rowadr[i] + remaining[i]] = src[adr];
+
+      remaining[j]--;
+      dst[d->D_rowadr[j] + remaining[j]] = src[adr];
+
+      adr++;
+    }
+  }
+
+  // check that none remaining
+  for (int i=0; i < nv; i++) {
+    if (remaining[i]) {
+      mjERROR("unassigned index");
+    }
+  }
+
+  mj_freeStack(d);
+}
+
+
+
+// integer valued dst[M] = src[D lower], handle different sparsity representations
+static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int* src) {
+  int nv = m->nv;
+
+  // copy data
+  for (int i = nv - 1; i >= 0; i--) {
+    // find diagonal in qDeriv
+    int j = 0;
+    while (d->D_colind[d->D_rowadr[i] + j] < i) {
+      j++;
+    }
+
+    // copy
+    int adr = m->dof_Madr[i];
+    while (j >= 0) {
+      dst[adr] = src[d->D_rowadr[i] + j];
+      adr++;
+      j--;
+    }
+  }
+}
+
+
+
+// construct index mappings between D <-> M
+static void makeDmap(const mjModel* m, mjData* d) {
+  int nM = m->nM, nD = m->nD;
+  mj_markStack(d);
+
+  // make mapM2D
+  int* M = mj_stackAllocInt(d, nM);
+  for (int i=0; i < nM; i++) M[i] = i;
+  for (int i=0; i < nD; i++) d->mapM2D[i] = -1;
+  copyM2DSparse(m, d, d->mapM2D, M);
+
+  // check that all indices are filled in
+  for (int i=0; i < nD; i++) {
+    if (d->mapM2D[i] < 0) {
+      mjERROR("unassigned index in mapM2D");
+    }
+  }
+
+  // make mapD2M
+  int* D = mj_stackAllocInt(d, nD);
+  for (int i=0; i < nD; i++) D[i] = i;
+  for (int i=0; i < nM; i++) d->mapD2M[i] = -1;
+  copyD2MSparse(m, d, d->mapD2M, D);
+
+  // check that all indices are filled in
+  for (int i=0; i < nM; i++) {
+    if (d->mapD2M[i] < 0) {
+      mjERROR("unassigned index in mapD2M");
+    }
+  }
+
+  mj_freeStack(d);
+}
+
+
 //----------------------------------- mjData construction ------------------------------------------
 
 // set pointers into mjData buffer
@@ -1707,6 +1808,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
     makeDSparse(m, d);
     makeBSparse(m, d);
     checkDBSparse(m, d);
+    makeDmap(m, d);
   }
 
   // restore pluginstate and plugindata

src/engine/engine_print.c

@@ -1000,6 +1000,20 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
   }
   fprintf(fp, "\n\n");
 
+  // mapM2D
+  fprintf(fp, NAME_FORMAT, "mapM2D");
+  for (int i = 0; i < m->nD; i++) {
+    fprintf(fp, " %d", d->mapM2D[i]);
+  }
+  fprintf(fp, "\n\n");
+
+  // mapD2M
+  fprintf(fp, NAME_FORMAT, "mapD2M");
+  for (int i = 0; i < m->nM; i++) {
+    fprintf(fp, " %d", d->mapD2M[i]);
+  }
+  fprintf(fp, "\n\n");
+
   // B_rownnz
   fprintf(fp, NAME_FORMAT, "B_rownnz");
   for (int i = 0; i < m->nbody; i++) {

src/engine/engine_support.c

@@ -1285,66 +1285,6 @@ void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst) {
 }
 
 
-//-------------------------- sparse system matrix conversion ---------------------------------------
-
-// dst[D] = src[M], handle different sparsity representations
-void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) {
-  int nv = m->nv;
-  mj_markStack(d);
-
-  // init remaining
-  int* remaining = mj_stackAllocInt(d, nv);
-  mju_copyInt(remaining, d->D_rownnz, nv);
-
-  // copy data
-  for (int i = nv - 1; i >= 0; i--) {
-    // init at diagonal
-    int adr = m->dof_Madr[i];
-    remaining[i]--;
-    dst[d->D_rowadr[i] + remaining[i]] = src[adr];
-    adr++;
-
-    // process below diagonal
-    int j = i;
-    while ((j = m->dof_parentid[j]) >= 0) {
-      remaining[i]--;
-      dst[d->D_rowadr[i] + remaining[i]] = src[adr];
-
-      remaining[j]--;
-      dst[d->D_rowadr[j] + remaining[j]] = src[adr];
-
-      adr++;
-    }
-  }
-
-  mj_freeStack(d);
-}
-
-
-
-// dst[M] = src[D lower], handle different sparsity representations
-void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) {
-  int nv = m->nv;
-
-  // copy data
-  for (int i = nv - 1; i >= 0; i--) {
-    // find diagonal in qDeriv
-    int j = 0;
-    while (d->D_colind[d->D_rowadr[i] + j] < i) {
-      j++;
-    }
-
-    // copy
-    int adr = m->dof_Madr[i];
-    while (j >= 0) {
-      dst[adr] = src[d->D_rowadr[i] + j];
-      adr++;
-      j--;
-    }
-  }
-}
-
-
 
 //-------------------------- perturbations ---------------------------------------------------------
 

src/engine/engine_support.h

@@ -148,15 +148,6 @@ MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
 MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);
 
 
-//-------------------------- sparse system matrix conversion ---------------------------------------
-
-// dst[D] = src[M], handle different sparsity representations
-MJAPI void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src);
-
-// dst[M] = src[D lower], handle different sparsity representations
-MJAPI void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src);
-
-
 //-------------------------- perturbations ---------------------------------------------------------
 
 // apply Cartesian force and torque

unity/Runtime/Bindings/MjBindings.cs

@@ -4924,6 +4924,8 @@ public unsafe struct mjData_ {
   public int* D_rownnz;
   public int* D_rowadr;
   public int* D_colind;
+  public int* mapM2D;
+  public int* mapD2M;
   public int* B_rownnz;
   public int* B_rowadr;
   public int* B_colind;