Add mju_addToSymSparse
- Internal engine function to add a symmetric sparse matrix to a dense matrix. - Also minor refactors to related functions. PiperOrigin-RevId: 757954336 Change-Id: I8d8a48bcbbd5d6c618ae097a87df2bc1e6f5ef1d
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
755564a348
commit
436b5a8e1f
@@ -1089,20 +1089,18 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
|
||||
|
||||
// add inertia matrix to sparse destination matrix
|
||||
void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
|
||||
int* rownnz, int* rowadr, int* colind, mjtNum* M,
|
||||
int* M_rownnz, int* M_rowadr, int* M_colind) {
|
||||
int* rownnz, int* rowadr, int* colind, const mjtNum* M,
|
||||
const int* M_rownnz, const int* M_rowadr, const int* M_colind) {
|
||||
int nv = m->nv;
|
||||
|
||||
mj_markStack(d);
|
||||
mjtNum* buf_val = mjSTACKALLOC(d, nv, mjtNum);
|
||||
int* buf_ind = mjSTACKALLOC(d, nv, int);
|
||||
mjtNum* sparse_buf = mjSTACKALLOC(d, nv, mjtNum);
|
||||
|
||||
// add to destination
|
||||
for (int i=0; i < nv; i++) {
|
||||
rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1,
|
||||
rownnz[i], M_rownnz[i], colind + rowadr[i],
|
||||
M_colind + M_rowadr[i], sparse_buf, buf_ind);
|
||||
}
|
||||
mju_addToMatSparse(dst, rownnz, rowadr, colind, nv,
|
||||
M, M_rownnz, M_rowadr, M_colind,
|
||||
buf_val, buf_ind);
|
||||
|
||||
mj_freeStack(d);
|
||||
}
|
||||
|
||||
|
||||
@@ -137,8 +137,8 @@ MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
|
||||
|
||||
// add inertia matrix to sparse destination matrix
|
||||
MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
|
||||
int* rownnz, int* rowadr, int* colind, mjtNum* M,
|
||||
int* M_rownnz, int* M_rowadr, int* M_colind);
|
||||
int* rownnz, int* rowadr, int* colind, const mjtNum* M,
|
||||
const int* M_rownnz, const int* M_rowadr, const int* M_colind);
|
||||
|
||||
// add inertia matrix to dense destination matrix
|
||||
MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);
|
||||
|
||||
@@ -843,9 +843,9 @@ void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
|
||||
}
|
||||
|
||||
|
||||
|
||||
// compute M'*diag*M (diag=NULL: compute M'*M)
|
||||
void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) {
|
||||
// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional
|
||||
void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag,
|
||||
int nr, int nc, int flg_upper) {
|
||||
mjtNum tmp;
|
||||
|
||||
// half of MatMat routine: only lower triangle
|
||||
@@ -870,15 +870,23 @@ void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, in
|
||||
}
|
||||
}
|
||||
|
||||
// make symmetric
|
||||
for (int i=0; i < nc; i++) {
|
||||
for (int j=i+1; j < nc; j++) {
|
||||
res[i*nc+j] = res[j*nc+i];
|
||||
// flg_upper is set: make symmetric
|
||||
if (flg_upper) {
|
||||
for (int i=0; i < nc; i++) {
|
||||
for (int j=i+1; j < nc; j++) {
|
||||
res[i*nc+j] = res[j*nc+i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// compute M'*diag*M (diag=NULL: compute M'*M)
|
||||
void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) {
|
||||
mju_sqrMatTD_impl(res, mat, diag, nr, nc, /*flg_upper=*/ 1);
|
||||
}
|
||||
|
||||
|
||||
|
||||
// multiply matrices, first argument transposed
|
||||
void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
|
||||
|
||||
@@ -223,6 +223,10 @@ MJAPI void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
|
||||
MJAPI void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
|
||||
int r1, int c1, int c2);
|
||||
|
||||
// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional
|
||||
void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc,
|
||||
int flg_upper);
|
||||
|
||||
// compute M'*diag*M (diag=NULL: compute M'*M)
|
||||
MJAPI void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc);
|
||||
|
||||
|
||||
@@ -186,6 +186,41 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
|
||||
}
|
||||
|
||||
|
||||
// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers
|
||||
void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr,
|
||||
const mjtNum* M, const int* M_rownnz, const int* M_rowadr,
|
||||
const int* M_colind,
|
||||
mjtNum* buf_val, int* buf_ind) {
|
||||
for (int i=0; i < nr; i++) {
|
||||
rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1,
|
||||
rownnz[i], M_rownnz[i], colind + rowadr[i],
|
||||
M_colind + M_rowadr[i], buf_val, buf_ind);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// add symmetric matrix (lower triangle) to dense matrix, upper triangle optional
|
||||
void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n,
|
||||
const int* rownnz, const int* rowadr, const int* colind, int flg_upper) {
|
||||
for (int i=0; i < n; i++) {
|
||||
int start = rowadr[i];
|
||||
int end = start + rownnz[i];
|
||||
for (int adr=start; adr < end; adr++) {
|
||||
mjtNum val = mat[adr];
|
||||
int j = colind[adr];
|
||||
|
||||
// lower + diagonal
|
||||
res[i*n + j] += val;
|
||||
|
||||
// strict upper
|
||||
if (flg_upper && j < i) {
|
||||
res[j*n + i] += val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
// multiply symmetric matrix (only lower triangle represented) by vector:
|
||||
// res = (mat + strict_upper(mat')) * vec
|
||||
@@ -214,7 +249,7 @@ void mju_mulSymVecSparse(mjtNum* restrict res, const mjtNum* restrict mat,
|
||||
|
||||
// off-diagonals
|
||||
const int* ind = colind + adr;
|
||||
for (int k=0; k < diag; k++) {
|
||||
for (int k=diag-1; k >= 0; k--) {
|
||||
int j = ind[k];
|
||||
mjtNum val = row[k];
|
||||
res[i] += val * vec[j]; // strict lower
|
||||
|
||||
@@ -50,6 +50,17 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec
|
||||
MJAPI void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc,
|
||||
const int* rownnz, const int* rowadr, const int* colind);
|
||||
|
||||
// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers
|
||||
MJAPI void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr,
|
||||
const mjtNum* M, const int* M_rownnz, const int* M_rowadr,
|
||||
const int* M_colind,
|
||||
mjtNum* buf_val, int* buf_ind);
|
||||
|
||||
// add symmetric matrix (only lower triangle represented) to dense matrix
|
||||
MJAPI void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n,
|
||||
const int* rownnz, const int* rowadr, const int* colind,
|
||||
int flg_upper);
|
||||
|
||||
// multiply symmetric matrix (only lower triangle represented) by vector:
|
||||
// res = (mat + strict_upper(mat')) * vec
|
||||
MJAPI void mju_mulSymVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
|
||||
|
||||
@@ -1036,6 +1036,40 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
|
||||
EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24));
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuAddToSymSparse) {
|
||||
// 1 2 4
|
||||
// M = 2 3 0
|
||||
// 4 0 5
|
||||
|
||||
// only lower triangle represented
|
||||
mjtNum mat[] = {1, 2, 3, 4, 5};
|
||||
int colind[] = {0, 0, 1, 0, 2};
|
||||
int rownnz[] = {1, 2, 2};
|
||||
int rowadr[] = {0, 1, 3};
|
||||
|
||||
// 0 0 0
|
||||
// A = 5 4 2
|
||||
// 4 3 2
|
||||
mjtNum A[] = {0, 0, 0,
|
||||
5, 4, 2,
|
||||
4, 3, 2};
|
||||
|
||||
mju_addToSymSparse(A, mat, 3, rownnz, rowadr, colind, /*flg_upper=*/1);
|
||||
EXPECT_THAT(AsVector(A, 9), ElementsAre(1, 2, 4,
|
||||
7, 7, 2,
|
||||
8, 3, 7));
|
||||
|
||||
// same as A
|
||||
mjtNum B[] = {0, 0, 0,
|
||||
5, 4, 2,
|
||||
4, 3, 2};
|
||||
|
||||
mju_addToSymSparse(B, mat, 3, rownnz, rowadr, colind, /*flg_upper=*/0);
|
||||
EXPECT_THAT(AsVector(B, 9), ElementsAre(1, 0, 0,
|
||||
7, 7, 2,
|
||||
8, 3, 7));
|
||||
}
|
||||
|
||||
TEST_F(EngineUtilSparseTest, MjuMulSymVecSparse) {
|
||||
constexpr int n = 4;
|
||||
constexpr int nnz = 9;
|
||||
|
||||
Reference in New Issue
Block a user