diff --git a/doc/includes/references.h b/doc/includes/references.h index 5769ba30..c68ea413 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -309,6 +309,8 @@ struct mjData_ { int* D_rownnz; // non-zeros in each row (nv x 1) int* D_rowadr; // address of each row in D_colind (nv x 1) int* D_colind; // 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* B_rownnz; // non-zeros in each row (nbody x 1) int* B_rowadr; // address of each row in B_colind (nbody x 1) int* B_colind; // column indices of non-zeros (nB x 1) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 23a4f019..9495f4ae 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -337,6 +337,8 @@ struct mjData_ { int* D_rownnz; // non-zeros in each row (nv x 1) int* D_rowadr; // address of each row in D_colind (nv x 1) int* D_colind; // 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* B_rownnz; // non-zeros in each row (nbody x 1) int* B_rowadr; // address of each row in B_colind (nbody x 1) int* B_colind; // column indices of non-zeros (nB x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 0aa22805..cb289aff 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -643,6 +643,8 @@ X ( int, D_rownnz, nv, 1 ) \ X ( int, D_rowadr, nv, 1 ) \ X ( int, D_colind, nD, 1 ) \ + X ( int, mapM2D, nD, 1 ) \ + X ( int, mapD2M, nM, 1 ) \ X ( int, B_rownnz, nbody, 1 ) \ X ( int, B_rowadr, nbody, 1 ) \ X ( int, B_colind, nB, 1 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index be3e6ad7..d7ebbd6d 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -4866,6 +4866,20 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='column indices of non-zeros (nD x 1)', # pylint: disable=line-too-long ), + StructFieldDecl( + name='mapM2D', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='index mapping from M to D (nD x 1)', # pylint: disable=line-too-long + ), + StructFieldDecl( + name='mapD2M', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='index mapping from D to M (nM x 1)', # pylint: disable=line-too-long + ), StructFieldDecl( name='B_rownnz', type=PointerType( diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 1a43c1d9..16d5d182 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -792,7 +792,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { mjtNum* MhB = mj_stackAllocNum(d, nM); // MhB = M + h*diag(B) - mju_copy(MhB, d->qM, m->nM); + mju_copy(MhB, d->qM, nM); for (int i=0; i < nv; i++) { MhB[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i]; } @@ -933,7 +933,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; + int nv = m->nv, nM = m->nM, nD = m->nD; mj_markStack(d); mjtNum* qfrc = mj_stackAllocNum(d, nv); @@ -949,7 +949,9 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { mjd_smooth_vel(m, d, /* flg_bias = */ 1); // set qLU = qM - mj_copyM2DSparse(m, d, d->qLU, d->qM); + for (int i=0; i < nD; i++) { + d->qLU[i] = d->qM[d->mapM2D[i]]; + } // set qLU = qM - dt*qDeriv mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); @@ -970,18 +972,20 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { mjd_smooth_vel(m, d, /* flg_bias = */ 0); // modified mass matrix MhB = qDeriv[Lower] - mjtNum* MhB = mj_stackAllocNum(d, m->nM); - mj_copyD2MSparse(m, d, MhB, d->qDeriv); + mjtNum* MhB = mj_stackAllocNum(d, nM); + for (int i=0; i < nM; i++) { + MhB[i] = d->qDeriv[d->mapD2M[i]]; + } // set MhB = M - dt*qDeriv - mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, m->nM); + mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM); // factorize mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL); } // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc - mju_copy(qacc, qfrc, m->nv); + mju_copy(qacc, qfrc, nv); mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv); } else { mjERROR("integrator must be implicit or implicitfast"); diff --git a/src/engine/engine_inverse.c b/src/engine/engine_inverse.c index f237b6af..8805f8e4 100644 --- a/src/engine/engine_inverse.c +++ b/src/engine/engine_inverse.c @@ -71,7 +71,7 @@ void mj_invVelocity(const mjModel* m, mjData* d) { // convert discrete-time qacc to continuous-time qacc static void mj_discreteAcc(const mjModel* m, mjData* d) { - int nv = m->nv, dof_damping; + int nv = m->nv, nM = m->nM, nD = m->nD, dof_damping; mjtNum *qacc = d->qacc; mj_markStack(d); @@ -114,7 +114,9 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) { mjd_smooth_vel(m, d, /* flg_bias = */ 1); // set qLU = qM - mj_copyM2DSparse(m, d, d->qLU, d->qM); + for (int i=0; i < nD; i++) { + d->qLU[i] = d->qM[d->mapM2D[i]]; + } // set qLU = qM - dt*qDeriv mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); @@ -134,7 +136,9 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) { // set M = M - dt*qDeriv (reduced to M nonzeros) mjtNum* qDerivReduced = mj_stackAllocNum(d, m->nM); - mj_copyD2MSparse(m, d, qDerivReduced, d->qDeriv); + for (int i=0; i < nM; i++) { + qDerivReduced[i] = d->qDeriv[d->mapD2M[i]]; + } mju_addToScl(d->qM, qDerivReduced, -m->opt.timestep, m->nM); // set qfrc = (M - dt*qDeriv) * qacc diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 07a72c7c..efe66939 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1059,6 +1059,107 @@ static void checkDBSparse(const mjModel* m, mjData* d) { +// integer valued dst[D] = src[M], handle different sparsity representations +static void copyM2DSparse(const mjModel* m, mjData* d, int* dst, const int* src) { + int nv = m->nv; + mj_markStack(d); + + // init remaining + int* remaining = mj_stackAllocInt(d, nv); + mju_copyInt(remaining, d->D_rownnz, nv); + + // copy data + for (int i = nv - 1; i >= 0; i--) { + // init at diagonal + int adr = m->dof_Madr[i]; + remaining[i]--; + dst[d->D_rowadr[i] + remaining[i]] = src[adr]; + adr++; + + // process below diagonal + int j = i; + while ((j = m->dof_parentid[j]) >= 0) { + remaining[i]--; + dst[d->D_rowadr[i] + remaining[i]] = src[adr]; + + remaining[j]--; + dst[d->D_rowadr[j] + remaining[j]] = src[adr]; + + adr++; + } + } + + // check that none remaining + for (int i=0; i < nv; i++) { + if (remaining[i]) { + mjERROR("unassigned index"); + } + } + + mj_freeStack(d); +} + + + +// integer valued dst[M] = src[D lower], handle different sparsity representations +static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int* src) { + int nv = m->nv; + + // copy data + for (int i = nv - 1; i >= 0; i--) { + // find diagonal in qDeriv + int j = 0; + while (d->D_colind[d->D_rowadr[i] + j] < i) { + j++; + } + + // copy + int adr = m->dof_Madr[i]; + while (j >= 0) { + dst[adr] = src[d->D_rowadr[i] + j]; + adr++; + j--; + } + } +} + + + +// construct index mappings between D <-> M +static void makeDmap(const mjModel* m, mjData* d) { + int nM = m->nM, nD = m->nD; + mj_markStack(d); + + // make mapM2D + int* M = mj_stackAllocInt(d, nM); + for (int i=0; i < nM; i++) M[i] = i; + for (int i=0; i < nD; i++) d->mapM2D[i] = -1; + copyM2DSparse(m, d, d->mapM2D, M); + + // check that all indices are filled in + for (int i=0; i < nD; i++) { + if (d->mapM2D[i] < 0) { + mjERROR("unassigned index in mapM2D"); + } + } + + // make mapD2M + int* D = mj_stackAllocInt(d, nD); + for (int i=0; i < nD; i++) D[i] = i; + for (int i=0; i < nM; i++) d->mapD2M[i] = -1; + copyD2MSparse(m, d, d->mapD2M, D); + + // check that all indices are filled in + for (int i=0; i < nM; i++) { + if (d->mapD2M[i] < 0) { + mjERROR("unassigned index in mapD2M"); + } + } + + mj_freeStack(d); +} + + //----------------------------------- mjData construction ------------------------------------------ // set pointers into mjData buffer @@ -1707,6 +1808,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { makeDSparse(m, d); makeBSparse(m, d); checkDBSparse(m, d); + makeDmap(m, d); } // restore pluginstate and plugindata diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 3202c7a2..019baaa4 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1000,6 +1000,20 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, } fprintf(fp, "\n\n"); + // mapM2D + fprintf(fp, NAME_FORMAT, "mapM2D"); + for (int i = 0; i < m->nD; i++) { + fprintf(fp, " %d", d->mapM2D[i]); + } + fprintf(fp, "\n\n"); + + // mapD2M + fprintf(fp, NAME_FORMAT, "mapD2M"); + for (int i = 0; i < m->nM; i++) { + fprintf(fp, " %d", d->mapD2M[i]); + } + fprintf(fp, "\n\n"); + // B_rownnz fprintf(fp, NAME_FORMAT, "B_rownnz"); for (int i = 0; i < m->nbody; i++) { diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 06c4380d..3e6dc212 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1285,66 +1285,6 @@ void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst) { } -//-------------------------- sparse system matrix conversion --------------------------------------- - -// dst[D] = src[M], handle different sparsity representations -void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) { - int nv = m->nv; - mj_markStack(d); - - // init remaining - int* remaining = mj_stackAllocInt(d, nv); - mju_copyInt(remaining, d->D_rownnz, nv); - - // copy data - for (int i = nv - 1; i >= 0; i--) { - // init at diagonal - int adr = m->dof_Madr[i]; - remaining[i]--; - dst[d->D_rowadr[i] + remaining[i]] = src[adr]; - adr++; - - // process below diagonal - int j = i; - while ((j = m->dof_parentid[j]) >= 0) { - remaining[i]--; - dst[d->D_rowadr[i] + remaining[i]] = src[adr]; - - remaining[j]--; - dst[d->D_rowadr[j] + remaining[j]] = src[adr]; - - adr++; - } - } - - mj_freeStack(d); -} - - - -// dst[M] = src[D lower], handle different sparsity representations -void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) { - int nv = m->nv; - - // copy data - for (int i = nv - 1; i >= 0; i--) { - // find diagonal in qDeriv - int j = 0; - while (d->D_colind[d->D_rowadr[i] + j] < i) { - j++; - } - - // copy - int adr = m->dof_Madr[i]; - while (j >= 0) { - dst[adr] = src[d->D_rowadr[i] + j]; - adr++; - j--; - } - } -} - - //-------------------------- perturbations --------------------------------------------------------- diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index f7cdf9b8..057fc6e2 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -148,15 +148,6 @@ MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst, MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst); -//-------------------------- sparse system matrix conversion --------------------------------------- - -// dst[D] = src[M], handle different sparsity representations -MJAPI void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src); - -// dst[M] = src[D lower], handle different sparsity representations -MJAPI void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src); - - //-------------------------- perturbations --------------------------------------------------------- // apply Cartesian force and torque diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index f0cc7180..96f4ed74 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4924,6 +4924,8 @@ public unsafe struct mjData_ { public int* D_rownnz; public int* D_rowadr; public int* D_colind; + public int* mapM2D; + public int* mapD2M; public int* B_rownnz; public int* B_rowadr; public int* B_colind;