diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index d7b8db51..64c50073 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -423,7 +423,8 @@ Get name of object with the specified mjtObj type and id, returns NULL if name n .. mujoco-include:: mj_fullM -Convert sparse inertia matrix M into full (i.e. dense) matrix. +Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix. +|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same size as ``mjData.qM``. .. _mj_mulM: diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 9d425623..03c1e7fb 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -233,6 +233,11 @@ found, the function will return ``distmax`` and ``fromto``, if given, will be se In order to determine whether a geom pair uses ``mjc_Convex``, inspect the table at the top of `engine_collision_driver.c `__. +.. _mj_fullM: + +Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix. +|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same size as ``mjData.qM``. + .. _mj_mulM: This function multiplies the joint-space inertia matrix stored in mjData.qM by a vector. qM has a custom sparse format diff --git a/doc/includes/references.h b/doc/includes/references.h index dcb3f22d..693a2a99 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -272,7 +272,7 @@ struct mjData_ { mjtNum* qM; // total inertia (sparse) (nM x 1) // computed by mj_fwdPosition/mj_factorM - mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) + mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) // computed by mj_collisionTree @@ -305,7 +305,7 @@ struct mjData_ { mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) // computed by mj_Euler or mj_implicit - mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) + mjtNum* qH; // L'*D*L factorization of modified M (nC x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) // computed by mj_resetData diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 3c4d889e..c573a2c7 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -300,7 +300,7 @@ struct mjData_ { mjtNum* qM; // total inertia (sparse) (nM x 1) // computed by mj_fwdPosition/mj_factorM - mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) + mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) // computed by mj_collisionTree @@ -333,7 +333,7 @@ struct mjData_ { mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) // computed by mj_Euler or mj_implicit - mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) + mjtNum* qH; // L'*D*L factorization of modified M (nC x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) // computed by mj_resetData diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 50f0f8e2..88b802b8 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -637,7 +637,7 @@ X ( mjtNum, actuator_moment, nJmom, 1 ) \ X ( mjtNum, crb, nbody, 10 ) \ X ( mjtNum, qM, nM, 1 ) \ - X ( mjtNum, qLD, nM, 1 ) \ + X ( mjtNum, qLD, nC, 1 ) \ X ( mjtNum, qLDiagInv, nv, 1 ) \ XMJV( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ XMJV( mjtByte, bvh_active, nbvh, 1 ) \ @@ -654,7 +654,7 @@ X ( mjtNum, qfrc_passive, nv, 1 ) \ X ( mjtNum, subtree_linvel, nbody, 3 ) \ X ( mjtNum, subtree_angmom, nbody, 3 ) \ - X ( mjtNum, qH, nM, 1 ) \ + X ( mjtNum, qH, nC, 1 ) \ X ( mjtNum, qHDiagInv, nv, 1 ) \ X ( int, B_rownnz, nbody, 1 ) \ X ( int, B_rowadr, nbody, 1 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index abe63628..605fcd4b 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -5261,7 +5261,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='mjtNum'), ), doc="L'*D*L factorization of M (sparse)", - array_extent=('nM',), + array_extent=('nC',), ), StructFieldDecl( name='qLDiagInv', @@ -5397,7 +5397,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='mjtNum'), ), doc="L'*D*L factorization of modified M", - array_extent=('nM',), + array_extent=('nC',), ), StructFieldDecl( name='qHDiagInv', diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 6acef205..24704b18 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -368,7 +368,7 @@ def make_data( 'efc_aref': (nefc, float), 'efc_force': (nefc, float), '_qM_sparse': (m.nM, float), - '_qLD_sparse': (m.nM, float), + '_qLD_sparse': (m.nC, float), '_qLDiagInv_sparse': (m.nv, float), } diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 11cf85fd..bf9176ba 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -91,7 +91,10 @@ class SmoothTest(absltest.TestCase): _assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse') # factor_m dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d)) - _assert_attr_eq(d, dx, 'qLD') + qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name + for i in range(m.nC): + qLDLegacy[d.mapM2C[i]] = d.qLD[i] + _assert_eq(qLDLegacy, dx.qLD, 'qLD') _assert_attr_eq(d, dx, 'qLDiagInv') _assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse') _assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse') diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 3a7f3f9b..a7fd3617 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -1348,7 +1348,7 @@ class Data(PyTreeNode): efc_aref: reference pseudo-acceleration (nefc,) efc_force: constraint force in constraint space (nefc,) _qM_sparse: qM in sparse representation (nM,) - _qLD_sparse: qLD in sparse representation (nM,) + _qLD_sparse: qLD in sparse representation (nC,) _qLDiagInv_sparse: qLDiagInv in sparse representation (nv,) """ # fmt: skip # constant sizes: diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index df37ab30..af773edd 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -2131,7 +2131,8 @@ 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++) { - sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); + int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1; + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); } // sparse @@ -2238,13 +2239,6 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { // === in-place sparse back-substitution: B <- B * M^-1/2 - // make qLD - int nC = m->nC; - mjtNum* qLD = mjSTACKALLOC(d, nC, mjtNum); - for (int i=0; i < nC; i++) { - qLD[i] = d->qLD[d->mapM2C[i]]; - } - // sparse backsubM2 (half of LD back-substitution) for (int r=0; r < nefc; r++) { int nnzB = B_rownnz[r]; @@ -2258,7 +2252,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { } int j = B_colind[i]; int adrC = d->C_rowadr[j]; - mju_addToSclSparseInc(B + adrB, qLD + adrC, + mju_addToSclSparseInc(B + adrB, d->qLD + adrC, nnzB, B_colind + adrB, d->C_rownnz[j]-1, d->C_colind + adrC, -b); } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index c7bffa00..9f73b092 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1465,7 +1465,11 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd void mj_factorM(const mjModel* m, mjData* d) { TM_START; - mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); + int nC = m->nC; + for (int i=0; i < nC; i++) { + d->qLD[i] = d->qM[d->mapM2C[i]]; + } + mj_factorIs(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); TM_ADD(mjTIMER_POS_INERTIA); } @@ -1709,18 +1713,20 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { if (x != y) { mju_copy(x, y, n*m->nv); } - mj_solveLD(m, x, n, d->qLD, d->qLDiagInv); + mj_solveLDs(x, d->qLD, d->qLDiagInv, m->nv, n, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } // in-place sparse backsubstitution for one island: x = inv(L'*D*L)*x // L is in lower triangle of qLD; D is on diagonal of qLD -void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int island) { +void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) { // if no islands, call mj_solveLD const mjtNum* qLD = d->qLD; const mjtNum* qLDiagInv = d->qLDiagInv; if (island < 0) { - mj_solveLD(m, x, 1, qLD, qLDiagInv); + mj_solveLDs(x, qLD, qLDiagInv, m->nv, 1, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); return; } @@ -1730,14 +1736,6 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan const int* colind = d->C_colind; const int* diagnum = m->dof_simplenum; - // temporary: make local CSR version of qLD - int nC = m->nC; - mj_markStack(d); - mjtNum* qLDs = mjSTACKALLOC(d, nC, mjtNum); - for (int i=0; i < nC; i++) { - qLDs[i] = d->qLD[d->mapM2C[i]]; - } - // local constants: island specific int ndof = d->island_dofnum[island]; const int* dofind = d->island_dofind + d->island_dofadr[island]; @@ -1751,7 +1749,7 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int adr=end-1; adr >= start; adr--) { - x[islandind[colind[adr]]] -= qLDs[adr] * x_k; + x[islandind[colind[adr]]] -= qLD[adr] * x_k; } } } @@ -1773,11 +1771,9 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int adr=end-1; adr >= start; adr--) { - x[k] -= x[islandind[colind[adr]]] * qLDs[adr]; + x[k] -= x[islandind[colind[adr]]] * qLD[adr]; } } - - mj_freeStack(d); } @@ -1785,23 +1781,18 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, const mjtNum* sqrtInvD, int n) { + int nv = m->nv; + // local copies of key variables - int nv = m->nv, nC = m->nC; const int* rownnz = d->C_rownnz; const int* rowadr = d->C_rowadr; const int* colind = d->C_colind; const int* diagnum = m->dof_simplenum; + const mjtNum* qLD = d->qLD; // x = y mju_copy(x, y, n * nv); - // temporary: make local CSR version of qLD - mj_markStack(d); - mjtNum* qLD = mjSTACKALLOC(d, nC, mjtNum); - for (int i=0; i < nC; i++) { - qLD[i] = d->qLD[d->mapM2C[i]]; - } - // x <- L^-T x for (int i=nv-1; i > 0; i--) { // skip diagonal rows @@ -1831,8 +1822,6 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, x[i+offset] *= invD_i; } } - - mj_freeStack(d); } diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index ec8d168d..a710416c 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -71,9 +71,8 @@ MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, i // sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); -// TODO(tassa): Restore mjData const-ness. // sparse backsubstitution for one island: x = inv(L'*D*L)*x, use factorization in d -MJAPI void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* x, int island); +MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island); // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 0ece01ac..34fd4da9 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -770,7 +770,7 @@ static void mj_advance(const mjModel* m, mjData* d, // Euler integrator, semi-implicit in velocity, possibly skipping factorisation void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { TM_START; - int nv = m->nv, nM = m->nM; + int nv = m->nv, nC = m->nC; mj_markStack(d); mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum); @@ -794,22 +794,23 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { // damping: integrate implicitly else { if (!skipfactor) { - mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum); - - // MhB = M + h*diag(B) - mju_copy(MhB, d->qM, nM); + // qH = M + h*diag(B) + for (int i=0; i < nC; i++) { + d->qH[i] = d->qM[d->mapM2C[i]]; + } for (int i=0; i < nv; i++) { - MhB[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i]; + d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i]; } - // factor - mj_factorI(m, d, MhB, d->qH, d->qHDiagInv); + // factorize in-place + mj_factorIs(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } // solve mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); mju_copy(qacc, qfrc, m->nv); - mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv); + mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } // advance state and time @@ -938,7 +939,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) { // fully implicit in velocity, possibly skipping factorization void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { TM_START; - int nv = m->nv, nM = m->nM, nD = m->nD; + int nv = m->nv, nM = m->nM, nD = m->nD, nC = m->nC; mj_markStack(d); mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); @@ -985,13 +986,20 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // set MhB = M - dt*qDeriv mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM); - // factorize - mj_factorI(m, d, MhB, d->qH, d->qHDiagInv); + // copy into qH + for (int i=0; i < nC; i++) { + d->qH[i] = MhB[d->mapM2C[i]]; + } + + // factorize in-place + mj_factorIs(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc mju_copy(qacc, qfrc, nv); - mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv); + mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + } else { mjERROR("integrator must be implicit or implicitfast"); } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 377234b1..729a2e54 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1127,7 +1127,8 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printInertia("QM", d->qM, m, fp, float_format); - printInertia("QLD", d->qLD, m, fp, float_format); + printSparse("QLD", d->qLD, m->nv, d->C_rownnz, + d->C_rowadr, d->C_colind, fp, float_format); printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format); // B sparse structure diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index e5237a61..32068a77 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1087,38 +1087,25 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { int nv = m->nv; const mjtNum* qLD = d->qLD; - const int* dofMadr = m->dof_Madr; mju_zero(res, nv); // res = L * vec for (int i=0; i < nv; i++) { - // simple: diagonal - if (m->dof_simplenum[i]) { - res[i] = vec[i]; - } + // diagonal + res[i] = vec[i]; - // regular: full multiplication - else { - // diagonal - res[i] += vec[i]; - - // off-diagonal - int j = m->dof_parentid[i]; - int adr = dofMadr[i] + 1; - while (j >= 0) { - res[i] += qLD[adr]*vec[j]; - - // advance to next element - j = m->dof_parentid[j]; - adr++; - } + // 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, /*flg_unc1=*/0); } } // res *= sqrt(D) for (int i=0; i < nv; i++) { - res[i] *= mju_sqrt(qLD[dofMadr[i]]); + int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1; + res[i] *= mju_sqrt(qLD[diag]); } } diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c index 6d1acf19..3366e731 100644 --- a/src/engine/engine_vis_interact.c +++ b/src/engine/engine_vis_interact.c @@ -556,7 +556,8 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur // compute average spatial inertia at selection point for (int i=0; i < nv; i++) { - sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); + int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1; + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); } mj_jac(m, d, jac, NULL, selpos, sel); mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3); diff --git a/test/benchmark/factorI_benchmark_test.cc b/test/benchmark/factorI_benchmark_test.cc index 55118a99..c1dedbd9 100644 --- a/test/benchmark/factorI_benchmark_test.cc +++ b/test/benchmark/factorI_benchmark_test.cc @@ -43,21 +43,23 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) { // allocate inputs and outputs mj_markStack(d); - // CSR matrices - mjtNum* Ms = mj_stackAllocNum(d, m->nC); - mjtNum* LDs = mj_stackAllocNum(d, m->nC); + // M: mass matrix in CSR format + mjtNum* M = mj_stackAllocNum(d, m->nC); for (int i=0; i < m->nC; i++) { - Ms[i] = d->qM[d->mapM2C[i]]; + M[i] = d->qM[d->mapM2C[i]]; } + // LDlegacy: legacy LD matrix (size nM) + mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); + // benchmark while (state.KeepRunningBatch(kNumBenchmarkSteps)) { for (int i=0; i < kNumBenchmarkSteps; i++) { if (legacy) { - mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); + mj_factorI(m, d, d->qM, LDlegacy, d->qLDiagInv); } else { - mju_copy(LDs, Ms, m->nC); - mj_factorIs(LDs, d->qLDiagInv, m->nv, + mju_copy(d->qLD, M, m->nC); + mj_factorIs(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } } diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc index 7ddc3468..df230fc8 100644 --- a/test/benchmark/inertia_benchmark_test.cc +++ b/test/benchmark/inertia_benchmark_test.cc @@ -45,13 +45,15 @@ static void BM_solve(benchmark::State& state, SolveType type) { // allocate input and output vectors mj_markStack(d); - // make CSR matrix - mjtNum* Ms = mj_stackAllocNum(d, m->nC); - mjtNum* LDs = mj_stackAllocNum(d, m->nC); + // M: mass matrix in CSR format + mjtNum* M = mj_stackAllocNum(d, m->nC); for (int i=0; i < m->nC; i++) { - Ms[i] = d->qM[d->mapM2C[i]]; + M[i] = d->qM[d->mapM2C[i]]; } + // LDlegacy: legacy LD matrix (size nM) + mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); + // arbitrary input vector mjtNum *res = mj_stackAllocNum(d, m->nv); mjtNum *vec = mj_stackAllocNum(d, m->nv); @@ -62,17 +64,18 @@ static void BM_solve(benchmark::State& state, SolveType type) { // benchmark while (state.KeepRunningBatch(kNumBenchmarkSteps)) { for (int i=0; i < kNumBenchmarkSteps; i++) { + mju_copy(res, vec, m->nv); switch (type) { case SolveType::kLegacy: - mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); + mj_factorI(m, d, d->qM, LDlegacy, d->qLDiagInv); + mj_solveLD(m, res, 1, LDlegacy, d->qLDiagInv); mj_solveM(m, d, res, vec, 1); break; case SolveType::kCsr: - mju_copy(LDs, Ms, m->nC); - mj_factorIs(LDs, d->qLDiagInv, m->nv, + mju_copy(d->qLD, M, m->nC); + mj_factorIs(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); - mju_copy(res, vec, m->nv); - mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1, + mj_solveLDs(res, d->qLD, d->qLDiagInv, m->nv, 1, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } } diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc index 18ebe69b..434d5d2e 100644 --- a/test/benchmark/solveLD_benchmark_test.cc +++ b/test/benchmark/solveLD_benchmark_test.cc @@ -50,20 +50,21 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) { vec[i] = 0.2 + 0.3*i; } - // make CSR matrix - mjtNum* LDs = mj_stackAllocNum(d, m->nC); + // make legacy matrix + mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); + mju_zero(LDlegacy, m->nM); for (int i=0; i < m->nC; i++) { - LDs[i] = d->qLD[d->mapM2C[i]]; + LDlegacy[d->mapM2C[i]] = d->qLD[i]; } // benchmark while (state.KeepRunningBatch(kNumBenchmarkSteps)) { for (int i=0; i < kNumBenchmarkSteps; i++) { + mju_copy(res, vec, m->nv); if (featherstone) { - mj_solveM(m, d, res, vec, 1); + mj_solveLD(m, res, 1, LDlegacy, d->qLDiagInv); } else { - mju_copy(res, vec, m->nv); - mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1, + mj_solveLDs(res, d->qLD, d->qLDiagInv, m->nv, 1, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 47e1ff49..05e9348e 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -449,9 +449,8 @@ TEST_F(CoreSmoothTest, SolveMIsland) { // expect corresponding values to match for (int j=0; j < dofnum; j++) { - EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-14)); + EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12)); } - mju_free(res_i); } @@ -475,21 +474,21 @@ TEST_F(CoreSmoothTest, FactorI) { // dense L matrix int nv = model->nv; - vector Ldense(nv*nv); - mj_fullM(model, Ldense.data(), data->qLD); - // clear upper triangle, set diagonal to 1 + vector Ldense(nv*nv, 0); + mju_sparse2dense(Ldense.data(), data->qLD, nv, nv, + data->C_rownnz, data->C_rowadr, data->C_colind); for (int i=0; i < nv; i++) { - for (int j=i; j < nv; j++) { - Ldense[i*nv+j] = i == j ? 1 : 0; - } + // set diagonal to 1 + Ldense[i*nv+i] = 1; } // dense D matrix vector Ddense(nv*nv); - mj_fullM(model, Ddense.data(), data->qLD); - // clear everything but the diagonal + mju_sparse2dense(Ddense.data(), data->qLD, nv, nv, + data->C_rownnz, data->C_rowadr, data->C_colind); for (int i=0; i < nv; i++) { for (int j=0; j < nv; j++) { + // zero everything except the diagonal if (i != j) Ddense[i*nv+j] = 0; } } @@ -521,20 +520,21 @@ TEST_F(CoreSmoothTest, SolveLDs) { mj_forward(m, d); int nv = m->nv; + int nM = m->nM; int nC = m->nC; - // copy LD into LDs: CSR format - vector LDs(nC); + // copy M into LD: Legacy format + vector LDlegacy(nM, 0); for (int i=0; i < nC; i++) { - LDs[i] = d->qLD[d->mapM2C[i]]; + LDlegacy[d->mapM2C[i]] = d->qLD[i]; } // compare LD and LDs densified matrices vector LDdense(nv*nv); - mju_sparse2dense(LDdense.data(), LDs.data(), nv, nv, + mju_sparse2dense(LDdense.data(), d->qLD, nv, nv, d->C_rownnz, d->C_rowadr, d->C_colind); vector LDdense2(nv*nv); - mj_fullM(m, LDdense2.data(), d->qLD); + mj_fullM(m, LDdense2.data(), LDlegacy.data()); // expect lower triangles to match exactly for (int i=0; i < nv; i++) { @@ -543,14 +543,14 @@ TEST_F(CoreSmoothTest, SolveLDs) { } } - // compare LD and LDs vector solve + // compare legacy and CSR LD vector solve vector vec(nv); vector vec2(nv); for (int i=0; i < nv; i++) vec[i] = vec2[i] = 20 + 30*i; for (int i=0; i < nv; i+=2) vec[i] = vec2[i] = 0; - mj_solveLD(m, vec.data(), 1, d->qLD, d->qLDiagInv); - mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, 1, + mj_solveLD(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv); + mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, 1, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); // expect vectors to match up to floating point precision @@ -572,12 +572,13 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) { mj_forward(m, d); int nv = m->nv; + int nM = m->nM; int nC = m->nC; - // copy LD into LDs: CSR format - vector LDs(nC); + // copy LD into LDlegacy: Legacy format + vector LDlegacy(nM, 0); for (int i=0; i < nC; i++) { - LDs[i] = d->qLD[d->mapM2C[i]]; + LDlegacy[d->mapM2C[i]] = d->qLD[i]; } // compare n LD and LDs vector solve @@ -587,8 +588,8 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) { for (int i=0; i < nv*n; i++) vec[i] = vec2[i] = 2 + 3*i; for (int i=0; i < nv*n; i+=3) vec[i] = vec2[i] = 0; - mj_solveLD(m, vec.data(), n, d->qLD, d->qLDiagInv); - mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n, + mj_solveLD(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv); + mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); // expect vectors to match up to floating point precision @@ -609,19 +610,12 @@ TEST_F(CoreSmoothTest, SolveM2) { mjData* d = mj_makeData(m); mj_forward(m, d); - int nv = m->nv; - int nC = m->nC; - - // copy LD into LDs: CSR format - vector LDs(nC); - for (int i=0; i < nC; i++) { - LDs[i] = d->qLD[d->mapM2C[i]]; - } - // inverse square root of D from inertia LDL decomposition + int nv = m->nv; vector sqrtInvD(nv); for (int i=0; i < nv; i++) { - sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); + int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1; + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); } // compare full solve and half solve @@ -633,7 +627,7 @@ TEST_F(CoreSmoothTest, SolveM2) { vector res(nv*n); mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n); - mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n, + mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); // expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v) @@ -655,25 +649,32 @@ TEST_F(CoreSmoothTest, FactorIs) { mjData* d = mj_makeData(m); mj_forward(m, d); - int nC = m->nC, nv = m->nv; + int nC = m->nC, nM = m->nM, nv = m->nv; - // copy qM into LDs, qLD into qLDexpected: CSR format - vector qLDsExpected(nC); - vector qLDs(nC); + // copy qM into into qLDlegacy and factorize + vector qLDlegacy(nM); + mj_factorI(m, d, d->qM, qLDlegacy.data(), d->qLDiagInv); + + // copy qLDlegacy into qLDexpected: CSR format + vector qLDexpected(nC); for (int i=0; i < nC; i++) { - int index = d->mapM2C[i]; - qLDs[i] = d->qM[index]; // mj_factorIs is in-place - qLDsExpected[i] = d->qLD[index]; + qLDexpected[i] = qLDlegacy[d->mapM2C[i]]; // mj_factorIs is in-place + } + + // copy qM into qLD: CSR format + vector qLD(nC); + for (int i=0; i < nC; i++) { + qLD[i] = d->qM[d->mapM2C[i]]; // mj_factorIs is in-place } vector qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv); vector qLDiagInv(nv, 0); - mj_factorIs(qLDs.data(), qLDiagInv.data(), nv, + mj_factorIs(qLD.data(), qLDiagInv.data(), nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); // expect outputs to match to floating point precision - EXPECT_THAT(qLDs, Pointwise(DoubleNear(1e-12), qLDsExpected)); + EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected)); EXPECT_THAT(qLDiagInv, Pointwise(DoubleNear(1e-12), qLDiagInvExpected)); /* uncomment for debugging diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 1db7731a..7aaf6ad3 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -436,7 +436,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { Ac[i*nv + i] = -m->jnt_stiffness[i]; Ac[nv*nv + i*nv + i] = -m->dof_damping[i]; } - mj_solveLD(m, Ac, 2*nv, d->qH, d->qHDiagInv); + mj_solveLDs(Ac, d->qH, d->qHDiagInv, nv, 2*nv, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); // A = [dt*Ac; Ac] mju_transpose(A, Ac, 2*nv, nv); @@ -463,7 +464,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { mjtNum *BcT = mj_stackAllocNum(d, nv*nu); mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz, d->moment_rowadr, d->moment_colind); - mj_solveLD(m, Bc, nu, d->qH, d->qHDiagInv); + mj_solveLDs(Bc, d->qH, d->qHDiagInv, nv, nu, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); mju_transpose(BcT, Bc, nu, nv); mju_scl(B, BcT, dt*dt, nu*nv); mju_scl(B+nu*nv, BcT, dt, nu*nv); diff --git a/test/pipeline_test.cc b/test/pipeline_test.cc index f9c40713..6f1eb737 100644 --- a/test/pipeline_test.cc +++ b/test/pipeline_test.cc @@ -61,7 +61,8 @@ TEST_F(PipelineTest, SparseDenseEquivalent) { std::vector qacc_sparse = AsVector(data->qacc, model->nv); // expect accelerations to be insignificantly different - EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse)); + EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse)) + << "failed equivalence for solver=" << solver; } mj_deleteData(data);