From bbb70d98a431343a87dcfcf03bdd39153d628be1 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 15 Aug 2025 11:21:43 -0700 Subject: [PATCH] Add utility functions for mapping between sparse matrix patterns. PiperOrigin-RevId: 795543910 Change-Id: I92be3afc3cd90020027737fffdba18c45e25f825 --- src/engine/engine_util_misc.c | 101 +++++++++++++++++++ src/engine/engine_util_misc.h | 15 +++ test/engine/engine_util_misc_test.cc | 140 ++++++++++++++++++++++++++- 3 files changed, 253 insertions(+), 3 deletions(-) diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index f959ce71..2e99cb64 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -1450,6 +1450,16 @@ void mju_gather(mjtNum* restrict res, const mjtNum* restrict vec, const int* res +// masked gather (set to 0 at negative indices) +void mju_gatherMasked(mjtNum* restrict res, const mjtNum* restrict vec, + const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[i] = ind[i] >= 0 ? vec[ind[i]] : 0; + } +} + + + // scatter void mju_scatter(mjtNum* restrict res, const mjtNum* restrict vec, const int* restrict ind, int n) { for (int i=0; i < n; i++) { @@ -1477,6 +1487,97 @@ void mju_scatterInt(int* restrict res, const int* restrict vec, const int* restr +// build gather indices mapping src to res, assumes pattern(res) \subseteq pattern(src) +void mju_sparseMap(int* map, int nr, + const int* res_rowadr, const int* res_rownnz, const int* res_colind, + const int* src_rowadr, const int* src_rownnz, const int* src_colind) { + for (int i = 0; i < nr; i++) { + int res_cursor = res_rowadr[i]; + int res_end = res_cursor + res_rownnz[i]; + int src_cursor = src_rowadr[i]; + int src_end = src_cursor + src_rownnz[i]; + + while (res_cursor < res_end) { + int res_col = res_colind[res_cursor]; + while (src_cursor < src_end && src_colind[src_cursor] < res_col) { + src_cursor++; + } + + // found match, set index and advance cursors + map[res_cursor++] = src_cursor++; + } + } +} + + + +// build masked-gather map to copy a lower-triangular src into symmetric res +// `cursor` is a preallocated buffer of size `nr` +void mju_lower2SymMap(int* map, int nr, + const int* res_rowadr, const int* res_rownnz, const int* res_colind, + const int* src_rowadr, const int* src_rownnz, const int* src_colind, + int* cursor) { + if (!nr) return; + + // default all map entries to "no source" + int nnz = res_rowadr[nr-1] + res_rownnz[nr-1]; + for (int i = 0; i < nnz; i++) { + map[i] = -1; + } + + // initialize per-row cursor + for (int i = 0; i < nr; i++) { + cursor[i] = res_rowadr[i]; + } + + // sweep src rows; for each lower (i,j) set res(i,j) and res(j,i) + for (int i = 0; i < nr; i++) { + int src_start = src_rowadr[i]; + int src_end = src_start + src_rownnz[i]; + + // sweep src row + for (int k = src_start; k < src_end; k++) { + int j = src_colind[k]; + if (j > i) break; // use only lower triangle of src + + // --- lower triangle: res(i, j) + int res_start = res_rowadr[i]; + int res_end = res_start + res_rownnz[i]; + int c = cursor[i]; + + // increment c until there is a match + while (c < res_end && res_colind[c] < j) c++; + + // found match, set index, advance and save cursor + if (c < res_end && res_colind[c] == j) { + map[c] = k; + c++; + } + cursor[i] = c; + + + // --- upper mirror: res(j, i) + if (j != i) { + res_start = res_rowadr[j]; + res_end = res_start + res_rownnz[j]; + c = cursor[j]; + + // increment c until there is a match + while (c < res_end && res_colind[c] < i) c++; + + // found match, set index and advance and save cursor + if (c < res_end && res_colind[c] == i) { + map[c] = k; + c++; + } + cursor[j] = c; + } + } + } +} + + + // insertion sort, increasing order void mju_insertionSort(mjtNum* list, int n) { for (int i=1; i < n; i++) { diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h index 3b25ddfb..c5e35e2c 100644 --- a/src/engine/engine_util_misc.h +++ b/src/engine/engine_util_misc.h @@ -166,6 +166,9 @@ MJAPI void mju_n2d(double* res, const mjtNum* vec, int n); // gather mjtNums MJAPI void mju_gather(mjtNum* res, const mjtNum* vec, const int* ind, int n); +// gather mjtNums, set to 0 at negative indices +MJAPI void mju_gatherMasked(mjtNum* res, const mjtNum* vec, const int* ind, int n); + // scatter mjtNums MJAPI void mju_scatter(mjtNum* res, const mjtNum* vec, const int* ind, int n); @@ -175,6 +178,18 @@ MJAPI void mju_gatherInt(int* res, const int* vec, const int* ind, int n); // scatter integers MJAPI void mju_scatterInt(int* res, const int* vec, const int* ind, int n); +// build gather indices mapping src to res, assumes pattern(res) \subseteq pattern(src) +MJAPI void mju_sparseMap(int* map, int nr, + const int* res_rowadr, const int* res_rownnz, const int* res_colind, + const int* src_rowadr, const int* src_rownnz, const int* src_colind); + +// build masked-gather map to copy a lower-triangular src into symmetric res +// `cursor` is a preallocated buffer of size `nr` +MJAPI void mju_lower2SymMap(int* map, int nr, + const int* res_rowadr, const int* res_rownnz, const int* res_colind, + const int* src_rowadr, const int* src_rownnz, const int* src_colind, + int* cursor); + // insertion sort, increasing order MJAPI void mju_insertionSort(mjtNum* list, int n); diff --git a/test/engine/engine_util_misc_test.cc b/test/engine/engine_util_misc_test.cc index 79a90e08..c4b4c4dc 100644 --- a/test/engine/engine_util_misc_test.cc +++ b/test/engine/engine_util_misc_test.cc @@ -15,6 +15,7 @@ // Tests for engine/engine_util_solve.c. #include +#include #include #include #include @@ -30,11 +31,11 @@ namespace mujoco { namespace { using ::testing::DoubleNear; +using ::testing::ElementsAre; +using ::testing::ElementsAreArray; using ::testing::HasSubstr; -using ::testing::Ne; using ::testing::Pointwise; using ::testing::StrEq; -using ::testing::ElementsAreArray; using UtilMiscTest = MujocoTest; @@ -211,6 +212,140 @@ TEST_F(UtilMiscTest, MuscleGainLength) { EXPECT_EQ(mju_muscleGainLength(2.0, lmin, lmax), 0); } +TEST_F(UtilMiscTest, MjuSparseMap) { + // nr = 3 + // src = [[1, 2, 0], + // [0, 3, 4], + // [5, 0, 6]] + constexpr int nr = 3; + const mjtNum mat_src[] = {1, 2, 3, 4, 5, 6}; + const int rownnz_src[] = {2, 2, 2}; + const int rowadr_src[] = {0, 2, 4}; + const int colind_src[] = {0, 1, 1, 2, 0, 2}; + + // res = [[1, 0, 0], + // [0, 3, 0], + // [5, 0, 6]] + constexpr int nnz_res = 4; + const int rownnz_res[] = {1, 1, 2}; + const int rowadr_res[] = {0, 1, 2}; + const int colind_res[] = {0, 1, 0, 2}; + + int map[nnz_res]; + mju_sparseMap(map, nr, rowadr_res, rownnz_res, colind_res, rowadr_src, + rownnz_src, colind_src); + + // Expected map: + // res[0] (1 at 0,0) -> src[0] (1 at 0,0) => map[0] = 0 + // res[1] (3 at 1,1) -> src[2] (3 at 1,1) => map[1] = 2 + // res[2] (5 at 2,0) -> src[4] (5 at 2,0) => map[2] = 4 + // res[3] (6 at 2,2) -> src[5] (6 at 2,2) => map[3] = 5 + EXPECT_THAT(map, ElementsAre(0, 2, 4, 5)); + + // Verify the map by checking values + mjtNum mat_res_gathered[nnz_res]; + mju_gather(mat_res_gathered, mat_src, map, nnz_res); + EXPECT_THAT(AsVector(mat_res_gathered, nnz_res), + ElementsAre(1, 3, 5, 6)); +} + + +TEST_F(UtilMiscTest, MjuSparseLower2SymMap) { + // nr = 3 + // src = [[1, 0, 0], + // [2, 3, 0], + // [4, 5, 6]] + constexpr int nr = 3; + const mjtNum mat_src[] = {1, 2, 3, 4, 5, 6}; + const int rownnz_src[] = {1, 2, 3}; + const int rowadr_src[] = {0, 1, 3}; + const int colind_src[] = {0, 0, 1, 0, 1, 2}; + + // res = [[*, *, *], + // [*, *, *], + // [*, *, *]] (dense symmetric) + constexpr int res_nnz = 9; + const int rownnz_res[] = {3, 3, 3}; + const int rowadr_res[] = {0, 3, 6}; + const int colind_res[] = {0, 1, 2, 0, 1, 2, 0, 1, 2}; + + int map[res_nnz]; + int cursor[nr]; + + mju_lower2SymMap(map, nr, rowadr_res, rownnz_res, colind_res, + rowadr_src, rownnz_src, colind_src, cursor); + + // Expected map: + // res(0,0) -> src(0,0) (k=0) => map[0] = 0 + // res(0,1) -> src(1,0) (k=1) => map[1] = 1 + // res(0,2) -> src(2,0) (k=3) => map[2] = 3 + // res(1,0) -> src(1,0) (k=1) => map[3] = 1 + // res(1,1) -> src(1,1) (k=2) => map[4] = 2 + // res(1,2) -> src(2,1) (k=4) => map[5] = 4 + // res(2,0) -> src(2,0) (k=3) => map[6] = 3 + // res(2,1) -> src(2,1) (k=4) => map[7] = 4 + // res(2,2) -> src(2,2) (k=5) => map[8] = 5 + EXPECT_THAT(map, ElementsAre(0, 1, 3, 1, 2, 4, 3, 4, 5)); + + // Verify the map by checking values + mjtNum mat_res[res_nnz]; + mju_gatherMasked(mat_res, mat_src, map, res_nnz); + + EXPECT_THAT(AsVector(mat_res, res_nnz), + ElementsAre(1, 2, 4, 2, 3, 5, 4, 5, 6)); +} + +TEST_F(UtilMiscTest, MjuSparseLower2SymMapPartial) { + // nr = 3 + // src = [[1, 0, 0], + // [2, 3, 0], + // [0, 0, 6]] + constexpr int nr = 3; + const mjtNum mat_src[] = {1, 2, 3, 6}; + const int rownnz_src[] = {1, 2, 1}; + const int rowadr_src[] = {0, 1, 3}; + const int colind_src[] = {0, 0, 1, 2}; + + // res with a sparse symmetric pattern + // res = [[*, *, *], + // [*, *, 0], + // [*, 0, *]] + constexpr int res_nnz = 7; + const int rownnz_res[] = {3, 2, 2}; + const int rowadr_res[] = {0, 3, 5}; + const int colind_res[] = {0, 1, 2, 0, 1, 0, 2}; + + int map[res_nnz]; + int cursor[nr]; + + mju_lower2SymMap(map, nr, rowadr_res, rownnz_res, colind_res, + rowadr_src, rownnz_src, colind_src, cursor); + + // Expected map for the non-zeros in res: + // res(0,0) -> src(0,0) (k=0) => map[0] = 0 + // res(0,1) -> src(1,0) (k=1) => map[1] = 1 + // res(0,2) -> Unmapped => map[2] = -1 + // res(1,0) -> src(1,0) (k=1) => map[3] = 1 + // res(1,1) -> src(1,1) (k=2) => map[4] = 2 + // res(2,0) -> Unmapped => map[5] = -1 + // res(2,2) -> src(2,2) (k=3) => map[6] = 3 + EXPECT_THAT(map, ElementsAre(0, 1, -1, 1, 2, -1, 3)); + + // Verify the map by checking values + mjtNum mat_res[res_nnz]; + mju_gatherMasked(mat_res, mat_src, map, res_nnz); + + // Expected res values based on map: + // mat_res[0] = mat_src[0] = 1 + // mat_res[1] = mat_src[1] = 2 + // mat_res[2] = 0 (unmapped) + // mat_res[3] = mat_src[1] = 2 + // mat_res[4] = mat_src[2] = 3 + // mat_res[5] = 0 (unmapped) + // mat_res[6] = mat_src[3] = 6 + EXPECT_THAT(AsVector(mat_res, res_nnz), ElementsAre(1, 2, 0, 2, 3, 0, 6)); +} + // --------------------------------- Interpolation ----------------------------- using InterpolationTest = MujocoTest; @@ -334,7 +469,6 @@ TEST_F(Base64Test, mju_encodeBase64_align1) { EXPECT_THAT(buffer.data(), StrEq("QUI=")); EXPECT_THAT(n, std::strlen(buffer.data()) + 1); EXPECT_THAT(n, buffer.size()); - } TEST_F(Base64Test, mju_encodeBase64_align2) {