diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index 251e4924..d1543567 100644 --- a/src/engine/engine_util_solve.h +++ b/src/engine/engine_util_solve.h @@ -33,9 +33,8 @@ MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus); // sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank' // mat must be lower-triangular, have preallocated space for fill-in -int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, - int* rownnz, const int* rowadr, int* colind, - mjData* d); +MJAPI int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, + int* rownnz, const int* rowadr, int* colind, mjData* d); // precount row non-zeros of reverse-Cholesky factor L, return total MJAPI int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, @@ -47,9 +46,9 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int // sparse reverse-order Cholesky rank-one update: L'*L +/i x*x'; return rank // x is sparse, change in sparsity pattern of mat is not allowed -int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus, - const int* rownnz, const int* rowadr, const int* colind, - int x_nnz, int* x_ind, mjData* d); +MJAPI int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus, + const int* rownnz, const int* rowadr, const int* colind, + int x_nnz, int* x_ind, mjData* d); // band-dense Cholesky decomposition // returns minimum value in the factorized diagonal, or 0 if rank-deficient diff --git a/test/engine/engine_util_solve_test.cc b/test/engine/engine_util_solve_test.cc index aad5b421..a2b1b473 100644 --- a/test/engine/engine_util_solve_test.cc +++ b/test/engine/engine_util_solve_test.cc @@ -17,10 +17,11 @@ #include "src/engine/engine_util_solve.h" #include +#include #include #include -#include #include +#include #include #include @@ -37,6 +38,7 @@ using ::testing::Pointwise; using ::testing::DoubleNear; using ::testing::ElementsAre; using ::std::string; +using ::std::vector; using ::std::setw; using QCQP2Test = MujocoTest; @@ -723,5 +725,218 @@ TEST_F(EngineUtilSolveTest, MjuCholFactorNNZ) { mj_deleteModel(model); } +// Test for mju_cholUpdate: rank-one Cholesky update L*L' +/- x*x' +// Verifies that after applying the update, reconstructing H = L*L' produces +// the expected result H_original +/- x*x'. +TEST_F(EngineUtilSolveTest, MjuCholUpdate) { + std::mt19937_64 rng; + rng.seed(42); + std::normal_distribution dist(0, 1); + + for (int n : {4, 6, 8}) { + vector H(n * n); + vector H_expected(n * n); + vector L_dense(n * n); + vector sqrtH(n * n); + vector x(n); + vector x_copy(n); + vector H_reconstructed(n * n); + + for (int flg_plus : {0, 1}) { + // generate random lower-triangular matrix for sqrtH + mju_zero(sqrtH.data(), n * n); + for (int i = 0; i < n; i++) { + for (int j = 0; j <= i; j++) { + sqrtH[n * i + j] = dist(rng); + } + sqrtH[n * i + i] = mju_abs(sqrtH[n * i + i]) + n; + } + + // create SPD matrix H = sqrtH * sqrtH' (forward-order Cholesky: L*L') + for (int i = 0; i < n; i++) { + for (int j = 0; j < n; j++) { + mjtNum sum = 0; + for (int k = 0; k <= mju_min(i, j); k++) { + sum += sqrtH[n * i + k] * sqrtH[n * j + k]; + } + H[n * i + j] = sum; + } + } + + // generate random update vector x + for (int i = 0; i < n; i++) { + x[i] = dist(rng) * 0.5; + } + + // compute expected result: H_expected = H +/- x*x' + mju_copy(H_expected.data(), H.data(), n * n); + for (int i = 0; i < n; i++) { + for (int j = 0; j < n; j++) { + if (flg_plus) { + H_expected[i * n + j] += x[i] * x[j]; + } else { + H_expected[i * n + j] -= x[i] * x[j]; + } + } + } + + // test using dense Cholesky (forward-order: L*L') + mju_copy(L_dense.data(), H.data(), n * n); + + // dense factorization + int rank_factor = mju_cholFactor(L_dense.data(), n, 0); + EXPECT_EQ(rank_factor, n) << "Initial factorization failed"; + + // make copy of x for update (it's modified in-place) + mju_copy(x_copy.data(), x.data(), n); + + // apply dense rank-one update + int rank_update = + mju_cholUpdate(L_dense.data(), x_copy.data(), n, flg_plus); + EXPECT_EQ(rank_update, n) + << "Dense update rank loss for n=" << n << ", flg_plus=" << flg_plus; + + // zero out upper triangle (Cholesky only uses lower triangle) + for (int i = 0; i < n; i++) { + for (int j = i + 1; j < n; j++) { + L_dense[i * n + j] = 0; + } + } + + // reconstruct H from L*L' and compare to expected + mju_mulMatMatT(H_reconstructed.data(), L_dense.data(), L_dense.data(), n, + n, n); + + // compare + mjtNum eps = 1e-8; + for (int i = 0; i < n; i++) { + for (int j = 0; j < n; j++) { + EXPECT_NEAR(H_reconstructed[i * n + j], H_expected[i * n + j], eps) + << "Dense mismatch at (" << i << "," << j << ") for n=" << n + << ", flg_plus=" << flg_plus; + } + } + } + } +} + +// Test for mju_cholUpdateSparse: sparse rank-one Cholesky update L'*L +/- x*x' +// Uses sparse reverse Cholesky factorization and sparse rank-one update. +TEST_F(EngineUtilSolveTest, MjuCholUpdateSparse) { + mjModel* model = LoadModelFromString(""); + mjData* d = mj_makeData(model); + + std::mt19937_64 rng; + rng.seed(123); + std::normal_distribution dist(0, 1); + + for (int n : {4, 6, 8}) { + int max_nnz = n * (n + 1) / 2; + + vector H(n * n); + vector H_lower(n * n); + vector H_expected(n * n); + vector sqrtH(n * n); + vector L_sparse(max_nnz); + vector rownnz(n); + vector rowadr(n); + vector colind(max_nnz); + vector x(n); + vector x_sparse(n); + vector x_ind(n); + vector L_sparse_dense(n * n); + vector H_sparse_reconstructed(n * n); + + for (int flg_plus : {0, 1}) { + // generate random lower-triangular matrix for sqrtH + mju_zero(sqrtH.data(), n * n); + for (int i = 0; i < n; i++) { + for (int j = 0; j <= i; j++) { + sqrtH[n * i + j] = dist(rng); + } + sqrtH[n * i + i] = mju_abs(sqrtH[n * i + i]) + n; + } + + // create SPD matrix H = sqrtH * sqrtH' (forward-order: L*L') + for (int i = 0; i < n; i++) { + for (int j = 0; j < n; j++) { + mjtNum sum = 0; + for (int k = 0; k <= mju_min(i, j); k++) { + sum += sqrtH[n * i + k] * sqrtH[n * j + k]; + } + H[n * i + j] = sum; + } + } + + // generate random update vector x + for (int i = 0; i < n; i++) { + x[i] = dist(rng) * 0.5; + } + + // compute expected result: H_expected = H +/- x*x' + mju_copy(H_expected.data(), H.data(), n * n); + for (int i = 0; i < n; i++) { + for (int j = 0; j < n; j++) { + if (flg_plus) { + H_expected[i * n + j] += x[i] * x[j]; + } else { + H_expected[i * n + j] -= x[i] * x[j]; + } + } + } + + // copy H to H_lower and zero upper triangle (sparse expects lower only) + mju_copy(H_lower.data(), H.data(), n * n); + for (int i = 0; i < n; i++) { + for (int j = i + 1; j < n; j++) { + H_lower[i * n + j] = 0; + } + } + + // convert lower-triangular H to sparse format + mju_dense2sparse(L_sparse.data(), H_lower.data(), n, n, rownnz.data(), + rowadr.data(), colind.data(), max_nnz); + + // prepare sparse update vector (all elements, fully dense) + int x_nnz = n; + mju_copy(x_sparse.data(), x.data(), n); + for (int i = 0; i < n; i++) { + x_ind[i] = i; + } + + // sparse Cholesky factorization (reverse-order: L'*L) + mju_cholFactorSparse(L_sparse.data(), n, 0, rownnz.data(), rowadr.data(), + colind.data(), d); + + // apply sparse rank-one update + int rank_sparse = mju_cholUpdateSparse( + L_sparse.data(), x_sparse.data(), n, flg_plus, rownnz.data(), + rowadr.data(), colind.data(), x_nnz, x_ind.data(), d); + EXPECT_EQ(rank_sparse, n) + << "Sparse update rank loss for n=" << n << ", flg_plus=" << flg_plus; + + // reconstruct H from L'*L and compare to expected + mju_sparse2dense(L_sparse_dense.data(), L_sparse.data(), n, n, + rownnz.data(), rowadr.data(), colind.data()); + mju_mulMatTMat(H_sparse_reconstructed.data(), L_sparse_dense.data(), + L_sparse_dense.data(), n, n, n); + + // compare + mjtNum eps = 1e-8; + for (int i = 0; i < n; i++) { + for (int j = 0; j < n; j++) { + EXPECT_NEAR(H_sparse_reconstructed[i * n + j], H_expected[i * n + j], + eps) + << "Sparse mismatch at (" << i << "," << j << ") for n=" << n + << ", flg_plus=" << flg_plus; + } + } + } + } + + mj_deleteData(d); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco