Add dense LU factorization and solve functions.

PiperOrigin-RevId: 909290647
Change-Id: I77ae2352b20abf96ef7b48ae32b03a1f4098604e

src/engine/engine_util_solve.c

@@ -751,7 +751,90 @@ void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 }
 
 
-//------------------------------ LU factorization --------------------------------------------------
+//------------------------------ dense LU factorization --------------------------------------------
+
+// dense LU factorization with partial pivoting
+//   factorizes n x n row-major matrix A in-place into L and U
+//   L has unit diagonal (not stored), U has explicit diagonal
+//   pivot stores row permutation: row i of original = row pivot[i] of result
+//   return 1 if successful, 0 if singular (diagonal element < mjMINVAL)
+int mju_factorLU(mjtNum* restrict A, int n, int* pivot) {
+  for (int k=0; k < n; k++) {
+    // initialize pivot
+    pivot[k] = k;
+
+    // find pivot: max absolute value in column k, rows k..n-1
+    mjtNum maxval = mju_abs(A[k*n+k]);
+    int maxrow = k;
+    for (int i=k+1; i < n; i++) {
+      mjtNum val = mju_abs(A[i*n+k]);
+      if (val > maxval) {
+        maxval = val;
+        maxrow = i;
+      }
+    }
+
+    // check singularity
+    if (maxval < mjMINVAL) {
+      return 0;
+    }
+
+    // swap rows k and maxrow
+    if (maxrow != k) {
+      pivot[k] = maxrow;
+      for (int j=0; j < n; j++) {
+        mjtNum tmp = A[k*n+j];
+        A[k*n+j] = A[maxrow*n+j];
+        A[maxrow*n+j] = tmp;
+      }
+    }
+
+    // compute multipliers and update trailing submatrix
+    mjtNum diaginv = 1.0 / A[k*n+k];
+    for (int i=k+1; i < n; i++) {
+      A[i*n+k] *= diaginv;
+      mjtNum Aik = A[i*n+k];
+      for (int j=k+1; j < n; j++) {
+        A[i*n+j] -= Aik * A[k*n+j];
+      }
+    }
+  }
+
+  return 1;
+}
+
+
+// solve A*x = b given LU factorization of A, LU and pivot are output of mju_factorLU
+void mju_solveLU(mjtNum* restrict x, const mjtNum* LU, const mjtNum* b, const int* pivot, int n) {
+  // copy b into x
+  mju_copy(x, b, n);
+
+  // apply row permutation and forward substitution: solve L*y = P*b
+  for (int i=0; i < n; i++) {
+    // apply pivot swap
+    if (pivot[i] != i) {
+      mjtNum tmp = x[i];
+      x[i] = x[pivot[i]];
+      x[pivot[i]] = tmp;
+    }
+
+    // subtract known terms
+    for (int j=0; j < i; j++) {
+      x[i] -= LU[i*n+j] * x[j];
+    }
+  }
+
+  // back substitution: solve U*x = y
+  for (int i=n-1; i >= 0; i--) {
+    for (int j=i+1; j < n; j++) {
+      x[i] -= LU[i*n+j] * x[j];
+    }
+    x[i] /= LU[i*n+i];
+  }
+}
+
+
+//------------------------------ sparse LU factorization -------------------------------------------
 
 // sparse reverse-order LU factorization, no fill-in (assuming tree topology)
 //   result: LU = L + U; original = (U+I) * L; scratch size is n

src/engine/engine_util_solve.h

@@ -94,6 +94,16 @@ MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 // address of diagonal element i in band-dense matrix representation
 MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
 
