Move mju_combineSparseCount into engine_util_sparse.
PiperOrigin-RevId: 561025741 Change-Id: I0bdc48554928fc4545a38a8cea3bdeb593bd7ecd
This commit is contained in:
committed by
Copybara-Service
parent
df1b3515c7
commit
ea5e00cad8
@@ -1351,33 +1351,6 @@ void mj_makeImpedance(const mjModel* m, mjData* d) {
|
|||||||
|
|
||||||
//------------------------------------- constraint counting ----------------------------------------
|
//------------------------------------- constraint counting ----------------------------------------
|
||||||
|
|
||||||
// count the number of non-zeros in the sum of two sparse vectors
|
|
||||||
int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) {
|
|
||||||
int a = 0, b = 0, c_nnz = 0;
|
|
||||||
|
|
||||||
// count c_nnz: nonzero indices common to both a and b
|
|
||||||
while (a < a_nnz && b < b_nnz) {
|
|
||||||
// common index, increment everything
|
|
||||||
if (a_ind[a] == b_ind[b]) {
|
|
||||||
c_nnz++;
|
|
||||||
a++;
|
|
||||||
b++;
|
|
||||||
}
|
|
||||||
|
|
||||||
// update smallest index
|
|
||||||
else if (a_ind[a] < b_ind[b]) {
|
|
||||||
a++;
|
|
||||||
} else {
|
|
||||||
b++;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// union minus the intersection
|
|
||||||
return a_nnz + b_nnz - c_nnz;
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
// count the non-zero columns in the Jacobian difference of two bodies
|
// count the non-zero columns in the Jacobian difference of two bodies
|
||||||
static int mj_jacDifPairCount(const mjModel* m, int* chain,
|
static int mj_jacDifPairCount(const mjModel* m, int* chain,
|
||||||
int b1, int b2, int issparse) {
|
int b1, int b2, int issparse) {
|
||||||
|
|||||||
@@ -93,13 +93,6 @@ void mj_diagApprox(const mjModel* m, mjData* d);
|
|||||||
// compute efc_R, efc_D, efc_KDIP, adjust diagApprox
|
// compute efc_R, efc_D, efc_KDIP, adjust diagApprox
|
||||||
void mj_makeImpedance(const mjModel* m, mjData* d);
|
void mj_makeImpedance(const mjModel* m, mjData* d);
|
||||||
|
|
||||||
|
|
||||||
//------------------------- constraint counting
|
|
||||||
|
|
||||||
// count the number of non-zeros in the sum of two sparse vectors
|
|
||||||
MJAPI int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind);
|
|
||||||
|
|
||||||
|
|
||||||
//---------------------------- top-level API for constraint construction ---------------------------
|
//---------------------------- top-level API for constraint construction ---------------------------
|
||||||
|
|
||||||
// main driver: call all functions above
|
// main driver: call all functions above
|
||||||
|
|||||||
@@ -207,6 +207,33 @@ static int mju_compare(const int* vec1, const int* vec2, int n) {
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
// count the number of non-zeros in the sum of two sparse vectors
|
||||||
|
int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) {
|
||||||
|
int a = 0, b = 0, c_nnz = 0;
|
||||||
|
|
||||||
|
// count c_nnz: nonzero indices common to both a and b
|
||||||
|
while (a < a_nnz && b < b_nnz) {
|
||||||
|
// common index, increment everything
|
||||||
|
if (a_ind[a] == b_ind[b]) {
|
||||||
|
c_nnz++;
|
||||||
|
a++;
|
||||||
|
b++;
|
||||||
|
}
|
||||||
|
|
||||||
|
// update smallest index
|
||||||
|
else if (a_ind[a] < b_ind[b]) {
|
||||||
|
a++;
|
||||||
|
} else {
|
||||||
|
b++;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// union minus the intersection
|
||||||
|
return a_nnz + b_nnz - c_nnz;
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
|
// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
|
||||||
int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
|
int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
|
||||||
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
|
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
|
||||||
|
|||||||
@@ -51,6 +51,9 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec
|
|||||||
MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc,
|
MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc,
|
||||||
int* rownnz, int* rowadr, int* colind);
|
int* rownnz, int* rowadr, int* colind);
|
||||||
|
|
||||||
|
// count the number of non-zeros in the sum of two sparse vectors
|
||||||
|
MJAPI int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind);
|
||||||
|
|
||||||
// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
|
// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
|
||||||
int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
|
int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
|
||||||
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
|
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
|
||||||
|
|||||||
@@ -198,77 +198,5 @@ TEST_F(CoreConstraintTest, JacobianPreAllocate) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
TEST_F(CoreConstraintTest, CombineSparseCount) {
|
|
||||||
{
|
|
||||||
std::array a_ind{0, 1};
|
|
||||||
std::array b_ind{2};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{2};
|
|
||||||
std::array b_ind{0, 1};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{0, 1};
|
|
||||||
std::array b_ind{2, 3, 4};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{5, 6};
|
|
||||||
std::array b_ind{1, 3, 8};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{1, 2, 3};
|
|
||||||
std::array b_ind{0, 4};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{1, 4};
|
|
||||||
std::array b_ind{2, 3};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{0, 1, 3};
|
|
||||||
std::array b_ind{0, 3, 4};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{1, 3, 5, 6};
|
|
||||||
std::array b_ind{1, 3, 5, 6};
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(
|
|
||||||
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4);
|
|
||||||
}
|
|
||||||
|
|
||||||
EXPECT_EQ(mju_combineSparseCount(0, 0, nullptr, nullptr), 0);
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array b_ind{1, 2};
|
|
||||||
EXPECT_EQ(
|
|
||||||
mju_combineSparseCount(0, b_ind.size(), nullptr, b_ind.data()), 2);
|
|
||||||
}
|
|
||||||
|
|
||||||
{
|
|
||||||
std::array a_ind{0};
|
|
||||||
EXPECT_EQ(
|
|
||||||
mju_combineSparseCount(a_ind.size(), 0, a_ind.data(), nullptr), 1);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
} // namespace mujoco
|
} // namespace mujoco
|
||||||
|
|||||||
@@ -14,6 +14,8 @@
|
|||||||
|
|
||||||
// Tests for engine/engine_util_sparse.c
|
// Tests for engine/engine_util_sparse.c
|
||||||
|
|
||||||
|
#include <array>
|
||||||
|
|
||||||
#include "src/engine/engine_util_sparse.h"
|
#include "src/engine/engine_util_sparse.h"
|
||||||
|
|
||||||
#include <gmock/gmock.h>
|
#include <gmock/gmock.h>
|
||||||
@@ -43,6 +45,78 @@ TEST_F(EngineUtilSparseTest, MjuDot) {
|
|||||||
EXPECT_EQ(mju_dotSparse(a, b, 7, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7);
|
EXPECT_EQ(mju_dotSparse(a, b, 7, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(EngineUtilSparseTest, CombineSparseCount) {
|
||||||
|
{
|
||||||
|
std::array a_ind{0, 1};
|
||||||
|
std::array b_ind{2};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{2};
|
||||||
|
std::array b_ind{0, 1};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{0, 1};
|
||||||
|
std::array b_ind{2, 3, 4};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{5, 6};
|
||||||
|
std::array b_ind{1, 3, 8};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{1, 2, 3};
|
||||||
|
std::array b_ind{0, 4};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{1, 4};
|
||||||
|
std::array b_ind{2, 3};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{0, 1, 3};
|
||||||
|
std::array b_ind{0, 3, 4};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{1, 3, 5, 6};
|
||||||
|
std::array b_ind{1, 3, 5, 6};
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(
|
||||||
|
a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4);
|
||||||
|
}
|
||||||
|
|
||||||
|
EXPECT_EQ(mju_combineSparseCount(0, 0, nullptr, nullptr), 0);
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array b_ind{1, 2};
|
||||||
|
EXPECT_EQ(
|
||||||
|
mju_combineSparseCount(0, b_ind.size(), nullptr, b_ind.data()), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
{
|
||||||
|
std::array a_ind{0};
|
||||||
|
EXPECT_EQ(
|
||||||
|
mju_combineSparseCount(a_ind.size(), 0, a_ind.data(), nullptr), 1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
TEST_F(EngineUtilSparseTest, MjuTranspose3by3) {
|
TEST_F(EngineUtilSparseTest, MjuTranspose3by3) {
|
||||||
// 1 2 0 1 0 0
|
// 1 2 0 1 0 0
|
||||||
// 0 1 0 --> 2 1 3
|
// 0 1 0 --> 2 1 3
|
||||||
|
|||||||
Reference in New Issue
Block a user