Move local function to top of test file.

PiperOrigin-RevId: 757722664
Change-Id: Ief7d8ca9fda990bf582ca0feaa4b645f0a4bb6e7
This commit is contained in:
Yuval Tassa
2025-05-12 05:35:49 -07:00
committed by Copybara-Service
parent af088adaee
commit 9ac5fff41d
+25 -32
View File
@@ -26,6 +26,29 @@
namespace mujoco { namespace mujoco {
namespace { namespace {
// permute the rows and columns of a dense matrix
inline void PermuteMat(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* perm_r, const int* perm_c,
bool scatter_r, bool scatter_c) {
for (int r = 0; r < nr; r++) {
for (int c = 0; c < nc; c++) {
if (scatter_r && scatter_c) {
// scatter both
res[perm_r[r] * nc + perm_c[c]] = mat[r * nc + c];
} else if (scatter_r && !scatter_c) {
// scatter rows, gather columns
res[perm_r[r] * nc + c] = mat[r * nc + perm_c[c]];
} else if (!scatter_r && scatter_c) {
// gather rows, scatter columns
res[r * nc + perm_c[c]] = mat[perm_r[r] * nc + c];
} else {
// gather both
res[r * nc + c] = mat[perm_r[r] * nc + perm_c[c]];
}
}
}
}
using ::testing::ElementsAre; using ::testing::ElementsAre;
using EngineUtilSparseTest = MujocoTest; using EngineUtilSparseTest = MujocoTest;
@@ -1128,12 +1151,10 @@ TEST_F(EngineUtilSparseTest, BlockDiag) {
// 4x5 matrix with 3 blocks // 4x5 matrix with 3 blocks
constexpr int nr = 4; constexpr int nr = 4;
constexpr int nc = 5; constexpr int nc = 5;
const mjtNum mat[nr*nc] = { const mjtNum mat[nr * nc] = {1, 2, 0, 0, 0,
1, 2, 0, 0, 0,
0, 0, 3, 4, 0, 0, 0, 3, 4, 0,
0, 0, 5, 6, 0, 0, 0, 5, 6, 0,
0, 0, 0, 0, 7 0, 0, 0, 0, 7};
};
// block structure // block structure
constexpr int nb = 3; constexpr int nb = 3;
@@ -1156,11 +1177,6 @@ TEST_F(EngineUtilSparseTest, BlockDiag) {
7, 0, 0, 0, 0)); 7, 0, 0, 0, 0));
} }
void PermuteMat(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* perm_r, const int* perm_c,
bool scatter_r, bool scatter_c);
TEST_F(EngineUtilSparseTest, BlockDiagPerm) { TEST_F(EngineUtilSparseTest, BlockDiagPerm) {
// 4x5 matrix with 3 blocks // 4x5 matrix with 3 blocks
constexpr int nr = 4; constexpr int nr = 4;
@@ -1323,28 +1339,5 @@ TEST_F(EngineUtilSparseTest, PermuteMat) {
0, 0, 5, 6)); 0, 0, 5, 6));
} }
// local function for permuting the rows and columns of a dense matrix
void PermuteMat(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* perm_r, const int* perm_c,
bool scatter_r, bool scatter_c) {
for (int r = 0; r < nr; r++) {
for (int c = 0; c < nc; c++) {
if (scatter_r && scatter_c) {
// scatter both
res[perm_r[r] * nc + perm_c[c]] = mat[r * nc + c];
} else if (scatter_r && !scatter_c) {
// scatter rows, gather columns
res[perm_r[r] * nc + c] = mat[r * nc + perm_c[c]];
} else if (!scatter_r && scatter_c) {
// gather rows, scatter columns
res[r * nc + perm_c[c]] = mat[perm_r[r] * nc + c];
} else {
// gather both
res[r * nc + c] = mat[perm_r[r] * nc + perm_c[c]];
}
}
}
}
} // namespace } // namespace
} // namespace mujoco } // namespace mujoco