+// dense LU factorization with partial pivoting
+//   factorizes n x n row-major matrix A in-place into L and U
+//   L has unit diagonal (not stored), U has explicit diagonal
+//   pivot stores row permutation: row i of original = row pivot[i] of result
+//   return 1 if successful, 0 if singular (diagonal element < mjMINVAL)
+MJAPI int mju_factorLU(mjtNum* A, int n, int* pivot);
+
+// solve A*x = b given LU factorization of A, LU and pivot are output of mju_factorLU
+MJAPI void mju_solveLU(mjtNum* x, const mjtNum* LU, const mjtNum* b, const int* pivot, int n);
+
 // sparse reverse-order LU factorization, assume tree topology (only dofs in index, if given)
 //  LU = L + U; original = (U+I) * L; scratch is size n
 void mju_factorLUSparse(mjtNum *LU, int n, int* scratch,

test/engine/engine_util_solve_test.cc

@@ -1041,5 +1041,147 @@ TEST_F(EngineUtilSolveTest, CholFactorSymbolicNumeric) {
   mj_deleteModel(model);
 }
 
+// ----------------------------- dense LU --------------------------------------
+
+using DenseLUTest = MujocoTest;
+
+// factor identity, solve recovers b exactly
+TEST_F(DenseLUTest, Identity) {
+  constexpr int n = 4;
+  mjtNum A[n*n] = {
+    1, 0, 0, 0,
+    0, 1, 0, 0,
+    0, 0, 1, 0,
+    0, 0, 0, 1,
+  };
+  int pivot[n];
+  mjtNum b[n] = {1, 2, 3, 4};
+  mjtNum x[n];
+
+  EXPECT_EQ(mju_factorLU(A, n, pivot), 1);
+  mju_solveLU(x, A, b, pivot, n);
+
+  for (int i = 0; i < n; i++) {
+    EXPECT_MJTNUM_EQ(x[i], b[i]);
+  }
+}
+
+// 3x3 with known solution
+TEST_F(DenseLUTest, SmallKnown) {
+  constexpr int n = 3;
+  // A = [2 1 1; 4 3 3; 8 7 9], b = [1; 1; 1]
+  // solution: x = [1; -1; 0]  (verified: A*x = [2-1; 4-3; 8-7] = [1;1;1])
+  mjtNum A[n*n] = {
+    2, 1, 1,
+    4, 3, 3,
+    8, 7, 9,
+  };
+  int pivot[n];
+  mjtNum b[n] = {1, 1, 1};
+  mjtNum x[n];
+
+  EXPECT_EQ(mju_factorLU(A, n, pivot), 1);
+  mju_solveLU(x, A, b, pivot, n);
+
+  mjtNum eps = MjTol(1e-14, 1e-6);
+  EXPECT_NEAR(x[0],  1, eps);
+  EXPECT_NEAR(x[1], -1, eps);
+  EXPECT_NEAR(x[2],  0, eps);
+}
+
+// random SPD matrices: compare LU solve against Cholesky solve
+TEST_F(DenseLUTest, RandomSPD) {
+  std::mt19937_64 rng;
+  rng.seed(7);
+  std::normal_distribution<double> dist(0, 1);
+
+  for (int n : {4, 8, 16}) {
+    vector<mjtNum> sqrtH(n * n);
+    vector<mjtNum> A(n * n);
+    vector<mjtNum> A_chol(n * n);
+    vector<mjtNum> b(n);
+    vector<mjtNum> x_lu(n);
+    vector<mjtNum> x_chol(n);
+    vector<int> pivot(n);
+
+    // generate random SPD matrix
+    for (int i = 0; i < n * n; i++) sqrtH[i] = dist(rng);
+    mju_mulMatTMat(A.data(), sqrtH.data(), sqrtH.data(), n, n, n);
+
+    // add diagonal regularizer
+    for (int i = 0; i < n; i++) A[i*n+i] += n;
+
+    // generate random rhs
+    for (int i = 0; i < n; i++) b[i] = dist(rng);
+
+    // solve with Cholesky
+    mju_copy(A_chol.data(), A.data(), n * n);
+    int rank = mju_cholFactor(A_chol.data(), n, 0);
+    EXPECT_EQ(rank, n);
+    mju_cholSolve(x_chol.data(), A_chol.data(), b.data(), n);
+
+    // solve with LU
+    int ok = mju_factorLU(A.data(), n, pivot.data());
+    EXPECT_EQ(ok, 1);
+    mju_solveLU(x_lu.data(), A.data(), b.data(), pivot.data(), n);
+
+    // compare
+    mjtNum eps = MjTol(1e-15, 1e-7);
+    EXPECT_THAT(AsVector(x_lu.data(), n),
+                Pointwise(MjNear(eps, eps),
+                          AsVector(x_chol.data(), n)));
+  }
+}
+
+// random non-symmetric matrices: verify A*x == b
+TEST_F(DenseLUTest, RandomGeneral) {
+  std::mt19937_64 rng;
+  rng.seed(42);
+  std::normal_distribution<double> dist(0, 1);
+
+  for (int n : {3, 5, 10, 20}) {
+    vector<mjtNum> A(n * n);
+    vector<mjtNum> A_orig(n * n);
+    vector<mjtNum> b(n);
+    vector<mjtNum> x(n);
+    vector<mjtNum> Ax(n);
+    vector<int> pivot(n);
+
+    // random non-symmetric matrix with diagonal dominance
+    for (int i = 0; i < n; i++) {
+      for (int j = 0; j < n; j++) {
+        A[i*n+j] = dist(rng);
+      }
+      A[i*n+i] += 2 * n;
+    }
+    mju_copy(A_orig.data(), A.data(), n * n);
+
+    // random rhs
+    for (int i = 0; i < n; i++) b[i] = dist(rng);
+
+    // factor and solve
+    int ok = mju_factorLU(A.data(), n, pivot.data());
+    EXPECT_EQ(ok, 1);
+    mju_solveLU(x.data(), A.data(), b.data(), pivot.data(), n);
+
+    // verify: A_orig * x == b
+    mju_mulMatVec(Ax.data(), A_orig.data(), x.data(), n, n);
+
+    mjtNum eps = MjTol(1e-14, 1e-5);
+    EXPECT_THAT(AsVector(Ax.data(), n),
+                Pointwise(MjNear(eps, eps), AsVector(b.data(), n)));
+  }
+}
+
+// near-singular matrix returns 0
+TEST_F(DenseLUTest, Singular) {
+  constexpr int n = 3;
+  // all zeros: maximally singular
+  mjtNum A[n*n] = {0};
+  int pivot[n];
+
+  EXPECT_EQ(mju_factorLU(A, n, pivot), 0);
+}
+
 }  // namespace
 }  // namespace mujoco