diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index cc2c7fd4..89403bca 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -206,6 +206,28 @@ void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, +// multiply transposed sparse matrix and dense vector: res = mat' * vec. +void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc, + const int* rownnz, const int* rowadr, const int* colind) { + // clear res + mju_zero(res, nc); + + for (int i=0; i < nr; i++) { + int nnz = rownnz[i]; + int adr = rowadr[i]; + const int* ind = colind + adr; + const mjtNum* row = mat + adr; + mjtNum scl = vec[i]; + + // add row scaled by the corresponding vector element + for (int j=0; j < nnz; j++) { + res[ind[j]] += row[j] * scl; + } + } +} + + + // res = res*scl1 + vec*scl2 static void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, int n) { #ifdef mjUSEAVX diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index c5576953..8b252dc0 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -46,6 +46,10 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec int nr, const int* rownnz, const int* rowadr, const int* colind, const int* rowsuper); +// multiply transposed sparse matrix and dense vector: res = mat' * 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); + // compress layout of sparse matrix MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind); diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index dfdadd82..92b8fc8d 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -30,8 +30,9 @@ namespace { using ::testing::ElementsAre; using EngineUtilSparseTest = MujocoTest; -std::vector AsVector(const int* array, int n) { - return std::vector(array, array + n); +template +std::vector AsVector(const T* array, int n) { + return std::vector(array, array + n); } TEST_F(EngineUtilSparseTest, MjuDot) { @@ -1043,5 +1044,25 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { mj_deleteModel(model); } +TEST_F(EngineUtilSparseTest, MjuMulMatTVec) { + int nr = 2; + int nc = 3; + mjtNum mat[] = {1, 2, 0, + 0, 3, 4}; + + mjtNum mat_sparse[6]; + int rownnz[2]; + int rowadr[2]; + int colind[4]; + mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind); + + // multiply: res = mat' * vec + mjtNum vec[] = {5, 6}; + mjtNum res[3]; + mju_mulMatTVecSparse(res, mat_sparse, vec, nr, nc, rownnz, rowadr, colind); + + EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24)); +} + } // namespace } // namespace mujoco