chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Provide improved mju_sqrMatTDSparse implementation that doesn't require dense memory allocation for sparse matrices.

Browse Source 
PiperOrigin-RevId: 516812783
Change-Id: Ieb43337831d8d3b2c7f18a7facd8e0a0f2b0eff5
This commit is contained in:
Kyle Bayes
2023-03-15 06:55:28 -07:00
committed by Copybara-Service
parent fe18e58ad2
commit 056e849273
7 changed files with 830 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+1 -1
View File
@@ -18,7 +18,7 @@ General
~500,000 ``mjtNum``'s, now only requires ~6000. Very large models can now load and run with the CG solver. ~500,000 ``mjtNum``'s, now only requires ~6000. Very large models can now load and run with the CG solver.
- Modified :ref:`mju_error` and :ref:`mju_warning` to be variadic functions (support for printf-like arguments). The - Modified :ref:`mju_error` and :ref:`mju_warning` to be variadic functions (support for printf-like arguments). The
functions :ref:`mju_error_i`, :ref:`mju_error_s`, :ref:`mju_warning_i`, and :ref:`mju_warning_s` are now deprecated. functions :ref:`mju_error_i`, :ref:`mju_error_s`, :ref:`mju_warning_i`, and :ref:`mju_warning_s` are now deprecated.
- Implemented a performant :ref:`mju_sqrMatTDSparse` function that doesn't require dense memory allocation.
Python bindings Python bindings
src/engine/engine_core_constraint.c
+1 -3
View File
@@ -1699,7 +1699,6 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
int* rownnzT = (int*)mj_stackAlloc(d, nv); int* rownnzT = (int*)mj_stackAlloc(d, nv);
int* rowadrT = (int*)mj_stackAlloc(d, nv); int* rowadrT = (int*)mj_stackAlloc(d, nv);
int* colindT = (int*)mj_stackAlloc(d, nv*nefc); int* colindT = (int*)mj_stackAlloc(d, nv*nefc);
int* rowsuperT = (int*)mj_stackAlloc(d, nv);
// construct JM2 = backsubM2(J')' by rows // construct JM2 = backsubM2(J')' by rows
for (int r=0; r<nefc; r++) { for (int r=0; r<nefc; r++) {
@@ -1786,12 +1785,11 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// construct supernodes // construct supernodes
mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind); mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind);
mju_superSparse(nv, rowsuperT, rownnzT, rowadrT, colindT);
// AR = JM2 * JM2', uncompressed layout // AR = JM2 * JM2', uncompressed layout
mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc, mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc,
d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind,
rownnzT, rowadrT, colindT, rowsuperT, rownnzT, rowadrT, colindT, NULL,
rownnz, rowadr, colind, rowsuper, d); rownnz, rowadr, colind, rowsuper, d);
// compress layout of AR // compress layout of AR
src/engine/engine_solver.c
+1 -1
View File
@@ -1365,7 +1365,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, D, nefc, nv, mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, D, nefc, nv,
ctx->rownnz, ctx->rowadr, ctx->colind, ctx->rownnz, ctx->rowadr, ctx->colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_rownnz, d->efc_J_rowadr,
d->efc_J_colind, d->efc_J_rowsuper, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper, d); d->efc_JT_colind, d->efc_JT_rowsuper, d);
src/engine/engine_util_sparse.c
+96 -123
View File
@@ -377,6 +377,7 @@ void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
} }
// construct row supernodes // construct row supernodes
void mju_superSparse(int nr, int* rowsuper, void mju_superSparse(int nr, int* rowsuper,
const int* rownnz, const int* rowadr, const int* colind) { const int* rownnz, const int* rowadr, const int* colind) {
@@ -417,149 +418,121 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
int* res_rownnz, int* res_rowadr, int* res_colind, int* res_rownnz, int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr, const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper, const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT, const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT, const int* colindT, const int* rowsuperT,
mjData* d) { mjData* d) {
// allocate space for accumulation buffer and matT // allocate space for accumulation buffer and matT
mjMARKSTACK; mjMARKSTACK;
int* chain = (int*) mj_stackAlloc(d, 2*nc);
mjtNum* buffer = mj_stackAlloc(d, nc);
// set uncompressed layout // set uncompressed layout (the following doesn't depend on this layout)
for (int r=0; r<nc; r++) { for (int r=0; r<nc; r++) {
res_rowadr[r] = r*nc; res_rowadr[r] = r*nc;
} }
// compute lower-triangular uncompressed layout (nc per row) // a dense row buffer that stores the current row in the resulting matrix
for (int r=0; r<nc; r++) { mjtNum* buffer = mj_stackAlloc(d, nc);
// copy chain from parent
if (rowsuperT && r>0 && rowsuperT[r-1]>0) {
// copy parent chain
res_rownnz[r] = res_rownnz[r-1];
memcpy(res_colind+res_rowadr[r], res_colind+res_rowadr[r-1],
res_rownnz[r]*sizeof(int));
// add diagonal if rowT is not empty // these mark the currently set columns in the dense row buffer,
if (rownnzT[r]) { // used for when creating the resulting sparse row
res_colind[res_rowadr[r]+res_rownnz[r]] = r; int* markers = (int*) mj_stackAlloc(d, nc);
res_rownnz[r]++;
} for (int i=0; i<nc; i++) {
int* cols = res_colind+res_rowadr[i];
res_rownnz[i] = 0;
buffer[i] = 0;
markers[i] = 0;
// if rowsuper, use the previous row sparsity structure
if (rowsuperT && i>0 && rowsuperT[i-1]) {
res_rownnz[i] = res_rownnz[i-1];
memcpy(cols, res_colind+res_rowadr[i-1], res_rownnz[i]*sizeof(int));
} }
// construct chain // iterate through each row of M'
else { int end = rowadrT[i] + rownnzT[i];
// clear chain accumulation buffers for (int r = rowadrT[i]; r<end; r++) {
int nchain = 0; int t = colindT[r];
int inew = 0, iold = nc; mjtNum v = diag ? matT[r] * diag[t] : matT[r];
int lastadded = -1; for (int c=rowadr[t]; c<rowadr[t]+rownnz[t]; c++) {
int cc = colind[c];
// for each nonzero c in matT_row(r), add nonzeros of mat_row(c) to chain(r) // ignore upper triangle
for (int i=0; i<rownnzT[r]; i++) { if (cc>i) {
// save c
int c = colindT[rowadrT[r]+i];
// skip if a chain from same supernode was already added
if (rowsuper && lastadded>=0 && (c-lastadded)<=rowsuper[lastadded]) {
continue;
} else {
lastadded = c;
}
// swap chains
int adr = inew;
inew = iold;
iold = adr;
// merge chains
int nnewchain = 0;
adr = 0;
int end = rowadr[c]+rownnz[c];
for (int adr1=rowadr[c]; adr1<end; adr1++) {
// save column index from mat
int col_mat = colind[adr1];
// skip column indices in chain smaller than col_mat
while (adr<nchain && chain[iold + adr]<col_mat && chain[iold + adr]<=r) {
chain[inew + nnewchain++] = chain[iold + adr++];
}
// only lower-triangular
if (col_mat>r) {
break;
}
// existing element: advance chain
if (adr<nchain && chain[iold + adr]==col_mat) {
adr++;
}
// add column index from matT
chain[inew + nnewchain++] = col_mat;
}
// append the rest of the master chain
while (adr<nchain && chain[iold + adr]<=r) {
chain[inew + nnewchain++] = chain[iold + adr++];
}
// assign newchain
nchain = nnewchain;
}
// copy chain
res_rownnz[r] = nchain;
if (nchain) {
memcpy(res_colind+res_rowadr[r], chain+inew, nchain*sizeof(int));
}
}
}
// compute matrix data given uncompressed layout
for (int r=0; r<nc; r++) {
// clear buffer[colind] for this chain
int adr = res_rowadr[r];
for (int i=0; i<res_rownnz[r]; i++) {
buffer[res_colind[adr+i]] = 0;
}
// res_row(r) = sum_c ( matT(r,c) * diag(c) * mat_row(c) )
for (int i=0; i<rownnzT[r]; i++) {
// save c and matT(r,c)*diag(c)
int c = colindT[rowadrT[r]+i];
mjtNum matTrc = matT[rowadrT[r]+i];
if (diag) {
matTrc *= diag[c];
}
// process row
int end = rowadr[c]+rownnz[c];
for (int adr=rowadr[c]; adr<end; adr++) {
// get column index from mat, only lower-triangular
int adr1;
if ((adr1=colind[adr])>r) {
break; break;
} }
// add to buffer buffer[cc] += v*mat[c];
buffer[adr1] += matTrc*mat[adr];
// only need to insert nnz if not marked
if (!markers[cc]) {
markers[cc] = 1;
// since i is the rightmost column, it can be inserted at the end
if (cc==i) {
cols[res_rownnz[i]++] = cc;
continue;
}
// insert col in order via binary search
int l = 0, h = res_rownnz[i];
while (l<h) {
int m = (l + h) >> 1;
if (cols[m]<cc) {
l = m + 1;
} else {
h = m;
}
}
// cc is the rightmost column so far, it can be inserted at the end
if (l==res_rownnz[i]) {
cols[l] = cc;
res_rownnz[i]++;
continue;
}
// move the cols to the right
h = res_rownnz[i];
while (l<h) {
cols[h] = cols[h-1];
h--;
}
// insert
cols[l] = cc;
res_rownnz[i]++;
}
} }
} }
// copy buffer end = res_rownnz[i];
adr = res_rowadr[r];
for (int i=0; i<res_rownnz[r]; i++) { // rowsuperT: reuse sparsity, copy into res
res[adr+i] = buffer[res_colind[adr+i]]; if (rowsuperT && rowsuperT[i]) {
for (int r=0; r<end; r++) {
res[res_rowadr[i] + r] = buffer[cols[r]];
buffer[cols[r]] = 0;
}
} else {
// clear out buffers since sparsity cannot be reused
for (int r=0; r<end; r++) {
int cc = cols[r];
res[res_rowadr[i] + r] = buffer[cc];
res_colind[res_rowadr[i] + r] = cc;
buffer[cc] = 0;
markers[cc] = 0;
}
} }
} }
// make symmetric; uncompressed layout
for (int r=1; r<nc; r++) { // fill upper triangle
int end = nc*r+res_rownnz[r]-1; for (int i=0; i<nc; i++) {
for (int adr=nc*r; adr<end; adr++) { int end = res_rowadr[i] + res_rownnz[i] - 1;
// add to row given by column index for (int j=res_rowadr[i]; j<end; j++) {
int adr1 = nc*res_colind[adr] + res_rownnz[res_colind[adr]]++; int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++;
res[adr1] = res[adr]; res[adr] = res[j];
res_colind[adr1] = r; res_colind[adr] = i;
} }
} }
src/engine/engine_util_sparse.h
+2 -2
View File
@@ -66,8 +66,8 @@ MJAPI void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind); const int* rownnz, const int* rowadr, const int* colind);
// construct row supernodes // construct row supernodes
void mju_superSparse(int nr, int* rowsuper, MJAPI void mju_superSparse(int nr, int* rowsuper,
const int* rownnz, const int* rowadr, const int* colind); const int* rownnz, const int* rowadr, const int* colind);
// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout // compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT, MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
test/benchmark/engine_util_sparse_benchmark_test.cc
+184 -1
View File
@@ -28,6 +28,7 @@ namespace {
using CombineFuncPtr = decltype(&mju_combineSparse); using CombineFuncPtr = decltype(&mju_combineSparse);
using TransposeFuncPtr = decltype(&mju_transposeSparse); using TransposeFuncPtr = decltype(&mju_transposeSparse);
using SqrMatTDFuncPtr = decltype(&mju_sqrMatTDSparse);
// number of steps to roll out before benchmarking // number of steps to roll out before benchmarking
static const int kNumWarmupSteps = 500; static const int kNumWarmupSteps = 500;
@@ -39,6 +40,117 @@ std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
// ----------------------------- old functions -------------------------------- // ----------------------------- 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) // transpose sparse matrix (uncompressed)
void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline( void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
mjtNum* res, const mjtNum* mat, int nr, int nc, int* res_rownnz, mjtNum* res, const mjtNum* mat, int nr, int nc, int* res_rownnz,
@@ -426,8 +538,79 @@ BM_transposeSparse_old(benchmark::State& state) {
MujocoErrorTestGuard guard; MujocoErrorTestGuard guard;
BM_transposeSparse(state, &transposeSparse_baseline); BM_transposeSparse(state, &transposeSparse_baseline);
} }
BENCHMARK(BM_transposeSparse_old); 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
} // namespace mujoco } // namespace mujoco
test/engine/engine_util_sparse_test.cc
+545
View File
@@ -18,6 +18,7 @@
#include <gmock/gmock.h> #include <gmock/gmock.h>
#include <gtest/gtest.h> #include <gtest/gtest.h>
#include <mujoco/mujoco.h>
#include "test/fixture.h" #include "test/fixture.h"
namespace mujoco { namespace mujoco {
@@ -180,5 +181,549 @@ TEST_F(EngineUtilSparseTest, MjuTransposeNullMatrix) {
EXPECT_THAT(rowadrT, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)); 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
} // namespace mujoco } // namespace mujoco