Add utility functions for mapping between sparse matrix patterns.
PiperOrigin-RevId: 795543910 Change-Id: I92be3afc3cd90020027737fffdba18c45e25f825
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
931036adda
commit
bbb70d98a4
@@ -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++) {
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
// Tests for engine/engine_util_solve.c.
|
||||
|
||||
#include <array>
|
||||
#include <cmath>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user