Move square root of Delassus matrix from stack to arena
PiperOrigin-RevId: 916852244 Change-Id: I4379553f808d0b238a1421606f806ea2d0a33d7c
This commit is contained in:
committed by
Copybara-Service
parent
2345663efb
commit
04042d8bf3
@@ -191,6 +191,7 @@ struct mjData_ {
|
||||
int nl; // number of limit constraints
|
||||
int nefc; // number of constraints
|
||||
int nJ; // number of non-zeros in constraint Jacobian
|
||||
int nY; // number of non-zeros in constraint inverse inertia square root
|
||||
int nA; // number of non-zeros in constraint inverse inertia matrix
|
||||
int nisland; // number of detected constraint islands
|
||||
int nidof; // number of dofs in all islands
|
||||
@@ -445,8 +446,12 @@ struct mjData_ {
|
||||
mjtNum* iefc_R; // inverse constraint mass (nefc x 1)
|
||||
|
||||
// computed by mj_projectConstraint (PGS solver)
|
||||
int* efc_Y_rownnz; // number of non-zeros in Y row (nefc x 1)
|
||||
int* efc_Y_rowadr; // row start address in Y colind array (nefc x 1)
|
||||
int* efc_Y_colind; // column indices in sparse Y (nY x 1)
|
||||
mjtNum* efc_Y; // whitened Jacobian Y = J*M^(-1/2) (nY x 1)
|
||||
int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1)
|
||||
int* efc_AR_rowadr; // row start address in colind array (nefc x 1)
|
||||
int* efc_AR_rowadr; // row start address in AR colind array (nefc x 1)
|
||||
int* efc_AR_colind; // column indices in sparse AR (nA x 1)
|
||||
mjtNum* efc_AR; // J*inv(M)*J' + R (nA x 1)
|
||||
|
||||
|
||||
@@ -225,6 +225,7 @@ struct mjData_ {
|
||||
int nl; // number of limit constraints
|
||||
int nefc; // number of constraints
|
||||
int nJ; // number of non-zeros in constraint Jacobian
|
||||
int nY; // number of non-zeros in constraint inverse inertia square root
|
||||
int nA; // number of non-zeros in constraint inverse inertia matrix
|
||||
int nisland; // number of detected constraint islands
|
||||
int nidof; // number of dofs in all islands
|
||||
@@ -479,8 +480,12 @@ struct mjData_ {
|
||||
mjtNum* iefc_R; // inverse constraint mass (nefc x 1)
|
||||
|
||||
// computed by mj_projectConstraint (PGS solver)
|
||||
int* efc_Y_rownnz; // number of non-zeros in Y row (nefc x 1)
|
||||
int* efc_Y_rowadr; // row start address in Y colind array (nefc x 1)
|
||||
int* efc_Y_colind; // column indices in sparse Y (nY x 1)
|
||||
mjtNum* efc_Y; // whitened Jacobian Y = J*M^(-1/2) (nY x 1)
|
||||
int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1)
|
||||
int* efc_AR_rowadr; // row start address in colind array (nefc x 1)
|
||||
int* efc_AR_rowadr; // row start address in AR colind array (nefc x 1)
|
||||
int* efc_AR_colind; // column indices in sparse AR (nA x 1)
|
||||
mjtNum* efc_AR; // J*inv(M)*J' + R (nA x 1)
|
||||
|
||||
|
||||
@@ -946,6 +946,10 @@
|
||||
|
||||
// array fields of mjData that are used in the dual problem
|
||||
#define MJDATA_ARENA_POINTERS_DUAL \
|
||||
XNV( int, efc_Y_rownnz, MJ_D(nefc), 1 ) \
|
||||
XNV( int, efc_Y_rowadr, MJ_D(nefc), 1 ) \
|
||||
XNV( int, efc_Y_colind, MJ_D(nY), 1 ) \
|
||||
XNV( mjtNum, efc_Y, MJ_D(nY), 1 ) \
|
||||
XNV( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \
|
||||
XNV( int, efc_AR_rowadr, MJ_D(nefc), 1 ) \
|
||||
XNV( int, efc_AR_colind, MJ_D(nA), 1 ) \
|
||||
@@ -1020,6 +1024,7 @@
|
||||
X( int, nl ) \
|
||||
X( int, nefc ) \
|
||||
X( int, nJ ) \
|
||||
X( int, nY ) \
|
||||
X( int, nA ) \
|
||||
X( int, nisland ) \
|
||||
X( int, nidof ) \
|
||||
|
||||
@@ -5492,6 +5492,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
|
||||
type=ValueType(name='int'),
|
||||
doc='number of non-zeros in constraint Jacobian',
|
||||
),
|
||||
StructFieldDecl(
|
||||
name='nY',
|
||||
type=ValueType(name='int'),
|
||||
doc='number of non-zeros in constraint inverse inertia square root', # pylint: disable=line-too-long
|
||||
),
|
||||
StructFieldDecl(
|
||||
name='nA',
|
||||
type=ValueType(name='int'),
|
||||
@@ -6726,6 +6731,38 @@ STRUCTS: Mapping[str, StructDecl] = dict([
|
||||
doc='inverse constraint mass',
|
||||
array_extent=('nefc',),
|
||||
),
|
||||
StructFieldDecl(
|
||||
name='efc_Y_rownnz',
|
||||
type=PointerType(
|
||||
inner_type=ValueType(name='int'),
|
||||
),
|
||||
doc='number of non-zeros in Y row',
|
||||
array_extent=('nefc',),
|
||||
),
|
||||
StructFieldDecl(
|
||||
name='efc_Y_rowadr',
|
||||
type=PointerType(
|
||||
inner_type=ValueType(name='int'),
|
||||
),
|
||||
doc='row start address in Y colind array',
|
||||
array_extent=('nefc',),
|
||||
),
|
||||
StructFieldDecl(
|
||||
name='efc_Y_colind',
|
||||
type=PointerType(
|
||||
inner_type=ValueType(name='int'),
|
||||
),
|
||||
doc='column indices in sparse Y',
|
||||
array_extent=('nY',),
|
||||
),
|
||||
StructFieldDecl(
|
||||
name='efc_Y',
|
||||
type=PointerType(
|
||||
inner_type=ValueType(name='mjtNum'),
|
||||
),
|
||||
doc='whitened Jacobian Y = J*M^(-1/2)',
|
||||
array_extent=('nY',),
|
||||
),
|
||||
StructFieldDecl(
|
||||
name='efc_AR_rownnz',
|
||||
type=PointerType(
|
||||
@@ -6739,7 +6776,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
|
||||
type=PointerType(
|
||||
inner_type=ValueType(name='int'),
|
||||
),
|
||||
doc='row start address in colind array',
|
||||
doc='row start address in AR colind array',
|
||||
array_extent=('nefc',),
|
||||
),
|
||||
StructFieldDecl(
|
||||
|
||||
@@ -812,6 +812,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
|
||||
X(ne);
|
||||
X(nf);
|
||||
X(nJ);
|
||||
X(nY);
|
||||
X(nA);
|
||||
X(nefc);
|
||||
X(nisland);
|
||||
@@ -891,6 +892,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
||||
X(ne);
|
||||
X(nf);
|
||||
X(nJ);
|
||||
X(nY);
|
||||
X(nA);
|
||||
X(nefc);
|
||||
X(nisland);
|
||||
|
||||
+197
-181
@@ -2591,12 +2591,137 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
|
||||
}
|
||||
|
||||
|
||||
// pre-count Y_rownnz, Y_rowadr, return total nonzeros nY
|
||||
// Y has the sparsity of J * inv(L'), where L is the Cholesky factor of M
|
||||
static int computeY_precount(int* Y_rownnz, int* Y_rowadr, int nefc, int nv,
|
||||
const int* J_rownnz, const int* J_rowadr, const int* J_colind,
|
||||
const int* M_rownnz, const int* M_rowadr, const int* M_colind,
|
||||
int* marker) {
|
||||
mju_fillInt(marker, -1, nv);
|
||||
|
||||
Y_rowadr[0] = 0;
|
||||
for (int r=0; r < nefc; r++) {
|
||||
int nnz = 0; // nonzeros in row r of Y
|
||||
|
||||
// traverse row r of J in reverse, count unique nonzeros
|
||||
int start = J_rowadr[r];
|
||||
int end = start + J_rownnz[r];
|
||||
for (int i=end-1; i >= start; i--) {
|
||||
int j = J_colind[i];
|
||||
|
||||
// if dof j is marked, it was already counted by a child dof: skip it
|
||||
if (marker[j] == r) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// traverse row j of M, marking new unique nonzeros
|
||||
int nnzM = M_rownnz[j];
|
||||
int adrM = M_rowadr[j];
|
||||
for (int k=0; k < nnzM; k++) {
|
||||
int c = M_colind[adrM + k];
|
||||
if (marker[c] != r) {
|
||||
marker[c] = r;
|
||||
nnz++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// update rownnz and rowadr
|
||||
Y_rownnz[r] = nnz;
|
||||
if (r < nefc - 1) {
|
||||
Y_rowadr[r+1] = Y_rowadr[r] + nnz;
|
||||
}
|
||||
}
|
||||
|
||||
// total non-zeros in Y
|
||||
return Y_rowadr[nefc-1] + Y_rownnz[nefc-1];
|
||||
}
|
||||
|
||||
|
||||
// fill Y column indices and values from J, chaining up the kinematic tree
|
||||
static void computeY_fill(mjtNum* Y, int* Y_colind,
|
||||
const int* Y_rownnz, const int* Y_rowadr, int nefc,
|
||||
const mjtNum* J, const int* J_rownnz, const int* J_rowadr,
|
||||
const int* J_colind, const int* dof_parentid) {
|
||||
for (int r=0; r < nefc; r++) {
|
||||
// init row
|
||||
int end = Y_rowadr[r] + Y_rownnz[r];
|
||||
int adrJ = J_rowadr[r];
|
||||
int remainJ = J_rownnz[r];
|
||||
int nnzY = 0;
|
||||
|
||||
// complete chain in reverse
|
||||
while (1) {
|
||||
// get previous dof in src and dst
|
||||
int prev_src = (remainJ > 0 ? J_colind[adrJ + remainJ - 1] : -1);
|
||||
int prev_dst = (nnzY > 0 ? dof_parentid[Y_colind[end - nnzY]] : -1);
|
||||
|
||||
// both finished: break
|
||||
if (prev_src < 0 && prev_dst < 0) {
|
||||
break;
|
||||
}
|
||||
|
||||
// add src
|
||||
else if (prev_src >= prev_dst) {
|
||||
nnzY++;
|
||||
remainJ--;
|
||||
Y_colind[end - nnzY] = prev_src;
|
||||
Y[end - nnzY] = J[adrJ + remainJ];
|
||||
}
|
||||
|
||||
// add dst
|
||||
else {
|
||||
nnzY++;
|
||||
Y_colind[end - nnzY] = prev_dst;
|
||||
Y[end - nnzY] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// compare with Y_rownnz: SHOULD NOT OCCUR
|
||||
if (nnzY != Y_rownnz[r]) {
|
||||
mjERROR("pre and post-count of Y_rownnz are not equal on row %d", r);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// in-place sparse back-substitution: Y <- Y * M^{-1/2}
|
||||
static void computeY_backsub(mjtNum* Y, const int* Y_rownnz, const int* Y_rowadr,
|
||||
const int* Y_colind, int nefc,
|
||||
const mjtNum* qLD, const int* M_rownnz, const int* M_rowadr,
|
||||
const int* M_colind, const mjtNum* sqrtInvD) {
|
||||
for (int r=0; r < nefc; r++) {
|
||||
int nnzY = Y_rownnz[r];
|
||||
int adrY = Y_rowadr[r];
|
||||
|
||||
// Y(r,:) <- inv(L') * Y(r,:), exploit sparsity of input vector
|
||||
for (int i=adrY + nnzY-1; i >= adrY; i--) {
|
||||
mjtNum val = Y[i];
|
||||
if (val == 0) {
|
||||
continue;
|
||||
}
|
||||
int j = Y_colind[i];
|
||||
int adrM = M_rowadr[j];
|
||||
mju_addToSclSparseInc(Y + adrY, qLD + adrM,
|
||||
nnzY, Y_colind + adrY,
|
||||
M_rownnz[j]-1, M_colind + adrM, -val);
|
||||
}
|
||||
|
||||
// Y(r,:) <- sqrt(inv(D)) * Y(r,:)
|
||||
for (int i=adrY; i < adrY + nnzY; i++) {
|
||||
int j = Y_colind[i];
|
||||
Y[i] *= sqrtInvD[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
//---------------------------- top-level API for constraint construction ---------------------------
|
||||
|
||||
// driver: call all functions above
|
||||
void mj_makeConstraint(const mjModel* m, mjData* d) {
|
||||
// clear sizes
|
||||
d->ne = d->nf = d->nl = d->nefc = d->nJ = d->nA = 0;
|
||||
d->ne = d->nf = d->nl = d->nefc = d->nJ = d->nA = d->nY = 0;
|
||||
|
||||
// disabled or Jacobian not allocated: return
|
||||
if (mjDISABLED(mjDSBL_CONSTRAINT)) {
|
||||
@@ -2705,177 +2830,60 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
|
||||
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
|
||||
}
|
||||
|
||||
// sparse
|
||||
// sparse Y = backsubM2(J')' and its transpose
|
||||
if (mj_isSparse(m)) {
|
||||
// compute B = backsubM2(J')' and its transpose
|
||||
|
||||
|
||||
// === pre-count B_rownnz, B_rowadr, nB (total nonzeros)
|
||||
|
||||
// allocate B rownnz and rowadr
|
||||
int* B_rownnz = mjSTACKALLOC(d, nefc, int);
|
||||
int* B_rowadr = mjSTACKALLOC(d, nefc, int);
|
||||
// arena-allocate Y rownnz and rowadr
|
||||
d->efc_Y_rownnz = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int));
|
||||
d->efc_Y_rowadr = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int));
|
||||
if (!d->efc_Y_rownnz || !d->efc_Y_rowadr) {
|
||||
mj_warning(d, mjWARN_CNSTRFULL, d->narena);
|
||||
mj_clearEfc(d);
|
||||
d->parena = d->ncon * sizeof(mjContact);
|
||||
mj_freeStack(d);
|
||||
return;
|
||||
}
|
||||
|
||||
// markers for merged dofs, initialized to -1
|
||||
int* marker = mjSTACKALLOC(d, nv, int);
|
||||
mju_fillInt(marker, -1, nv);
|
||||
|
||||
B_rowadr[0] = 0;
|
||||
for (int r=0; r < nefc; r++) {
|
||||
// supernode: same sparsity as previous row
|
||||
if (r > 0 && d->efc_J_rowsuper[r-1] > 0) {
|
||||
B_rownnz[r] = B_rownnz[r-1];
|
||||
}
|
||||
// pre-count Y_rownnz, Y_rowadr, nY (total nonzeros)
|
||||
d->nY = computeY_precount(d->efc_Y_rownnz, d->efc_Y_rowadr, nefc, nv,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
|
||||
m->M_rownnz, m->M_rowadr, m->M_colind, marker);
|
||||
|
||||
// first row in supernode block: full chain traversal
|
||||
else {
|
||||
int nnz = 0;
|
||||
|
||||
// traverse row r of J in reverse, count unique nonzeros
|
||||
int start = d->efc_J_rowadr[r];
|
||||
int end = start + d->efc_J_rownnz[r];
|
||||
for (int i=end-1; i >= start; i--) {
|
||||
int j = d->efc_J_colind[i];
|
||||
|
||||
// if dof j is marked, it was already counted by a child dof: skip it
|
||||
if (marker[j] == r) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// traverse row j of M, marking new unique nonzeros
|
||||
int nnzM = m->M_rownnz[j];
|
||||
int adrM = m->M_rowadr[j];
|
||||
for (int k=0; k < nnzM; k++) {
|
||||
int c = m->M_colind[adrM + k];
|
||||
if (marker[c] != r) {
|
||||
marker[c] = r;
|
||||
nnz++;
|
||||
}
|
||||
}
|
||||
}
|
||||
B_rownnz[r] = nnz;
|
||||
}
|
||||
|
||||
// update rowadr
|
||||
if (r < nefc - 1) {
|
||||
B_rowadr[r+1] = B_rowadr[r] + B_rownnz[r];
|
||||
}
|
||||
// arena-allocate values and column indices
|
||||
d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum));
|
||||
d->efc_Y_colind = mj_arenaAllocByte(d, sizeof(int) * d->nY, _Alignof(int));
|
||||
if (!d->efc_Y || !d->efc_Y_colind) {
|
||||
mj_warning(d, mjWARN_CNSTRFULL, d->narena);
|
||||
mj_clearEfc(d);
|
||||
d->parena = d->ncon * sizeof(mjContact);
|
||||
mj_freeStack(d);
|
||||
return;
|
||||
}
|
||||
|
||||
// total non-zeros in B
|
||||
int nB = B_rowadr[nefc-1] + B_rownnz[nefc-1];
|
||||
// fill in Y column indices, copy values from J
|
||||
computeY_fill(d->efc_Y, d->efc_Y_colind, d->efc_Y_rownnz, d->efc_Y_rowadr, nefc,
|
||||
d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
|
||||
m->dof_parentid);
|
||||
|
||||
|
||||
// === fill in B column indices, copy values from J
|
||||
// in-place sparse back-substitution: Y <- Y * M^-1/2
|
||||
computeY_backsub(d->efc_Y, d->efc_Y_rownnz, d->efc_Y_rowadr,
|
||||
d->efc_Y_colind, nefc,
|
||||
d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind, sqrtInvD);
|
||||
|
||||
// allocate values and column indices
|
||||
mjtNum* B = mjSTACKALLOC(d, nB, mjtNum);
|
||||
int* B_colind = mjSTACKALLOC(d, nB, int);
|
||||
// Y supernodes are identical to J supernodes
|
||||
const int* Y_rowsuper = d->efc_J_rowsuper;
|
||||
|
||||
for (int r=0; r < nefc; r++) {
|
||||
// supernode: copy column indices, only update values from J
|
||||
if (r > 0 && d->efc_J_rowsuper[r-1] > 0) {
|
||||
int prevAdr = B_rowadr[r-1];
|
||||
int adrB = B_rowadr[r];
|
||||
int nnzB = B_rownnz[r];
|
||||
mju_copyInt(B_colind + adrB, B_colind + prevAdr, nnzB);
|
||||
mju_zero(B + adrB, nnzB);
|
||||
|
||||
// copy J values into correct positions
|
||||
int adrJ = d->efc_J_rowadr[r];
|
||||
int jnnz = d->efc_J_rownnz[r];
|
||||
int bi = 0, ji = 0;
|
||||
while (ji < jnnz && bi < nnzB) {
|
||||
if (B_colind[adrB+bi] == d->efc_J_colind[adrJ+ji]) {
|
||||
B[adrB+bi] = d->efc_J[adrJ+ji];
|
||||
bi++;
|
||||
ji++;
|
||||
} else {
|
||||
bi++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// first row in supernode block: full chain completion
|
||||
else {
|
||||
int end = B_rowadr[r] + B_rownnz[r];
|
||||
int adrJ = d->efc_J_rowadr[r];
|
||||
int remainJ = d->efc_J_rownnz[r];
|
||||
int nnzB = 0;
|
||||
|
||||
// complete chain in reverse
|
||||
while (1) {
|
||||
// get previous dof in src and dst
|
||||
int prev_src = (remainJ > 0 ? d->efc_J_colind[adrJ + remainJ - 1] : -1);
|
||||
int prev_dst = (nnzB > 0 ? m->dof_parentid[B_colind[end - nnzB]] : -1);
|
||||
|
||||
// both finished: break
|
||||
if (prev_src < 0 && prev_dst < 0) {
|
||||
break;
|
||||
}
|
||||
|
||||
// add src
|
||||
else if (prev_src >= prev_dst) {
|
||||
nnzB++;
|
||||
remainJ--;
|
||||
B_colind[end - nnzB] = prev_src;
|
||||
B[end - nnzB] = d->efc_J[adrJ + remainJ];
|
||||
}
|
||||
|
||||
// add dst
|
||||
else {
|
||||
nnzB++;
|
||||
B_colind[end - nnzB] = prev_dst;
|
||||
B[end - nnzB] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// compare with B_rownnz: SHOULD NOT OCCUR
|
||||
if (nnzB != B_rownnz[r]) {
|
||||
mjERROR("pre and post-count of B_rownnz are not equal on row %d", r);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// === in-place sparse back-substitution: B <- B * M^-1/2
|
||||
|
||||
// sparse backsubM2 (half of LD back-substitution)
|
||||
for (int r=0; r < nefc; r++) {
|
||||
int nnzB = B_rownnz[r];
|
||||
int adrB = B_rowadr[r];
|
||||
|
||||
// B(r,:) <- inv(L') * B(r,:), exploit sparsity of input vector
|
||||
for (int i=adrB + nnzB-1; i >= adrB; i--) {
|
||||
mjtNum b = B[i];
|
||||
if (b == 0) {
|
||||
continue;
|
||||
}
|
||||
int j = B_colind[i];
|
||||
int adrC = m->M_rowadr[j];
|
||||
mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
|
||||
nnzB, B_colind + adrB,
|
||||
m->M_rownnz[j]-1, m->M_colind + adrC, -b);
|
||||
}
|
||||
|
||||
// B(r,:) <- sqrt(inv(D)) * B(r,:)
|
||||
for (int i=adrB; i < adrB + nnzB; i++) {
|
||||
int j = B_colind[i];
|
||||
B[i] *= sqrtInvD[j];
|
||||
}
|
||||
}
|
||||
|
||||
// B supernodes are identical to J supernodes
|
||||
const int* B_rowsuper = d->efc_J_rowsuper;
|
||||
|
||||
// construct B transposed
|
||||
int* BT_rownnz = mjSTACKALLOC(d, nv, int);
|
||||
int* BT_rowadr = mjSTACKALLOC(d, nv, int);
|
||||
int* BT_colind = mjSTACKALLOC(d, nB, int);
|
||||
mjtNum* BT = mjSTACKALLOC(d, nB, mjtNum);
|
||||
mju_transposeSparse(BT, B, nefc, nv,
|
||||
BT_rownnz, BT_rowadr, BT_colind, NULL,
|
||||
B_rownnz, B_rowadr, B_colind);
|
||||
// construct Y transposed
|
||||
int* YT_rownnz = mjSTACKALLOC(d, nv, int);
|
||||
int* YT_rowadr = mjSTACKALLOC(d, nv, int);
|
||||
int* YT_colind = mjSTACKALLOC(d, d->nY, int);
|
||||
mjtNum* YT = mjSTACKALLOC(d, d->nY, mjtNum);
|
||||
mju_transposeSparse(YT, d->efc_Y, nefc, nv,
|
||||
YT_rownnz, YT_rowadr, YT_colind, NULL,
|
||||
d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind);
|
||||
|
||||
// allocate AR row nonzeros and addresses on arena
|
||||
d->efc_AR_rownnz = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int));
|
||||
@@ -2891,8 +2899,8 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
|
||||
int* diagind = mjSTACKALLOC(d, nefc, int);
|
||||
d->nA = mju_sqrMatTDSparseSymbolic(
|
||||
d->efc_AR_rownnz, d->efc_AR_rowadr, NULL, diagind,
|
||||
nv, nefc, BT_rownnz, BT_rowadr, BT_colind,
|
||||
B_rownnz, B_rowadr, B_colind, B_rowsuper, d);
|
||||
nv, nefc, YT_rownnz, YT_rowadr, YT_colind,
|
||||
d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind, Y_rowsuper, d);
|
||||
|
||||
// allocate A values and column indices on arena
|
||||
d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum));
|
||||
@@ -2905,17 +2913,18 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
|
||||
return;
|
||||
}
|
||||
|
||||
// A = B * B': symbolic phase
|
||||
// A = Y * Y': symbolic phase
|
||||
mju_sqrMatTDSparseSymbolic(
|
||||
d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, diagind,
|
||||
nv, nefc, BT_rownnz, BT_rowadr, BT_colind,
|
||||
B_rownnz, B_rowadr, B_colind, B_rowsuper, d);
|
||||
nv, nefc, YT_rownnz, YT_rowadr, YT_colind,
|
||||
d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind, Y_rowsuper, d);
|
||||
|
||||
// A = B * B': numeric phase
|
||||
// A = Y * Y': numeric phase
|
||||
mju_sqrMatTDSparseNumeric(
|
||||
d->efc_AR, nefc, d->efc_AR_rownnz, d->efc_AR_rowadr,
|
||||
d->efc_AR_colind, diagind, BT, BT_rownnz, BT_rowadr,
|
||||
BT_colind, B, B_rownnz, B_rowadr, B_colind, B_rowsuper, NULL, d);
|
||||
d->efc_AR_colind, diagind, YT, YT_rownnz, YT_rowadr,
|
||||
YT_colind, d->efc_Y, d->efc_Y_rownnz, d->efc_Y_rowadr,
|
||||
d->efc_Y_colind, Y_rowsuper, NULL, d);
|
||||
|
||||
// AR = A + diag(R)
|
||||
for (int i=0; i < nefc; i++) {
|
||||
@@ -2923,11 +2932,24 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
|
||||
}
|
||||
}
|
||||
|
||||
// dense
|
||||
// dense Y = backsubM2(J')' and its transpose
|
||||
else {
|
||||
d->nA = nefc * nefc;
|
||||
// arena-allocate efc_Y
|
||||
d->nY = nefc * nv;
|
||||
d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum));
|
||||
if (!d->efc_Y) {
|
||||
mj_warning(d, mjWARN_CNSTRFULL, d->narena);
|
||||
mj_clearEfc(d);
|
||||
d->parena = d->ncon * sizeof(mjContact);
|
||||
mj_freeStack(d);
|
||||
return;
|
||||
}
|
||||
|
||||
// Y = backsubM2(J')'
|
||||
mj_solveM2(m, d, d->efc_Y, d->efc_J, sqrtInvD, nefc);
|
||||
|
||||
// arena-allocate efc_AR
|
||||
d->nA = nefc * nefc;
|
||||
d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum));
|
||||
if (!d->efc_AR) {
|
||||
mj_warning(d, mjWARN_CNSTRFULL, d->narena);
|
||||
@@ -2937,18 +2959,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
|
||||
return;
|
||||
}
|
||||
|
||||
// space for B = backsubM2(J')' and its transpose
|
||||
mjtNum* B = mjSTACKALLOC(d, nefc*nv, mjtNum);
|
||||
mjtNum* BT = mjSTACKALLOC(d, nv*nefc, mjtNum);
|
||||
// construct YT on stack
|
||||
mjtNum* YT = mjSTACKALLOC(d, nv*nefc, mjtNum);
|
||||
mju_transpose(YT, d->efc_Y, nefc, nv);
|
||||
|
||||
// B = backsubM2(J')'
|
||||
mj_solveM2(m, d, B, d->efc_J, sqrtInvD, nefc);
|
||||
|
||||
// construct BT
|
||||
mju_transpose(BT, B, nefc, nv);
|
||||
|
||||
// AR = B * B'
|
||||
mju_sqrMatTD(d->efc_AR, BT, NULL, nv, nefc);
|
||||
// AR = Y * Y'
|
||||
mju_sqrMatTD(d->efc_AR, YT, NULL, nv, nefc);
|
||||
|
||||
// add R to diagonal of AR
|
||||
for (int r=0; r < nefc; r++) {
|
||||
|
||||
@@ -1320,6 +1320,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
|
||||
d->nl = 0;
|
||||
d->nefc = 0;
|
||||
d->nJ = 0;
|
||||
d->nY = 0;
|
||||
d->nA = 0;
|
||||
d->nisland = 0;
|
||||
d->nidof = 0;
|
||||
|
||||
@@ -1192,7 +1192,7 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
|
||||
|
||||
// BVHs
|
||||
fprintf(fp, "BVH:\n");
|
||||
fprintf(fp, " %-8s%-8s%-8s%-10s%-s\n","id", "depth", "nodeid", "child[0]" ,"child[1]");
|
||||
fprintf(fp, " %-8s%-8s%-8s%-10s%-s\n", "id", "depth", "nodeid", "child[0]", "child[1]");
|
||||
for (int i=0; i < m->nbvh; i++) {
|
||||
fprintf(fp, " %-8d%-8d% -8d% -10d% -d\n",
|
||||
i, m->bvh_depth[i], m->bvh_nodeid[i], m->bvh_child[2*i], m->bvh_child[2*i+1]);
|
||||
@@ -1204,7 +1204,6 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// print mjModel to text file
|
||||
void mj_printModel(const mjModel* m, const char* filename) {
|
||||
mj_printFormattedModel(m, filename, FLOAT_FORMAT);
|
||||
@@ -1545,7 +1544,6 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
|
||||
mjtNum force[6] = {0};
|
||||
mj_contactForce(m, d, i, force);
|
||||
printVector(" force ", force, 6, fp, float_format);
|
||||
|
||||
}
|
||||
if (d->ncon) fprintf(fp, "\n");
|
||||
|
||||
@@ -1568,6 +1566,11 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
|
||||
d->efc_J_rowadr, d->efc_J_colind, fp, float_format);
|
||||
mj_printSparsity("J: constraint Jacobian", d->nefc, m->nv, d->efc_J_rowadr, NULL,
|
||||
d->efc_J_rownnz, d->efc_J_rowsuper, d->efc_J_colind, fp);
|
||||
if (d->nY) {
|
||||
mj_printSparsity("EFC_Y: inverse constraint inertia square root", d->nefc, m->nv,
|
||||
d->efc_Y_rowadr, NULL, d->efc_Y_rownnz, d->efc_J_rowsuper,
|
||||
d->efc_Y_colind, fp);
|
||||
}
|
||||
if (d->nisland) {
|
||||
mj_printBlockSparsity("IEFC_J: block-diagonalized constraint Jacobian (nnzs are island ids)",
|
||||
d->nefc, d->nidof, d->nisland,
|
||||
@@ -1582,7 +1585,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
|
||||
printArray2dInt("EFC_AR_ROWADR", d->nefc, 1, d->efc_AR_rowadr, fp);
|
||||
printSparse("EFC_AR", d->efc_AR, d->nefc, d->efc_AR_rownnz,
|
||||
d->efc_AR_rowadr, d->efc_AR_colind, fp, float_format);
|
||||
mj_printSparsity("efc_AR: inverse constraint inertia", d->nefc, d->nefc, d->efc_AR_rowadr,
|
||||
mj_printSparsity("EFC_AR: inverse constraint inertia", d->nefc, d->nefc, d->efc_AR_rowadr,
|
||||
NULL, d->efc_AR_rownnz, NULL, d->efc_AR_colind, fp);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5724,6 +5724,7 @@ public unsafe struct mjData_ {
|
||||
public int nl;
|
||||
public int nefc;
|
||||
public int nJ;
|
||||
public int nY;
|
||||
public int nA;
|
||||
public int nisland;
|
||||
public int nidof;
|
||||
@@ -5883,6 +5884,10 @@ public unsafe struct mjData_ {
|
||||
public double* iefc_frictionloss;
|
||||
public double* iefc_D;
|
||||
public double* iefc_R;
|
||||
public int* efc_Y_rownnz;
|
||||
public int* efc_Y_rowadr;
|
||||
public int* efc_Y_colind;
|
||||
public double* efc_Y;
|
||||
public int* efc_AR_rownnz;
|
||||
public int* efc_AR_rowadr;
|
||||
public int* efc_AR_colind;
|
||||
|
||||
@@ -6494,6 +6494,12 @@ struct MjData {
|
||||
void set_nJ(int value) {
|
||||
ptr_->nJ = value;
|
||||
}
|
||||
int nY() const {
|
||||
return ptr_->nY;
|
||||
}
|
||||
void set_nY(int value) {
|
||||
ptr_->nY = value;
|
||||
}
|
||||
int nA() const {
|
||||
return ptr_->nA;
|
||||
}
|
||||
@@ -7005,6 +7011,18 @@ struct MjData {
|
||||
emscripten::val iefc_R() const {
|
||||
return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_R));
|
||||
}
|
||||
emscripten::val efc_Y_rownnz() const {
|
||||
return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_Y_rownnz));
|
||||
}
|
||||
emscripten::val efc_Y_rowadr() const {
|
||||
return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_Y_rowadr));
|
||||
}
|
||||
emscripten::val efc_Y_colind() const {
|
||||
return emscripten::val(emscripten::typed_memory_view(ptr_->nY, ptr_->efc_Y_colind));
|
||||
}
|
||||
emscripten::val efc_Y() const {
|
||||
return emscripten::val(emscripten::typed_memory_view(ptr_->nY, ptr_->efc_Y));
|
||||
}
|
||||
emscripten::val efc_AR_rownnz() const {
|
||||
return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_AR_rownnz));
|
||||
}
|
||||
@@ -11547,6 +11565,10 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
|
||||
.property("efc_J_rowsuper", &MjData::efc_J_rowsuper)
|
||||
.property("efc_KBIP", &MjData::efc_KBIP)
|
||||
.property("efc_R", &MjData::efc_R)
|
||||
.property("efc_Y", &MjData::efc_Y)
|
||||
.property("efc_Y_colind", &MjData::efc_Y_colind)
|
||||
.property("efc_Y_rowadr", &MjData::efc_Y_rowadr)
|
||||
.property("efc_Y_rownnz", &MjData::efc_Y_rownnz)
|
||||
.property("efc_aref", &MjData::efc_aref)
|
||||
.property("efc_b", &MjData::efc_b)
|
||||
.property("efc_diagApprox", &MjData::efc_diagApprox)
|
||||
@@ -11626,6 +11648,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
|
||||
.property("moment_rownnz", &MjData::moment_rownnz)
|
||||
.property("nA", &MjData::nA, &MjData::set_nA, reference())
|
||||
.property("nJ", &MjData::nJ, &MjData::set_nJ, reference())
|
||||
.property("nY", &MjData::nY, &MjData::set_nY, reference())
|
||||
.property("narena", &MjData::narena, &MjData::set_narena, reference())
|
||||
.property("nbody_awake", &MjData::nbody_awake, &MjData::set_nbody_awake, reference())
|
||||
.property("nbuffer", &MjData::nbuffer, &MjData::set_nbuffer, reference())
|
||||
|
||||
@@ -315,6 +315,10 @@ MJDATA_SIZES: tuple[str, ...] = (
|
||||
"efc_AR_colind",
|
||||
"efc_AR_rowadr",
|
||||
"efc_AR_rownnz",
|
||||
"efc_Y",
|
||||
"efc_Y_colind",
|
||||
"efc_Y_rowadr",
|
||||
"efc_Y_rownnz",
|
||||
"efc_D",
|
||||
"efc_J",
|
||||
"efc_JT",
|
||||
|
||||
Reference in New Issue
Block a user