Provide improved mju_sqrMatTDSparse implementation that doesn't require dense memory allocation for sparse matrices.
PiperOrigin-RevId: 516812783 Change-Id: Ieb43337831d8d3b2c7f18a7facd8e0a0f2b0eff5
This commit is contained in:
Kyle Bayes
committed by
Copybara-Service
Copybara-Service
parent
fe18e58ad2
commit
056e849273
@@ -28,6 +28,7 @@ namespace {
|
||||
|
||||
using CombineFuncPtr = decltype(&mju_combineSparse);
|
||||
using TransposeFuncPtr = decltype(&mju_transposeSparse);
|
||||
using SqrMatTDFuncPtr = decltype(&mju_sqrMatTDSparse);
|
||||
|
||||
// number of steps to roll out before benchmarking
|
||||
static const int kNumWarmupSteps = 500;
|
||||
@@ -39,6 +40,117 @@ std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
|
||||
|
||||
// ----------------------------- old functions --------------------------------
|
||||
|
||||
void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
|
||||
mjtNum* res, const mjtNum* mat, const mjtNum* matT, const mjtNum* diag,
|
||||
int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind,
|
||||
const int* rownnz, const int* rowadr, const int* colind,
|
||||
const int* rowsuper, const int* rownnzT, const int* rowadrT,
|
||||
const int* colindT, const int* rowsuperT, mjData* d) {
|
||||
mjMARKSTACK;
|
||||
int* chain = (int*)mj_stackAlloc(d, 2 * nc);
|
||||
mjtNum* buffer = mj_stackAlloc(d, nc);
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
res_rowadr[r] = r * nc;
|
||||
}
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
if (rowsuperT && r > 0 && rowsuperT[r - 1] > 0) {
|
||||
res_rownnz[r] = res_rownnz[r - 1];
|
||||
memcpy(res_colind + res_rowadr[r], res_colind + res_rowadr[r - 1],
|
||||
res_rownnz[r] * sizeof(int));
|
||||
|
||||
if (rownnzT[r]) {
|
||||
res_colind[res_rowadr[r] + res_rownnz[r]] = r;
|
||||
res_rownnz[r]++;
|
||||
}
|
||||
} else {
|
||||
int nchain = 0;
|
||||
int inew = 0, iold = nc;
|
||||
int lastadded = -1;
|
||||
for (int i = 0; i < rownnzT[r]; i++) {
|
||||
int c = colindT[rowadrT[r] + i];
|
||||
if (rowsuper && lastadded >= 0 &&
|
||||
(c - lastadded) <= rowsuper[lastadded]) {
|
||||
continue;
|
||||
} else {
|
||||
lastadded = c;
|
||||
}
|
||||
|
||||
int adr = inew;
|
||||
inew = iold;
|
||||
iold = adr;
|
||||
|
||||
int nnewchain = 0;
|
||||
adr = 0;
|
||||
int end = rowadr[c] + rownnz[c];
|
||||
for (int adr1 = rowadr[c]; adr1 < end; adr1++) {
|
||||
int col_mat = colind[adr1];
|
||||
while (adr < nchain && chain[iold + adr] < col_mat &&
|
||||
chain[iold + adr] <= r) {
|
||||
chain[inew + nnewchain++] = chain[iold + adr++];
|
||||
}
|
||||
|
||||
if (col_mat > r) {
|
||||
break;
|
||||
}
|
||||
|
||||
if (adr < nchain && chain[iold + adr] == col_mat) {
|
||||
adr++;
|
||||
}
|
||||
chain[inew + nnewchain++] = col_mat;
|
||||
}
|
||||
|
||||
while (adr < nchain && chain[iold + adr] <= r) {
|
||||
chain[inew + nnewchain++] = chain[iold + adr++];
|
||||
}
|
||||
nchain = nnewchain;
|
||||
}
|
||||
res_rownnz[r] = nchain;
|
||||
if (nchain) {
|
||||
memcpy(res_colind + res_rowadr[r], chain + inew, nchain * sizeof(int));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
int adr = res_rowadr[r];
|
||||
for (int i = 0; i < res_rownnz[r]; i++) {
|
||||
buffer[res_colind[adr + i]] = 0;
|
||||
}
|
||||
for (int i = 0; i < rownnzT[r]; i++) {
|
||||
int c = colindT[rowadrT[r] + i];
|
||||
mjtNum matTrc = matT[rowadrT[r] + i];
|
||||
if (diag) {
|
||||
matTrc *= diag[c];
|
||||
}
|
||||
|
||||
int end = rowadr[c] + rownnz[c];
|
||||
for (int adr = rowadr[c]; adr < end; adr++) {
|
||||
int adr1;
|
||||
if ((adr1 = colind[adr]) > r) {
|
||||
break;
|
||||
}
|
||||
buffer[adr1] += matTrc * mat[adr];
|
||||
}
|
||||
}
|
||||
adr = res_rowadr[r];
|
||||
for (int i = 0; i < res_rownnz[r]; i++) {
|
||||
res[adr + i] = buffer[res_colind[adr + i]];
|
||||
}
|
||||
}
|
||||
for (int r = 1; r < nc; r++) {
|
||||
int end = res_rowadr[r] + res_rownnz[r] - 1;
|
||||
for (int adr = res_rowadr[r]; adr < end; adr++) {
|
||||
int adr1 = res_rowadr[res_colind[adr]] + res_rownnz[res_colind[adr]]++;
|
||||
res[adr1] = res[adr];
|
||||
res_colind[adr1] = r;
|
||||
}
|
||||
}
|
||||
|
||||
mjFREESTACK;
|
||||
}
|
||||
|
||||
// transpose sparse matrix (uncompressed)
|
||||
void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
|
||||
mjtNum* res, const mjtNum* mat, int nr, int nc, int* res_rownnz,
|
||||
@@ -426,8 +538,79 @@ BM_transposeSparse_old(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse(state, &transposeSparse_baseline);
|
||||
}
|
||||
|
||||
BENCHMARK(BM_transposeSparse_old);
|
||||
|
||||
static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
|
||||
static mjModel* m = LoadModelFromPath("humanoid100/humanoid100.xml");
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
// force use of sparse matrices
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
|
||||
// warm-up rollout to get a typical state
|
||||
while (d->time < 2) {
|
||||
mj_step(m, d);
|
||||
}
|
||||
|
||||
// allocate
|
||||
mjMARKSTACK;
|
||||
mjtNum* H = mj_stackAlloc(d, m->nv * m->nv);
|
||||
int* rownnz = (int*)mj_stackAlloc(d, m->nv);
|
||||
int* rowadr = (int*)mj_stackAlloc(d, m->nv);
|
||||
int* colind = (int*)mj_stackAlloc(d, m->nv * m->nv);
|
||||
|
||||
// compute D corresponding to quad states
|
||||
mjtNum* D = mj_stackAlloc(d, d->nefc);
|
||||
for (int i = 0; i < d->nefc; i++) {
|
||||
if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) {
|
||||
D[i] = d->efc_D[i];
|
||||
} else {
|
||||
D[i] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// time benchmark
|
||||
if (func) {
|
||||
for (auto s : state) {
|
||||
// compute H = J'*D*J, uncompressed layout
|
||||
func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
|
||||
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
|
||||
d->efc_JT_rowsuper, d);
|
||||
}
|
||||
} else {
|
||||
for (auto s : state) {
|
||||
// baseline depends on efc_J_rowsuper
|
||||
mju_superSparse(d->nefc, d->efc_J_rowsuper,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
|
||||
|
||||
// compute H = J'*D*J, uncompressed layout
|
||||
mju_sqrMatTDSparse_baseline(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper,
|
||||
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
|
||||
d->efc_JT_rowsuper, d);
|
||||
}
|
||||
}
|
||||
|
||||
// finalize
|
||||
mjFREESTACK;
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_sqrMatTDSparse_new(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTDSparse(state, &mju_sqrMatTDSparse);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTDSparse_new);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_sqrMatTDSparse_old(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTDSparse(state, nullptr);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTDSparse_old);
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
|
||||
@@ -18,6 +18,7 @@
|
||||
|
||||
#include <gmock/gmock.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "test/fixture.h"
|
||||
|
||||
namespace mujoco {
|
||||
@@ -180,5 +181,549 @@ TEST_F(EngineUtilSparseTest, MjuTransposeNullMatrix) {
|
||||
EXPECT_THAT(rowadrT, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0, 0));
|
||||
}
|
||||
|
||||
static constexpr char modelStr[] = R"(<mujoco/>)";
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
|
||||
// 0 0 0
|
||||
// M = 0 0 0
|
||||
// 0 0 0
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colind[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnz[] = {3, 3, 3};
|
||||
int rowadr[] = {0, 3, 6};
|
||||
|
||||
mjtNum matT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindT[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnzT[] = {3, 3, 3};
|
||||
int rowadrT[] = {0, 3, 6};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) {
|
||||
// 2 -1 1
|
||||
// M = 1 2 -1
|
||||
// 2 2 3
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {2, -1, 1, 2, -1, 2, 2, 2, 3};
|
||||
int colind[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnz[] = {3, 3, 3};
|
||||
int rowadr[] = {0, 3, 6};
|
||||
|
||||
mjtNum matT[] = {2, 2, 2, -1, -1, 2, 1, 2, 3};
|
||||
int colindT[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnzT[] = {3, 3, 3};
|
||||
int rowadrT[] = {0, 3, 6};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(12, 0, 12, 0, 6, 3, 12, 3, 14));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) {
|
||||
// 1 2 0
|
||||
// M = 0 3 0
|
||||
// 4 0 0
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 2, 3, 4};
|
||||
int colind[] = {0, 1, 1, 0};
|
||||
int rownnz[] = {2, 1, 1};
|
||||
int rowadr[] = {0, 2, 3};
|
||||
|
||||
mjtNum matT[] = {1, 4, 2, 3};
|
||||
int colindT[] = {0, 2, 0, 1};
|
||||
int rownnzT[] = {2, 2, 0};
|
||||
int rowadrT[] = {0, 2, 4};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {2, 3, 4};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(66, 4, 0, 4, 35, 0, 0, 0, 0));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 0, 1, 0, 0, 0, 0));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) {
|
||||
// 1 0 2
|
||||
// M = 0 0 3
|
||||
// 4 0 0
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 2, 3, 4};
|
||||
int colind[] = {0, 2, 2, 0};
|
||||
int rownnz[] = {2, 1, 1};
|
||||
int rowadr[] = {0, 2, 3};
|
||||
|
||||
mjtNum matT[] = {1, 4, 2, 3};
|
||||
int colindT[] = {0, 2, 0, 1};
|
||||
int rownnzT[] = {2, 0, 2};
|
||||
int rowadrT[] = {0, 2, 2};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {2, 3, 4};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(66, 4, 0, 0, 0, 0, 4, 35, 0));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 0, 0, 0, 0, 2, 0));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(2, 0, 2));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) {
|
||||
// 1 0 4
|
||||
// M = 0 0 0
|
||||
// 2 3 0
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 4, 2, 3};
|
||||
int colind[] = {0, 2, 0, 1};
|
||||
int rownnz[] = {2, 0, 2};
|
||||
int rowadr[] = {0, 2, 2};
|
||||
|
||||
mjtNum matT[] = {1, 2, 3, 4};
|
||||
int colindT[] = {0, 2, 2, 0};
|
||||
int rownnzT[] = {2, 1, 1};
|
||||
int rowadrT[] = {0, 2, 3};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(5, 6, 4, 6, 9, 0, 4, 16, 0));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) {
|
||||
// 1 0 2
|
||||
// M = 0 2 0
|
||||
// 0 0 3
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 2, 2, 3};
|
||||
int colind[] = {0, 2, 1, 2};
|
||||
int rownnz[] = {2, 1, 1};
|
||||
int rowadr[] = {0, 2, 3};
|
||||
|
||||
mjtNum matT[] = {1, 2, 2, 3};
|
||||
int colindT[] = {0, 1, 0, 2};
|
||||
int rownnzT[] = {1, 1, 2};
|
||||
int rowadrT[] = {0, 1, 2};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(1, 2, 0, 4, 0, 0, 2, 13, 0));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 1, 0, 0, 0, 2, 0));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) {
|
||||
// 1 2
|
||||
// M = 0 3
|
||||
// 4 0
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 2, 3, 4};
|
||||
int colind[] = {0, 1, 1, 0};
|
||||
int rownnz[] = {2, 1, 1};
|
||||
int rowadr[] = {0, 2, 3};
|
||||
|
||||
mjtNum matT[] = {1, 4, 2, 3};
|
||||
int colindT[] = {0, 2, 0, 1};
|
||||
int rownnzT[] = {2, 2};
|
||||
int rowadrT[] = {0, 2};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0};
|
||||
int rowadrH[] = {0, 0};
|
||||
|
||||
mjtNum diag[] = {2, 3, 4};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 2, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(66, 4, 4, 35));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 1));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(2, 2));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 2));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) {
|
||||
// M = 1 0 4
|
||||
// 2 3 0
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 4, 2, 3};
|
||||
int colind[] = {0, 2, 0, 1};
|
||||
int rownnz[] = {2, 2};
|
||||
int rowadr[] = {0, 2};
|
||||
|
||||
mjtNum matT[] = {1, 2, 3, 4};
|
||||
int colindT[] = {0, 1, 1, 0};
|
||||
int rownnzT[] = {2, 1, 1};
|
||||
int rowadrT[] = {0, 2, 3};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {2, 3};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 2, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(14, 18, 8, 18, 27, 0, 8, 32, 0));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) {
|
||||
// 1 2 2
|
||||
// M = 1 3 4
|
||||
// 4 4 4
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 2, 2, 1, 3, 4, 4, 4, 4};
|
||||
int colind[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnz[] = {3, 3, 3};
|
||||
int rowadr[] = {0, 3, 6};
|
||||
|
||||
mjtNum matT[] = {1, 1, 4, 2, 3, 4, 2, 4, 4};
|
||||
int colindT[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnzT[] = {3, 3, 3};
|
||||
int rowadrT[] = {0, 3, 6};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {2, 3, 4};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
NULL, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(69, 77, 80, 77, 99, 108, 80, 108, 120));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
|
||||
// 1 1 1
|
||||
// M = 2 2 2
|
||||
// 3 3 3
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 1, 1, 2, 2, 2, 3, 3, 3};
|
||||
int colind[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnz[] = {3, 3, 3};
|
||||
int rowadr[] = {0, 3, 6};
|
||||
|
||||
mjtNum matT[] = {1, 2, 3, 1, 2, 3, 1, 2, 3};
|
||||
int colindT[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
|
||||
int rownnzT[] = {3, 3, 3};
|
||||
int rowadrT[] = {0, 3, 6};
|
||||
int superowT[] = {2, 1, 0};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {1, 1, 1};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
superowT, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(14, 14, 14, 14, 14, 14, 14, 14, 14));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
|
||||
// 1 1 1
|
||||
// M = 0 0 0
|
||||
// 0 3 3
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 1, 1, 3, 3};
|
||||
int colind[] = {0, 1, 2, 1, 2};
|
||||
int rownnz[] = {3, 0, 2};
|
||||
int rowadr[] = {0, 3, 3};
|
||||
|
||||
mjtNum matT[] = {1, 1, 3, 1, 3};
|
||||
int colindT[] = {0, 0, 2, 0, 2};
|
||||
int rownnzT[] = {1, 2, 2};
|
||||
int rowadrT[] = {0, 1, 3};
|
||||
int superowT[] = {0, 1, 0};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {1, 1, 1};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
superowT, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(1, 1, 1, 1, 10, 10, 1, 10, 10));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
|
||||
// 1 1 1 1
|
||||
// M = 0 0 0 0
|
||||
// 0 0 3 3
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 1, 1, 1, 3, 3};
|
||||
int colind[] = {0, 1, 2, 3, 2, 3};
|
||||
int rownnz[] = {4, 0, 2};
|
||||
int rowadr[] = {0, 4, 4};
|
||||
|
||||
mjtNum matT[] = {1, 1, 1, 3, 1, 3};
|
||||
int colindT[] = {0, 0, 0, 2, 0, 2};
|
||||
int rownnzT[] = {1, 1, 2, 2};
|
||||
int rowadrT[] = {0, 1, 2, 4};
|
||||
int superowT[] = {1, 0, 1, 0};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {1, 1, 1};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 4, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
superowT, data);
|
||||
|
||||
EXPECT_THAT(matH,
|
||||
ElementsAre(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 10, 10, 1, 1, 10, 10));
|
||||
EXPECT_THAT(colindH,
|
||||
ElementsAre(0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(4, 4, 4, 4));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 4, 8, 12));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
|
||||
// 1 1 0 0 0
|
||||
// M = 1 1 0 0 0
|
||||
// 1 1 0 0 0
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 1, 1, 1, 1, 1};
|
||||
int colind[] = {0, 1, 0, 1, 0, 1};
|
||||
int rownnz[] = {2, 2, 2};
|
||||
int rowadr[] = {0, 2, 4};
|
||||
|
||||
mjtNum matT[] = {1, 1, 1, 1, 1, 1};
|
||||
int colindT[] = {0, 1, 2, 0, 1, 2};
|
||||
int rownnzT[] = {3, 3, 0, 0, 0};
|
||||
int rowadrT[] = {0, 3, 6, 6, 6};
|
||||
int superowT[] = {1, 0, 2, 1, 0};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0, 0, 0};
|
||||
|
||||
mjtNum diag[] = {1, 1, 1};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 5, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
superowT, data);
|
||||
|
||||
EXPECT_THAT(matH, ElementsAre(3, 3, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0));
|
||||
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0, 0, 0));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 5, 10, 15, 20));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
|
||||
// M = 1 1 1 1 2 2 2
|
||||
|
||||
mjModel* model = LoadModelFromString(modelStr);
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
mjtNum mat[] = {1, 1, 1, 1, 2, 2, 2};
|
||||
int colind[] = {0, 1, 2, 3, 4, 5, 6};
|
||||
int rownnz[] = {7};
|
||||
int rowadr[] = {0};
|
||||
|
||||
mjtNum matT[] = {1, 1, 1, 1, 2, 2, 2};
|
||||
int colindT[] = {0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzT[] = {1, 1, 1, 1, 1, 1, 1};
|
||||
int rowadrT[] = {0, 1, 2, 3, 4, 5, 6};
|
||||
int superowT[] = {3, 2, 1, 0, 2, 1, 0};
|
||||
|
||||
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
|
||||
int rownnzH[] = {0, 0, 0, 0, 0, 0, 0};
|
||||
int rowadrH[] = {0, 0, 0, 0, 0, 0, 0};
|
||||
|
||||
mju_sqrMatTDSparse(matH, mat, matT, NULL, 1, 7, rownnzH, rowadrH, colindH,
|
||||
rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
|
||||
superowT, data);
|
||||
|
||||
EXPECT_THAT(
|
||||
matH, ElementsAre(1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2,
|
||||
2, 2, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 2, 2, 2,
|
||||
2, 4, 4, 4, 2, 2, 2, 2, 4, 4, 4));
|
||||
EXPECT_THAT(colindH,
|
||||
ElementsAre(0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3,
|
||||
4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0,
|
||||
1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6
|
||||
|
||||
));
|
||||
EXPECT_THAT(rownnzH, ElementsAre(7, 7, 7, 7, 7, 7, 7));
|
||||
EXPECT_THAT(rowadrH, ElementsAre(0, 7, 14, 21, 28, 35, 42));
|
||||
|
||||
mj_deleteData(data);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
|
||||
Reference in New Issue
Block a user