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

Refactor Newton solver: move Hessian from arena back to the stack.

Browse Source 
The changes in 9a0dc20821, moving the Hessian from the stack to the arena, should be rolled back: they prevent future threading over islands. Unlike the stack, arena allocations are not thread-friendly. The reason for the original move was to save memory, but the savings are small: `O(ctx->nH)` and only linear in `nv`. The significant reduction of the Cholesky factor, from quadratic in `nv` to quadratic in the largest dof island — the reduction afforded by 2dd518734f — remains in place.

Also refactor and improve readability.

PiperOrigin-RevId: 685717356
Change-Id: Ia9ec3e44a62d459a3b9cffb578cf6479d5aa1d7f
This commit is contained in:
Yuval Tassa
2024-10-14 08:32:44 -07:00
committed by Copybara-Service
parent ac91a7639d
commit 7ec94f46d4
4 changed files with 260 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+232 -197
View File
@@ -774,7 +774,7 @@ struct _mjCGContext {
int* dofind; // dof indices of this island, NULL if monolithic
int* efcind; // constraint indices of this island, NULL if monolithic
// arrays
// common arrays (CGallocate)
mjtNum* Jaref; // Jac*qacc - aref (nefc x 1)
mjtNum* Jv; // Jac*search (nefc x 1)
mjtNum* Ma; // M*qacc (nv x 1)
@@ -784,6 +784,24 @@ struct _mjCGContext {
mjtNum* search; // linesearch vector (nv x 1)
mjtNum* quad; // quadratic polynomials for constraint costs (nefc x 3)
// Newton arrays, known-size (CGallocate)
mjtNum* D; // constraint inertia (nefc x 1)
mjtNum* C; // reduced sparse inertia matrix (nC x 1)
int* H_rowadr; // Hessian row addresses (nv x 1)
int* H_rownnz; // Hessian row nonzeros (nv x 1)
int* H_lowernnz; // Hessian lower triangle row nonzeros (nv x 1)
int* L_rownnz; // Hessian factor row nonzeros (nv x 1)
int* L_rowadr; // Hessian factor row addresses (nv x 1)
// Newton arrays, computed-size (HessianMake)
int nH; // number of nonzeros in Hessian H
int* H_colind; // Hessian column indices (nH x 1)
mjtNum* H; // Hessian (nH x 1)
int nL; // number of nonzeros in Cholesky factor L
int* L_colind; // Cholesky factor column indices (nL x 1)
mjtNum* L; // Cholesky factor (nL x 1)
mjtNum* Lcone; // Cholesky factor with cone contributions (nL x 1)
// globals
mjtNum cost; // constraint + Gauss cost
mjtNum quadGauss[3]; // quadratic polynomial for Gauss cost
@@ -800,17 +818,17 @@ struct _mjCGContext {
typedef struct _mjCGContext mjCGContext;
// allocate mjCGContext: mjMARK/FREE in caller function!
// allocate fixed-size arrays in mjCGContext
// mj_{mark/free}Stack in calling function!
static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
int island, int flg_Newton) {
// clear everything
memset(ctx, 0, sizeof(mjCGContext));
// get sizes
int nv = island < 0 ? m->nv : d->island_dofnum[island];
int nefc = island < 0 ? d->nefc : d->island_efcnum[island];
// clear everything
memset(ctx, 0, sizeof(mjCGContext));
// island-related
ctx->island = island;
ctx->nv = nv;
@@ -828,30 +846,19 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
ctx->search = mj_stackAllocNum(d, nv);
ctx->quad = mj_stackAllocNum(d, nefc*3);
// Hessian (Newton only)
// Newton only, known-size arrays
ctx->flg_Newton = flg_Newton;
if (flg_Newton) {
// sparse: allocate L_rowadr, L_rownnz
ctx->D = mj_stackAllocNum(d, nefc);
// sparse Newton only
if (mj_isSparse(m)) {
d->L_rowadr = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
if (!d->L_rowadr) mjERROR("failed to allocate L_rowadr");
d->L_rownnz = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
if (!d->L_rownnz) mjERROR("failed to allocate L_rownnz");
// zero nnzL, clear pointers (compute and allocate later in HessianDirect)
d->nnzL = 0;
d->L_colind = NULL;
d->L = NULL;
d->Lcone = NULL;
}
// dense: allocate L
else if (d->nnzL != nv*nv) {
d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nv*nv, _Alignof(mjtNum));
if (!d->L) mjERROR("failed to allocate L");
// set dense nnzL
d->nnzL = nv*nv;
ctx->C = mj_stackAllocNum(d, m->nC);
ctx->H_rowadr = mj_stackAllocInt(d, nv);
ctx->H_rownnz = mj_stackAllocInt(d, nv);
ctx->H_lowernnz = mj_stackAllocInt(d, nv);
ctx->L_rownnz = mj_stackAllocInt(d, nv);
ctx->L_rowadr = mj_stackAllocInt(d, nv);
}
}
}
@@ -904,10 +911,10 @@ static void CGupdateGradient(const mjModel* m, const mjData* d, mjCGContext* ctx
// TODO: b/295296178 - add island support to Newton solver
if (ctx->flg_Newton) {
if (mj_isSparse(m)) {
mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? d->Lcone : d->L),
ctx->grad, nv, d->L_rownnz, d->L_rowadr, d->L_colind);
mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L),
ctx->grad, nv, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind);
} else {
mju_cholSolve(ctx->Mgrad, (ctx->ncone ? d->Lcone : d->L), ctx->grad, nv);
mju_cholSolve(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L), ctx->grad, nv);
}
}
@@ -1375,20 +1382,192 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
// allocate and compute Hessian given efc_state
// mj_{mark/free}Stack in caller function!
static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
// compute constraint inertia
for (int i=0; i < nefc; i++) {
ctx->D[i] = d->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? d->efc_D[i] : 0;
}
// sparse
if (mj_isSparse(m)) {
// copy values of reduced sparse inertia matrix C
for (int i=0; i < m->nC; i++) {
ctx->C[i] = d->qM[d->mapM2C[i]];
}
// initialize Hessian rowadr, rownnz
mju_sqrMatTDSparseInit(ctx->H_rownnz, ctx->H_rowadr, nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
// add nC to Hessian total nonzeros (unavoidable overcounting since H_colind is still unknown)
ctx->nH = m->nC + ctx->H_rowadr[nv - 1] + ctx->H_rownnz[nv - 1];
// shift H row adresses to make room for C
int shift = 0;
for (int r = 0; r < nv - 1; r++) {
shift += d->C_rownnz[r];
ctx->H_rowadr[r + 1] += shift;
}
// allocate H_colind and H
ctx->H_colind = mj_stackAllocInt(d, ctx->nH);
ctx->H = mj_stackAllocNum(d, ctx->nH);
// compute H = J'*D*J
mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, ctx->D, nefc, nv,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
// add mass matrix: H = J'*D*J + C
mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
ctx->C, d->C_rownnz, d->C_rowadr, d->C_colind);
// count total and row non-zeros of reverse-Cholesky factor L
ctx->nL = mju_cholFactorNNZ(ctx->L_rownnz, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, d);
// compute L row adresses: rowadr = cumsum(rownnz)
ctx->L_rowadr[0] = 0;
for (int r=1; r < nv; r++) {
ctx->L_rowadr[r] = ctx->L_rowadr[r-1] + ctx->L_rownnz[r-1];
}
// allocate L_colind, L, Lcone
ctx->L_colind = mj_stackAllocInt(d, ctx->nL);
ctx->L = mj_stackAllocNum(d, ctx->nL);
if (m->opt.cone == mjCONE_ELLIPTIC) {
ctx->Lcone = mj_stackAllocNum(d, ctx->nL);
}
// count nonzeros in rows of H lower triangle
for (int r = 0; r < nv; r++) {
const int* colind = ctx->H_colind + ctx->H_rowadr[r];
int rownnz = ctx->H_rownnz[r];
// count nonzeros up to diagonal (inclusive) for row r
int nnz = 1;
while (nnz < rownnz && colind[nnz - 1] < r) {
nnz++;
}
// last row element is not the diagonal; SHOULD NOT OCCUR
if (colind[nnz - 1] != r) {
mjERROR("Newton solver Hessian has zero diagonal on row %d", r);
}
// save row nonzeros
ctx->H_lowernnz[r] = nnz;
}
}
// dense
else {
// allocate L, Lcone
ctx->nL = nv*nv;
ctx->L = mj_stackAllocNum(d, ctx->nL);
if (m->opt.cone == mjCONE_ELLIPTIC) {
ctx->Lcone = mj_stackAllocNum(d, ctx->nL);
}
// compute H = M + J'*D*J
mju_sqrMatTD(ctx->L, d->efc_J, ctx->D, nefc, nv);
mj_addMDense(m, d, ctx->L);
}
}
// forward declaration of HessianCone (readability)
static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx);
// factorize Hessian: L = chol(H), maybe (re)compute H given efc_state
static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx,
int flg_recompute) {
int nv = m->nv, nefc = d->nefc;
// maybe compute constraint inertia
if (flg_recompute) {
for (int i=0; i < nefc; i++) {
ctx->D[i] = d->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? d->efc_D[i] : 0;
}
}
// sparse
if (mj_isSparse(m)) {
// maybe compute H = M + J'*D*J
if (flg_recompute) {
// compute H = J'*D*J
mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, ctx->D, nefc, nv,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
// add mass matrix: H = J'*D*J + C
mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
ctx->C, d->C_rownnz, d->C_rowadr, d->C_colind);
}
// copy H lower-triangle into L, fill-in already accounted for
for (int r = 0; r < nv; r++) {
int nnz = ctx->H_lowernnz[r];
mju_copy(ctx->L + ctx->L_rowadr[r], ctx->H + ctx->H_rowadr[r], nnz);
mju_copyInt(ctx->L_colind + ctx->L_rowadr[r], ctx->H_colind + ctx->H_rowadr[r], nnz);
ctx->L_rownnz[r] = nnz;
}
// in-place sparse factorization: L = chol(H)
int rank = mju_cholFactorSparse(ctx->L, nv, mjMINVAL,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, d);
// rank-deficient; SHOULD NOT OCCUR
if (rank != nv) {
mjERROR("rank-deficient sparse Hessian");
}
// pre-counted nL does not match post-factorization nL; SHOULD NOT OCCUR
if (ctx->nL != ctx->L_rowadr[nv-1] + ctx->L_rownnz[nv-1]) {
mjERROR("mismatch between pre-counted and post-factorization L nonzeros");
}
}
// dense
else {
// maybe compute H = M + J'*D*J
if (flg_recompute) {
mju_sqrMatTD(ctx->L, d->efc_J, ctx->D, nefc, nv);
mj_addMDense(m, d, ctx->L);
}
// factorize H
mju_cholFactor(ctx->L, nv, mjMINVAL);
}
// add cones to factor if present
if (ctx->ncone) {
HessianCone(m, d, ctx);
}
// mark full update
ctx->nupdate = nefc;
}
// elliptic case: Hcone = H + cone_contributions
// TODO: b/295296178 - add island support to Newton solver
static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
mjtNum local[36];
// allocate Lcone if required
if (!d->Lcone) {
d->Lcone = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nnzL, _Alignof(mjtNum));
}
if (!d->Lcone) mjERROR("failed to allocate Lcone");
// start with Hcone = H
mju_copy(d->Lcone, d->L, d->nnzL);
mju_copy(ctx->Lcone, ctx->L, ctx->nL);
mj_markStack(d);
@@ -1427,8 +1606,8 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
mju_copyInt(LTJ_ind, d->efc_J_colind+d->efc_J_rowadr[i+r], nnz);
// update
mju_cholUpdateSparse(d->Lcone, LTJ_row, nv, 1,
d->L_rownnz, d->L_rowadr, d->L_colind, nnz, LTJ_ind, d);
mju_cholUpdateSparse(ctx->Lcone, LTJ_row, nv, 1,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, LTJ_ind, d);
}
}
@@ -1444,7 +1623,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
// update
for (int r=0; r < dim; r++) {
mju_cholUpdate(d->Lcone, LTJ+r*nv, nv, 1);
mju_cholUpdate(ctx->Lcone, LTJ+r*nv, nv, 1);
}
}
@@ -1461,155 +1640,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
// compute and factorize Hessian: direct method
// TODO: b/295296178 - add island support to Newton solver
static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
mj_markStack(d);
// compute D corresponding to quad states
mjtNum* D = mj_stackAllocNum(d, nefc);
for (int i=0; i < nefc; i++) {
if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) {
D[i] = d->efc_D[i];
} else {
D[i] = 0;
}
}
// sparse
if (mj_isSparse(m)) {
// copy values of reduced sparse inertia matrix C, get nnz
int nnz_C = m->nC;
mjtNum* C = mj_stackAllocNum(d, nnz_C);
for (int i=0; i < nnz_C; i++) {
C[i] = d->qM[d->mapM2C[i]];
}
// allocate and initialize Hessian rowadr, rownnz; get nnz for J'*J
int* H_rowadr = mj_stackAllocInt(d, nv);
int* H_rownnz = mj_stackAllocInt(d, nv);
mju_sqrMatTDSparseInit(H_rownnz, H_rowadr, nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
int nnz_JTJ = H_rowadr[nv-1] + H_rownnz[nv-1];
// shift H rowadr to make room for C
int shift = 0;
for (int r = 0; r < nv - 1; r++) {
shift += d->C_rownnz[r];
H_rowadr[r + 1] += shift;
}
// allocate Hessian H, colind
int nnz_H = nnz_C + nnz_JTJ;
mjtNum* H = mj_stackAllocNum(d, nnz_H);
int* H_colind = mj_stackAllocInt(d, nnz_H);
// compute H = J'*D*J
mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, nefc, nv,
H_rownnz, H_rowadr, H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
// add mass matrix; H = J'*D*J + C
mj_addMSparse(m, d, H, H_rownnz, H_rowadr, H_colind,
C, d->C_rownnz, d->C_rowadr, d->C_colind);
// count row and total non-zeros of reverse-Cholesky factor L
int* parent = mj_stackAllocInt(d, nv);
int* flag = mj_stackAllocInt(d, nv);
int nnz_L = mju_cholFactorNNZ(d->L_rownnz, parent, flag, H_rownnz, H_rowadr, H_colind, nv);
// allocate L_colind, L on arena if required
if (!d->nnzL) {
// nnzL is 0 but pointers are allocated; SHOULD NOT OCCUR
if (d->L_colind || d->L) {
mjERROR("nnzL is 0 but L_colind or L or Lcone are allocated");
}
// allocate on arena
d->L_colind = mj_arenaAllocByte(d, sizeof(int) * nnz_L, _Alignof(int));
if (!d->L_colind) mjERROR("failed to allocate L_colind");
d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nnz_L, _Alignof(mjtNum));
if (!d->L) mjERROR("failed to allocate L");
// set nnzL
d->nnzL = nnz_L;
} else if (d->nnzL != nnz_L) {
// nnzL is nonzero but not equal to computed value; SHOULD NOT OCCUR
mjERROR("nnzL is nonzero but not equal to computed value");
}
// compute L row adresses: L_rowadr = cumsum(L_rownnz)
d->L_rowadr[0] = 0;
for (int r=1; r < nv; r++) {
d->L_rowadr[r] = d->L_rowadr[r-1] + d->L_rownnz[r-1];
}
// copy H lower-triangle into L
for (int r = 0; r < nv; r++) {
// count H non-zeros up to diagonal (inclusive) for row r
const int* colind = H_colind + H_rowadr[r];
int rownnz = 1;
while (rownnz < nv && colind[rownnz - 1] < r) {
rownnz++;
}
// last row element is not the diagonal; SHOULD NOT OCCUR
if (colind[rownnz - 1] != r) {
mjERROR("Newton solver Hessian has zero diagonal on row %d", r);
}
// copy values and column indices
mju_copy(d->L + d->L_rowadr[r], H + H_rowadr[r], rownnz);
mju_copyInt(d->L_colind + d->L_rowadr[r], H_colind + H_rowadr[r], rownnz);
// set L_rownnz
d->L_rownnz[r] = rownnz;
}
// in-place sparse factorization L = chol(H)
int rank = mju_cholFactorSparse(d->L, nv, mjMINVAL, d->L_rownnz, d->L_rowadr, d->L_colind, d);
// rank-deficient; SHOULD NOT OCCUR
if (rank != nv) {
mjERROR("rank-deficient Hessian");
}
// pre-counted nnzL does not match post-factorization nnzL; SHOULD NOT OCCUR
if (d->nnzL != d->L_rowadr[nv-1] + d->L_rownnz[nv-1]) {
mjERROR("mismatch between pre-counted and post-factorization L nonzeros");
}
}
// dense
else {
// compute H = M + J'*D*J
mju_sqrMatTD(d->L, d->efc_J, D, nefc, nv);
mj_addMDense(m, d, d->L);
// factorize H
mju_cholFactor(d->L, nv, mjMINVAL);
}
mj_freeStack(d);
// add cones if present
if (ctx->ncone) {
HessianCone(m, d, ctx);
}
// mark full update
ctx->nupdate = nefc;
}
// incremental update to Hessian
// TODO: b/295296178 - add island support to Newton solver
// incremental update to Hessian factor due to changes in efc_state
static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, const int* oldstate) {
int rank, nv = m->nv, nefc = d->nefc;
mj_markStack(d);
@@ -1646,20 +1677,20 @@ static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, co
mju_scl(vec, d->efc_J+adr, mju_sqrt(d->efc_D[i]), nnz);
mju_copyInt(vec_ind, d->efc_J_colind+adr, nnz);
// sparse update
rank = mju_cholUpdateSparse(d->L, vec, nv, flag_update,
d->L_rownnz, d->L_rowadr, d->L_colind, nnz, vec_ind,
// sparse update or downdate
rank = mju_cholUpdateSparse(ctx->L, vec, nv, flag_update,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, vec_ind,
d);
} else {
mju_scl(vec, d->efc_J+i*nv, mju_sqrt(d->efc_D[i]), nv);
rank = mju_cholUpdate(d->L, vec, nv, flag_update);
rank = mju_cholUpdate(ctx->L, vec, nv, flag_update);
}
ctx->nupdate++;
// recompute H directly if accuracy lost
if (rank < nv) {
mj_freeStack(d);
HessianDirect(m, d, ctx);
FactorizeHessian(m, d, ctx, /*flg_recompute=*/1);
// nothing else to do
return;
@@ -1718,7 +1749,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// first update
CGupdateConstraint(m, d, &ctx);
if (flg_Newton) {
HessianDirect(m, d, &ctx);
// compute and factorize Hessian
MakeHessian(m, d, &ctx);
FactorizeHessian(m, d, &ctx, /*flg_recompute=*/0);
}
CGupdateGradient(m, d, &ctx);
@@ -1833,7 +1866,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// set solver_nnz
if (flg_Newton) {
if (mj_isSparse(m)) {
d->solver_nnz[island_stat] = 2*d->nnzL - nv;
// two L factors if Lcone is present
int num_factors = 1 + (ctx.Lcone != NULL);
d->solver_nnz[island_stat] = num_factors * ctx.nL + ctx.nH;
} else {
d->solver_nnz[island_stat] = nv*nv;
}
src/engine/engine_util_sparse.c
+11 -6
View File
@@ -738,7 +738,7 @@ void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc) {
// res_rowadr is required to be precomputed
void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const mjtNum* diag, int nr, int nc,
int* res_rownnz, int* res_rowadr, int* res_colind,
int* res_rownnz, const int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
@@ -859,13 +859,17 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
// compute row non-zeros of reverse-Cholesky factor L, return total
// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package'
int mju_cholFactorNNZ(int* L_rownnz, int* parent, int* flag, const int* rownnz,
const int* rowadr, const int* colind, int n) {
int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind,
int n, mjData* d) {
mj_markStack(d);
int* parent = mj_stackAllocInt(d, n);
int* flag = mj_stackAllocInt(d, n);
// loop over rows in reverse order
for (int r = n - 1; r >= 0; r--) {
parent[r] = -1;
flag[r] = r;
L_rownnz[r] = 0;
L_rownnz[r] = 1; // start with 1 for diagonal
int start = rowadr[r];
int end = start + rownnz[r];
// loop over non-zero columns
@@ -886,10 +890,11 @@ int mju_cholFactorNNZ(int* L_rownnz, int* parent, int* flag, const int* rownnz,
}
}
// add 1 for diagonal, accumulate sum
mj_freeStack(d);
// accumulate sum
int sum = 0;
for (int r = 0; r < n; r++) {
L_rownnz[r]++;
sum += L_rownnz[r];
}
src/engine/engine_util_sparse.h
+3 -3
View File
@@ -89,7 +89,7 @@ MJAPI void mju_superSparse(int nr, int* rowsuper,
// res_rowadr is required to be precomputed
MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const mjtNum* diag, int nr, int nc,
int* res_rownnz, int* res_rowadr, int* res_colind,
int* res_rownnz, const int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
@@ -106,8 +106,8 @@ MJAPI void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc);
// compute row non-zeros of reverse-Cholesky factor L, return total
MJAPI int mju_cholFactorNNZ(int* L_rownnz, int* parent, int* flag, const int* rownnz,
const int* rowadr, const int* colind, int n);
MJAPI int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind,
int n, mjData* d);
#ifdef __cplusplus
}
test/engine/engine_util_sparse_test.cc
+14 -16
View File
@@ -972,6 +972,9 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
}
TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
mjModel* model = LoadModelFromString(modelStr);
mjData* d = mj_makeData(model);
// A = [[1, 0],
// [0, 1]]
int nA = 2;
@@ -981,11 +984,9 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrA[2];
int colindA[4];
int rownnzA_factor[2];
int parentA[2];
int workspaceA[2];
mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA);
int nnzA = mju_cholFactorNNZ(rownnzA_factor, parentA, workspaceA, rownnzA,
rowadrA, colindA, nA);
int nnzA = mju_cholFactorNNZ(rownnzA_factor,
rownnzA, rowadrA, colindA, nA, d);
EXPECT_EQ(nnzA, 2);
EXPECT_THAT(AsVector(rownnzA_factor, 2), ElementsAre(1, 1));
@@ -1000,11 +1001,9 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrB[3];
int colindB[9];
int rownnzB_factor[3];
int parentB[3];
int workspaceB[3];
mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB);
int nnzB = mju_cholFactorNNZ(rownnzB_factor, parentB, workspaceB, rownnzB,
rowadrB, colindB, nB);
int nnzB = mju_cholFactorNNZ(rownnzB_factor,
rownnzB, rowadrB, colindB, nB, d);
EXPECT_EQ(nnzB, 5);
EXPECT_THAT(AsVector(rownnzB_factor, 3), ElementsAre(1, 2, 2));
@@ -1019,11 +1018,9 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrC[3];
int colindC[9];
int rownnzC_factor[3];
int parentC[3];
int workspaceC[3];
mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC);
int nnzC = mju_cholFactorNNZ(rownnzC_factor, parentC, workspaceC, rownnzC,
rowadrC, colindC, nC);
int nnzC = mju_cholFactorNNZ(rownnzC_factor,
rownnzC, rowadrC, colindC, nC, d);
EXPECT_EQ(nnzC, 4);
EXPECT_THAT(AsVector(rownnzC_factor, 3), ElementsAre(1, 2, 1));
@@ -1039,14 +1036,15 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrD[4];
int colindD[16];
int rownnzD_factor[4];
int parentD[4];
int workspaceD[4];
mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD);
int nnzD = mju_cholFactorNNZ(rownnzD_factor, parentD, workspaceD, rownnzD,
rowadrD, colindD, nD);
int nnzD = mju_cholFactorNNZ(rownnzD_factor,
rownnzD, rowadrD, colindD, nD, d);
EXPECT_EQ(nnzD, 8);
EXPECT_THAT(AsVector(rownnzD_factor, 4), ElementsAre(1, 2, 2, 3));
mj_deleteData(d);
mj_deleteModel(model);
}
} // namespace