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

Rename C sparse structure to M in mjdata, improve docstrings.

Browse Source 
PiperOrigin-RevId: 758636638
Change-Id: If78acc423601d2911f514929b27f7b6d0af9ef58
This commit is contained in:
Yuval Tassa
2025-05-14 05:25:11 -07:00
committed by Copybara-Service
parent 4ba04586ae
commit dd28b887d4
21 changed files with 128 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+12 -12
View File
@@ -269,8 +269,8 @@ struct mjData_ {
// computed by mj_fwdPosition/mj_crb
mjtNum* crb; // com-based composite inertia and mass (nbody x 10)
mjtNum* qM; // total inertia (sparse) (nM x 1)
mjtNum* M; // total inertia (compressed sparse row) (nC x 1)
mjtNum* qM; // inertia (sparse) (nM x 1)
mjtNum* M; // reduced inertia (compressed sparse row) (nC x 1)
// computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1)
@@ -313,16 +313,16 @@ struct mjData_ {
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
int* mapM2C; // index mapping from M to C (nC x 1)
int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1)
int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1)
int* D_diag; // dof-dof: index of diagonal element (nv x 1)
int* D_colind; // dof-dof: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nM x 1)
int* M_rownnz; // reduced inertia: non-zeros in each row (nv x 1)
int* M_rowadr; // reduced inertia: address of each row in M_colind (nv x 1)
int* M_colind; // reduced inertia: column indices of non-zeros (nC x 1)
int* mapM2M; // index mapping from qM to M (nC x 1)
int* D_rownnz; // full inertia: non-zeros in each row (nv x 1)
int* D_rowadr; // full inertia: address of each row in D_colind (nv x 1)
int* D_diag; // full inertia: index of diagonal element (nv x 1)
int* D_colind; // full inertia: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from qM to D (nD x 1)
int* mapD2M; // index mapping from D to qM (nM x 1)
// computed by mj_implicit/mj_derivative
mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1)
include/mujoco/mjdata.h
+12 -12
View File
@@ -297,8 +297,8 @@ struct mjData_ {
// computed by mj_fwdPosition/mj_crb
mjtNum* crb; // com-based composite inertia and mass (nbody x 10)
mjtNum* qM; // total inertia (sparse) (nM x 1)
mjtNum* M; // total inertia (compressed sparse row) (nC x 1)
mjtNum* qM; // inertia (sparse) (nM x 1)
mjtNum* M; // reduced inertia (compressed sparse row) (nC x 1)
// computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1)
@@ -341,16 +341,16 @@ struct mjData_ {
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
int* mapM2C; // index mapping from M to C (nC x 1)
int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1)
int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1)
int* D_diag; // dof-dof: index of diagonal element (nv x 1)
int* D_colind; // dof-dof: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nM x 1)
int* M_rownnz; // reduced inertia: non-zeros in each row (nv x 1)
int* M_rowadr; // reduced inertia: address of each row in M_colind (nv x 1)
int* M_colind; // reduced inertia: column indices of non-zeros (nC x 1)
int* mapM2M; // index mapping from qM to M (nC x 1)
int* D_rownnz; // full inertia: non-zeros in each row (nv x 1)
int* D_rowadr; // full inertia: address of each row in D_colind (nv x 1)
int* D_diag; // full inertia: index of diagonal element (nv x 1)
int* D_colind; // full inertia: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from qM to D (nD x 1)
int* mapD2M; // index mapping from D to qM (nM x 1)
// computed by mj_implicit/mj_derivative
mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1)
include/mujoco/mjxmacro.h
+4 -4
View File
@@ -679,10 +679,10 @@
X ( int, B_rownnz, nbody, 1 ) \
X ( int, B_rowadr, nbody, 1 ) \
X ( int, B_colind, nB, 1 ) \
X ( int, C_rownnz, nv, 1 ) \
X ( int, C_rowadr, nv, 1 ) \
X ( int, C_colind, nC, 1 ) \
X ( int, mapM2C, nC, 1 ) \
X ( int, M_rownnz, nv, 1 ) \
X ( int, M_rowadr, nv, 1 ) \
X ( int, M_colind, nC, 1 ) \
X ( int, mapM2M, nC, 1 ) \
X ( int, D_rownnz, nv, 1 ) \
X ( int, D_rowadr, nv, 1 ) \
X ( int, D_diag, nv, 1 ) \
mjx/mujoco/mjx/_src/io.py
+4 -4
View File
@@ -629,10 +629,10 @@ def _make_data_c(
'B_rownnz': (m.nbody, np.int32),
'B_rowadr': (m.nbody, np.int32),
'B_colind': (m.nB, np.int32),
'C_rownnz': (m.nv, np.int32),
'C_rowadr': (m.nv, np.int32),
'C_colind': (m.nC, np.int32),
'mapM2C': (m.nC, np.int32),
'M_rownnz': (m.nv, np.int32),
'M_rowadr': (m.nv, np.int32),
'M_colind': (m.nC, np.int32),
'mapM2M': (m.nC, np.int32),
'D_rownnz': (m.nv, np.int32),
'D_rowadr': (m.nv, np.int32),
'D_diag': (m.nv, np.int32),
mjx/mujoco/mjx/_src/smooth_test.py
+1 -1
View File
@@ -92,7 +92,7 @@ class SmoothTest(absltest.TestCase):
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name
for i in range(m.nC):
qLDLegacy[d.mapM2C[i]] = d.qLD[i]
qLDLegacy[d.mapM2M[i]] = d.qLD[i]
_assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
_assert_attr_eq(d, dx._impl, 'qLDiagInv')
# com_vel
mjx/mujoco/mjx/_src/types.py
+4 -4
View File
@@ -985,10 +985,10 @@ class DataC(PyTreeNode):
B_rownnz: jax.Array # pylint:disable=invalid-name
B_rowadr: jax.Array # pylint:disable=invalid-name
B_colind: jax.Array # pylint:disable=invalid-name
C_rownnz: jax.Array # pylint:disable=invalid-name
C_rowadr: jax.Array # pylint:disable=invalid-name
C_colind: jax.Array # pylint:disable=invalid-name
mapM2C: jax.Array # pylint:disable=invalid-name
M_rownnz: jax.Array # pylint:disable=invalid-name
M_rowadr: jax.Array # pylint:disable=invalid-name
M_colind: jax.Array # pylint:disable=invalid-name
mapM2M: jax.Array # pylint:disable=invalid-name
D_rownnz: jax.Array # pylint:disable=invalid-name
D_rowadr: jax.Array # pylint:disable=invalid-name
D_diag: jax.Array # pylint:disable=invalid-name
python/mujoco/introspect/structs.py
+16 -16
View File
@@ -5413,7 +5413,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='total inertia (sparse)',
doc='inertia (sparse)',
array_extent=('nM',),
),
StructFieldDecl(
@@ -5421,7 +5421,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='total inertia (compressed sparse row)',
doc='reduced inertia (compressed sparse row)',
array_extent=('nC',),
),
StructFieldDecl(
@@ -5601,35 +5601,35 @@ STRUCTS: Mapping[str, StructDecl] = dict([
array_extent=('nB',),
),
StructFieldDecl(
name='C_rownnz',
name='M_rownnz',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='reduced dof-dof: non-zeros in each row',
doc='reduced inertia: non-zeros in each row',
array_extent=('nv',),
),
StructFieldDecl(
name='C_rowadr',
name='M_rowadr',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='reduced dof-dof: address of each row in C_colind',
doc='reduced inertia: address of each row in M_colind',
array_extent=('nv',),
),
StructFieldDecl(
name='C_colind',
name='M_colind',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='reduced dof-dof: column indices of non-zeros',
doc='reduced inertia: column indices of non-zeros',
array_extent=('nC',),
),
StructFieldDecl(
name='mapM2C',
name='mapM2M',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='index mapping from M to C',
doc='index mapping from qM to M',
array_extent=('nC',),
),
StructFieldDecl(
@@ -5637,7 +5637,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='dof-dof: non-zeros in each row',
doc='full inertia: non-zeros in each row',
array_extent=('nv',),
),
StructFieldDecl(
@@ -5645,7 +5645,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='dof-dof: address of each row in D_colind',
doc='full inertia: address of each row in D_colind',
array_extent=('nv',),
),
StructFieldDecl(
@@ -5653,7 +5653,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='dof-dof: index of diagonal element',
doc='full inertia: index of diagonal element',
array_extent=('nv',),
),
StructFieldDecl(
@@ -5661,7 +5661,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='dof-dof: column indices of non-zeros',
doc='full inertia: column indices of non-zeros',
array_extent=('nD',),
),
StructFieldDecl(
@@ -5669,7 +5669,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='index mapping from M to D',
doc='index mapping from qM to D',
array_extent=('nD',),
),
StructFieldDecl(
@@ -5677,7 +5677,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='index mapping from D to M',
doc='index mapping from D to qM',
array_extent=('nM',),
),
StructFieldDecl(
src/engine/engine_core_constraint.c
+6 -6
View File
@@ -2039,7 +2039,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// inverse square root of D from inertia LDL decomposition
mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
for (int i=0; i < nv; i++) {
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
}
@@ -2076,10 +2076,10 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
}
// traverse row j of C, marking new unique nonzeros
int nnzC = d->C_rownnz[j];
int adrC = d->C_rowadr[j];
int nnzC = d->M_rownnz[j];
int adrC = d->M_rowadr[j];
for (int k=0; k < nnzC; k++) {
int c = d->C_colind[adrC + k];
int c = d->M_colind[adrC + k];
if (marker[c] != r) {
marker[c] = r;
nnz++;
@@ -2159,10 +2159,10 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
continue;
}
int j = B_colind[i];
int adrC = d->C_rowadr[j];
int adrC = d->M_rowadr[j];
mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
nnzB, B_colind + adrB,
d->C_rownnz[j]-1, d->C_colind + adrC, -b);
d->M_rownnz[j]-1, d->M_colind + adrC, -b);
}
// B(r,:) <- sqrt(inv(D)) * B(r,:)
src/engine/engine_core_smooth.c
+6 -6
View File
@@ -1579,7 +1579,7 @@ void mj_makeM(const mjModel* m, mjData* d) {
TM_START;
mj_crb(m, d);
mj_tendonArmature(m, d);
mju_gather(d->M, d->qM, d->mapM2C, m->nC);
mju_gather(d->M, d->qM, d->mapM2M, m->nC);
TM_END(mjTIMER_POS_INERTIA);
}
@@ -1653,7 +1653,7 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD
void mj_factorM(const mjModel* m, mjData* d) {
TM_START;
mju_copy(d->qLD, d->M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
TM_ADD(mjTIMER_POS_INERTIA);
}
@@ -1895,7 +1895,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
mju_copy(x, y, n*m->nv);
}
mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
@@ -1906,9 +1906,9 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
int nv = m->nv;
// local copies of key variables
const int* rownnz = d->C_rownnz;
const int* rowadr = d->C_rowadr;
const int* colind = d->C_colind;
const int* rownnz = d->M_rownnz;
const int* rowadr = d->M_rowadr;
const int* colind = d->M_colind;
const int* diagnum = m->dof_simplenum;
const mjtNum* qLD = d->qLD;
src/engine/engine_forward.c
+6 -6
View File
@@ -869,18 +869,18 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
// qH = M + h*diag(B)
mju_copy(d->qH, d->M, nC);
for (int i=0; i < nv; i++) {
d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
d->qH[d->M_rowadr[i] + d->M_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
}
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
// solve
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
mju_copy(qacc, qfrc, m->nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
// advance state and time
@@ -1053,16 +1053,16 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
// gather qH <- MhB (legacy to CSR)
mju_gather(d->qH, MhB, d->mapM2C, nC);
mju_gather(d->qH, MhB, d->mapM2M, nC);
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
mju_copy(qacc, qfrc, nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
} else {
mjERROR("integrator must be implicit or implicitfast");
src/engine/engine_io.c
+6 -6
View File
@@ -1139,8 +1139,8 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
const int* rownnz;
const int* rowadr;
if (reduced && !upper) {
rownnz = d->C_rownnz;
rowadr = d->C_rowadr;
rownnz = d->M_rownnz;
rowadr = d->M_rowadr;
} else if (!reduced && upper) {
rownnz = d->D_rownnz;
rowadr = d->D_rowadr;
@@ -1248,12 +1248,12 @@ static void makeDofDofmaps(const mjModel* m, mjData* d) {
}
// make mapM2C
for (int i=0; i < nC; i++) d->mapM2C[i] = -1;
copyM2Sparse(m, d, d->mapM2C, M, /*reduced=*/1, /*upper=*/0);
for (int i=0; i < nC; i++) d->mapM2M[i] = -1;
copyM2Sparse(m, d, d->mapM2M, M, /*reduced=*/1, /*upper=*/0);
// check that all indices are filled in
for (int i=0; i < nC; i++) {
if (d->mapM2C[i] < 0) {
if (d->mapM2M[i] < 0) {
mjERROR("unassigned index in mapM2C");
}
}
@@ -1976,7 +1976,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
checkDBSparse(m, d);
// make C
makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, NULL, d->C_colind, /*reduced=*/1, /*upper=*/0);
makeDofDofSparse(m, d, d->M_rownnz, d->M_rowadr, NULL, d->M_colind, /*reduced=*/1, /*upper=*/0);
// make index mappings: mapM2D, mapD2M, mapM2C, mapM2M
makeDofDofmaps(m, d);
src/engine/engine_island.c
+1 -1
View File
@@ -537,7 +537,7 @@ void mj_island(const mjModel* m, mjData* d) {
// inertia: block-diagonalize both iLD <- qLD and iM <- qM
mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind,
d->qLD, d->C_rownnz, d->C_rowadr, d->C_colind,
d->qLD, d->M_rownnz, d->M_rowadr, d->M_colind,
nidof, nisland,
d->map_idof2dof, d->map_dof2idof,
d->island_idofadr, d->island_idofadr,
src/engine/engine_print.c
+17 -17
View File
@@ -1124,14 +1124,14 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
d->moment_rowadr, d->moment_colind, fp, float_format);
printArray("CRB", m->nbody, 10, d->crb, fp, float_format);
printInertia("QM", d->qM, m, fp, float_format);
printSparse("M", d->M, m->nv, d->C_rownnz,
d->C_rowadr, d->C_colind, fp, float_format);
printSparse("QLD", d->qLD, m->nv, d->C_rownnz,
d->C_rowadr, d->C_colind, fp, float_format);
printSparse("M", d->M, m->nv, d->M_rownnz,
d->M_rowadr, d->M_colind, fp, float_format);
printSparse("QLD", d->qLD, m->nv, d->M_rownnz,
d->M_rowadr, d->M_colind, fp, float_format);
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
if (!mju_isZero(d->qHDiagInv, m->nv)) {
printSparse("QH", d->qH, m->nv, d->C_rownnz, d->C_rowadr, d->C_colind, fp, float_format);
printSparse("QH", d->qH, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind, fp, float_format);
printArray("QHDIAGINV", m->nv, 1, d->qHDiagInv, fp, float_format);
}
@@ -1160,34 +1160,34 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
}
fprintf(fp, "\n\n");
// C sparse structure
mj_printSparsity("C: reduced dof-dof matrix", m->nv, m->nv, d->C_rowadr, NULL, d->C_rownnz,
NULL, d->C_colind, fp);
// M sparse structure
mj_printSparsity("M: reduced inertia matrix", m->nv, m->nv, d->M_rowadr, NULL, d->M_rownnz,
NULL, d->M_colind, fp);
fprintf(fp, NAME_FORMAT, "C_rownnz");
fprintf(fp, NAME_FORMAT, "M_rownnz");
for (int i = 0; i < m->nv; i++) {
fprintf(fp, " %d", d->C_rownnz[i]);
fprintf(fp, " %d", d->M_rownnz[i]);
}
fprintf(fp, "\n\n");
// C_rowadr
fprintf(fp, NAME_FORMAT, "C_rowadr");
fprintf(fp, NAME_FORMAT, "M_rowadr");
for (int i = 0; i < m->nv; i++) {
fprintf(fp, " %d", d->C_rowadr[i]);
fprintf(fp, " %d", d->M_rowadr[i]);
}
fprintf(fp, "\n\n");
// C_colind
fprintf(fp, NAME_FORMAT, "C_colind");
fprintf(fp, NAME_FORMAT, "M_colind");
for (int i = 0; i < m->nC; i++) {
fprintf(fp, " %d", d->C_colind[i]);
fprintf(fp, " %d", d->M_colind[i]);
}
fprintf(fp, "\n\n");
// mapM2C
fprintf(fp, NAME_FORMAT, "mapM2C");
// mapM2M
fprintf(fp, NAME_FORMAT, "mapM2M");
for (int i = 0; i < m->nC; i++) {
fprintf(fp, " %d", d->mapM2C[i]);
fprintf(fp, " %d", d->mapM2M[i]);
}
fprintf(fp, "\n\n");
src/engine/engine_solver.c
+3 -3
View File
@@ -886,10 +886,10 @@ static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int
ctx->qacc = d->qacc;
// inertia
ctx->M_rownnz = d->C_rownnz;
ctx->M_rowadr = d->C_rowadr;
ctx->M_rownnz = d->M_rownnz;
ctx->M_rowadr = d->M_rowadr;
ctx->M_diagnum = m->dof_simplenum;
ctx->M_colind = d->C_colind;
ctx->M_colind = d->M_colind;
ctx->M = d->M;
ctx->qLD = d->qLD;
ctx->qLDiagInv = d->qLDiagInv;
src/engine/engine_support.c
+5 -5
View File
@@ -1047,14 +1047,14 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
// non-simple: add off-diagonals
if (!m->dof_simplenum[i]) {
int adr = d->C_rowadr[i];
res[i] += mju_dotSparse(qLD+adr, vec, d->C_rownnz[i] - 1, d->C_colind+adr);
int adr = d->M_rowadr[i];
res[i] += mju_dotSparse(qLD+adr, vec, d->M_rownnz[i] - 1, d->M_colind+adr);
}
}
// res *= sqrt(D)
for (int i=0; i < nv; i++) {
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
res[i] *= mju_sqrt(qLD[diag]);
}
}
@@ -1072,10 +1072,10 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
// gather C <- qM (legacy to CSR)
mjtNum* C = mjSTACKALLOC(d, nC, mjtNum);
mju_gather(C, d->qM, d->mapM2C, nC);
mju_gather(C, d->qM, d->mapM2M, nC);
// add to dst
mj_addMSparse(m, d, dst, rownnz, rowadr, colind, C, d->C_rownnz, d->C_rowadr, d->C_colind);
mj_addMSparse(m, d, dst, rownnz, rowadr, colind, C, d->M_rownnz, d->M_rowadr, d->M_colind);
mj_freeStack(d);
}
test/benchmark/factorI_benchmark_test.cc
+2 -2
View File
@@ -46,7 +46,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
// M: mass matrix in CSR format
mjtNum* M = mj_stackAllocNum(d, m->nC);
mju_gather(M, d->qM, d->mapM2C, m->nC);
mju_gather(M, d->qM, d->mapM2M, m->nC);
// LDlegacy: legacy LD matrix (size nM)
mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
@@ -59,7 +59,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
} else {
mju_copy(d->qLD, M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
}
}
test/benchmark/inertia_benchmark_test.cc
+3 -3
View File
@@ -48,7 +48,7 @@ static void BM_solve(benchmark::State& state, SolveType type) {
// M: mass matrix in CSR format
mjtNum* M = mj_stackAllocNum(d, m->nC);
mju_gather(M, d->qM, d->mapM2C, m->nC);
mju_gather(M, d->qM, d->mapM2M, m->nC);
// LDlegacy: legacy LD matrix (size nM)
mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
@@ -73,9 +73,9 @@ static void BM_solve(benchmark::State& state, SolveType type) {
case SolveType::kCsr:
mju_copy(d->qLD, M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
}
}
test/benchmark/solveLD_benchmark_test.cc
+2 -2
View File
@@ -54,7 +54,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
// scatter into legacy matrix
mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
mju_zero(LDlegacy, m->nM);
mju_scatter(LDlegacy, d->qLD, d->mapM2C, m->nC);
mju_scatter(LDlegacy, d->qLD, d->mapM2M, m->nC);
// benchmark
while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
@@ -64,7 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv);
} else {
mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
}
}
test/engine/engine_core_smooth_test.cc
+12 -12
View File
@@ -649,7 +649,7 @@ TEST_F(CoreSmoothTest, FactorI) {
int nv = model->nv;
vector<mjtNum> Ldense(nv*nv, 0);
mju_sparse2dense(Ldense.data(), data->qLD, nv, nv,
data->C_rownnz, data->C_rowadr, data->C_colind);
data->M_rownnz, data->M_rowadr, data->M_colind);
for (int i=0; i < nv; i++) {
// set diagonal to 1
Ldense[i*nv+i] = 1;
@@ -658,7 +658,7 @@ TEST_F(CoreSmoothTest, FactorI) {
// dense D matrix
vector<mjtNum> Ddense(nv*nv);
mju_sparse2dense(Ddense.data(), data->qLD, nv, nv,
data->C_rownnz, data->C_rowadr, data->C_colind);
data->M_rownnz, data->M_rowadr, data->M_colind);
for (int i=0; i < nv; i++) {
for (int j=0; j < nv; j++) {
// zero everything except the diagonal
@@ -698,12 +698,12 @@ TEST_F(CoreSmoothTest, SolveLDs) {
// copy M into LD: Legacy format
vector<mjtNum> LDlegacy(nM, 0);
mju_scatter(LDlegacy.data(), d->qLD, d->mapM2C, nC);
mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nC);
// compare LD and LDs densified matrices
vector<mjtNum> LDdense(nv*nv);
mju_sparse2dense(LDdense.data(), d->qLD, nv, nv,
d->C_rownnz, d->C_rowadr, d->C_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
vector<mjtNum> LDdense2(nv*nv);
mj_fullM(m, LDdense2.data(), LDlegacy.data());
@@ -722,7 +722,7 @@ TEST_F(CoreSmoothTest, SolveLDs) {
mj_solveLD_legacy(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
// expect vectors to match up to floating point precision
for (int i=0; i < nv; i++) {
@@ -746,7 +746,7 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
// copy LD into LDlegacy: Legacy format
vector<mjtNum> LDlegacy(m->nM, 0);
mju_scatter(LDlegacy.data(), d->qLD, d->mapM2C, m->nC);
mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, m->nC);
// compare n LD and LDs vector solve
int n = 3;
@@ -757,7 +757,7 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
mj_solveLD_legacy(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
// expect vectors to match up to floating point precision
for (int i=0; i < nv*n; i++) {
@@ -781,7 +781,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
int nv = m->nv;
vector<mjtNum> sqrtInvD(nv);
for (int i=0; i < nv; i++) {
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
}
@@ -795,7 +795,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
// expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v)
for (int i=0; i < n; i++) {
@@ -824,17 +824,17 @@ TEST_F(CoreSmoothTest, FactorIs) {
// copy qLDlegacy into qLDexpected: CSR format
vector<mjtNum> qLDexpected(nC);
mju_gather(qLDexpected.data(), qLDlegacy.data(), d->mapM2C, nC);
mju_gather(qLDexpected.data(), qLDlegacy.data(), d->mapM2M, nC);
// copy qM into qLD: CSR format
vector<mjtNum> qLD(nC);
mju_gather(qLD.data(), d->qM, d->mapM2C, nC);
mju_gather(qLD.data(), d->qM, d->mapM2M, nC);
vector<mjtNum> qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv);
vector<mjtNum> qLDiagInv(nv, 0);
mj_factorI(qLD.data(), qLDiagInv.data(), nv,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
// expect outputs to match to floating point precision
EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected));
test/engine/engine_derivative_test.cc
+2 -2
View File
@@ -436,7 +436,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
Ac[nv*nv + i*nv + i] = -m->dof_damping[i];
}
mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2*nv,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
// A = [dt*Ac; Ac]
mju_transpose(A, Ac, 2*nv, nv);
@@ -464,7 +464,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
d->moment_rowadr, d->moment_colind);
mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
mju_transpose(BcT, Bc, nu, nv);
mju_scl(B, BcT, dt*dt, nu*nv);
mju_scl(B+nu*nv, BcT, dt, nu*nv);
unity/Runtime/Bindings/MjBindings.cs
+4 -4
View File
@@ -4952,10 +4952,10 @@ public unsafe struct mjData_ {
public int* B_rownnz;
public int* B_rowadr;
public int* B_colind;
public int* C_rownnz;
public int* C_rowadr;
public int* C_colind;
public int* mapM2C;
public int* M_rownnz;
public int* M_rowadr;
public int* M_colind;
public int* mapM2M;
public int* D_rownnz;
public int* D_rowadr;
public int* D_diag;