Add CSR implementation of mj_factorI

PiperOrigin-RevId: 712498431
Change-Id: I13b52e53482ed97da8788875d4d95e2beb5ca7c1

src/engine/engine_core_smooth.c

@@ -1440,10 +1440,9 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu
     }
   }
 
-  // compute 1/diag(D), 1/sqrt(diag(D))
+  // compute 1/diag(D)
   for (int i=0; i < nv; i++) {
-    mjtNum qLDi = qLD[dof_Madr[i]];
-    qLDiagInv[i] = 1.0/qLDi;
+    qLDiagInv[i] = 1.0 / qLD[dof_Madr[i]];
   }
 }
 
@@ -1458,6 +1457,40 @@ void mj_factorM(const mjModel* m, mjData* d) {
 
 
 
+// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
+//  like mj_factorI, but using CSR representation
+void mj_factorIs(mjtNum* mat, mjtNum* diaginv, int nv,
+                 const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
+  // backward loop over rows
+  for (int k=nv-1; k >= 0; k--) {
+    // get row k's address, diagonal index, inverse diagonal value
+    int rowadr_k = rowadr[k];
+    int diag_k = rowadr_k + rownnz[k] - 1;
+    mjtNum invD = 1 / mat[diag_k];
+    if (diaginv) diaginv[k] = invD;
+
+    // skip if simple
+    if (diagnum[k]) {
+      continue;
+    }
+
+    // update triangle above row k, inclusive
+    for (int adr=diag_k - 1; adr >= rowadr_k; adr--) {
+      // tmp = L(k, i) / L(k, k)
+      mjtNum tmp = mat[adr] * invD;
+
+      // update row i < k:  L(i, 0..i) -= L(i, 0..i) * L(k, i) / L(k, k)
+      int i = colind[adr];
+      mju_addToScl(mat + rowadr[i], mat + rowadr_k, -tmp, rownnz[i]);
+
+      // update ith element of row k:  L(k, i) /= L(k, k)
+      mat[adr] = tmp;
+    }
+  }
+}
+
+
+
 // in-place sparse backsubstitution:  x = inv(L'*D*L)*x
 //  L is in lower triangle of qLD; D is on diagonal of qLD
 //  handle n vectors at once
@@ -1575,8 +1608,7 @@ void mj_solveLD(const mjModel* m, mjtNum* restrict x, int n,
 // in-place sparse backsubstitution:  x = inv(L'*D*L)*x
 //  like mj_solveLD, but using the CSR representation of L
 void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
-                 const int* rownnz, const int* rowadr, const int* diagind, const int* diagnum,
-                 const int* colind) {
+                 const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
   // x <- L^-T x
   for (int i=nv-1; i > 0; i--) {
     // skip diagonal (simple) rows, exploit sparsity of input vector
@@ -1584,7 +1616,7 @@ void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv
       continue;
     }
 
-    int d = diagind[i];
+    int d = rownnz[i] - 1;
     int adr_i = rowadr[i];
     mjtNum x_i = x[i];
     for (int j=0; j < d; j++) {
@@ -1607,7 +1639,7 @@ void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv
     }
 
     int adr = rowadr[i];
-    x[i] -= mju_dotSparse(qLDs+adr, x, diagind[i], colind+adr, /*flg_unc1=*/0);
+    x[i] -= mju_dotSparse(qLDs+adr, x, rownnz[i] - 1, colind+adr, /*flg_unc1=*/0);
   }
 }
 

src/engine/engine_core_smooth.h

@@ -51,6 +51,11 @@ MJAPI void mj_crb(const mjModel* m, mjData* d);
 // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
 MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv);
 
+// sparse L'*D*L factorizaton of inertia-like matrix
+//  like mj_factorI, but using CSR representation
+MJAPI void mj_factorIs(mjtNum* mat, mjtNum* diaginv, int nv,
+                       const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
+
 // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
 MJAPI void mj_factorM(const mjModel* m, mjData* d);
 
@@ -61,8 +66,7 @@ MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n,
 // in-place sparse backsubstitution:  x = inv(L'*D*L)*x
 //  like mj_solveLD, but using the CSR representation of L
 MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
-                       const int* rownnz, const int* rowadr, const int* diagind, const int* diagnum,
-                       const int* colind);
+                       const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
 
 // sparse backsubstitution:  x = inv(L'*D*L)*y, use factorization in d
 MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);