chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add dense LU factorization and solve functions.

Browse Source 
PiperOrigin-RevId: 909290647
Change-Id: I77ae2352b20abf96ef7b48ae32b03a1f4098604e
This commit is contained in:
Yuval Tassa
2026-05-02 13:26:14 -07:00
committed by Copybara-Service
parent dbd451138c
commit 25751a7b98
3 changed files with 236 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_solve.c
+84 -1
View File
@@ -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
+10
View File
@@ -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
+142
View File
@@ -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