From 25751a7b98099b1b9841c4ee3acb5d60b6d4d315 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sat, 2 May 2026 13:26:14 -0700 Subject: [PATCH] Add dense LU factorization and solve functions. PiperOrigin-RevId: 909290647 Change-Id: I77ae2352b20abf96ef7b48ae32b03a1f4098604e --- src/engine/engine_util_solve.c | 85 ++++++++++++++- src/engine/engine_util_solve.h | 10 ++ test/engine/engine_util_solve_test.cc | 142 ++++++++++++++++++++++++++ 3 files changed, 236 insertions(+), 1 deletion(-) diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 1f8b839b..a6e5541c 100644 --- a/src/engine/engine_util_solve.c +++ b/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 diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index a50ab4ee..8cdb8ea9 100644 --- a/src/engine/engine_util_solve.h +++ b/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, diff --git a/test/engine/engine_util_solve_test.cc b/test/engine/engine_util_solve_test.cc index 232297a3..d80abde1 100644 --- a/test/engine/engine_util_solve_test.cc +++ b/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 dist(0, 1); + + for (int n : {4, 8, 16}) { + vector sqrtH(n * n); + vector A(n * n); + vector A_chol(n * n); + vector b(n); + vector x_lu(n); + vector x_chol(n); + vector 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 dist(0, 1); + + for (int n : {3, 5, 10, 20}) { + vector A(n * n); + vector A_orig(n * n); + vector b(n); + vector x(n); + vector Ax(n); + vector 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