Add dense LU factorization and solve functions.

PiperOrigin-RevId: 909290647 Change-Id: I77ae2352b20abf96ef7b48ae32b03a1f4098604e
2026-05-02 13:26:14 -07:00
parent dbd451138c
commit 25751a7b98
3 changed files with 236 additions and 1 deletions
@@ -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