Actual source code: brgn.c

      1: #include <../src/tao/leastsquares/impls/brgn/brgn.h>

      3: const char *const TaoBRGNRegularizationTypes[] = {"user", "l2prox", "l2pure", "l1dict", "lm", "TaoBRGNRegularizationType", "TAOBRGN_REGULARIZATION_", NULL};

      5: static PetscErrorCode GNHessianProd(Mat H, Vec in, Vec out)
      6: {
      7:   TAO_BRGN *gn;

      9:   PetscFunctionBegin;
     10:   PetscCall(MatShellGetContext(H, &gn));
     11:   PetscCall(MatMult(gn->subsolver->ls_jac, in, gn->r_work));
     12:   PetscCall(MatMultTranspose(gn->subsolver->ls_jac, gn->r_work, out));
     13:   switch (gn->reg_type) {
     14:   case TAOBRGN_REGULARIZATION_USER:
     15:     PetscCall(MatMult(gn->Hreg, in, gn->x_work));
     16:     PetscCall(VecAXPY(out, gn->lambda, gn->x_work));
     17:     break;
     18:   case TAOBRGN_REGULARIZATION_L2PURE:
     19:     PetscCall(VecAXPY(out, gn->lambda, in));
     20:     break;
     21:   case TAOBRGN_REGULARIZATION_L2PROX:
     22:     PetscCall(VecAXPY(out, gn->lambda, in));
     23:     break;
     24:   case TAOBRGN_REGULARIZATION_L1DICT:
     25:     /* out = out + lambda*D'*(diag.*(D*in)) */
     26:     if (gn->D) {
     27:       PetscCall(MatMult(gn->D, in, gn->y)); /* y = D*in */
     28:     } else {
     29:       PetscCall(VecCopy(in, gn->y));
     30:     }
     31:     PetscCall(VecPointwiseMult(gn->y_work, gn->diag, gn->y)); /* y_work = diag.*(D*in), where diag = epsilon^2 ./ sqrt(x.^2+epsilon^2).^3 */
     32:     if (gn->D) {
     33:       PetscCall(MatMultTranspose(gn->D, gn->y_work, gn->x_work)); /* x_work = D'*(diag.*(D*in)) */
     34:     } else {
     35:       PetscCall(VecCopy(gn->y_work, gn->x_work));
     36:     }
     37:     PetscCall(VecAXPY(out, gn->lambda, gn->x_work));
     38:     break;
     39:   case TAOBRGN_REGULARIZATION_LM:
     40:     PetscCall(VecPointwiseMult(gn->x_work, gn->damping, in));
     41:     PetscCall(VecAXPY(out, 1, gn->x_work));
     42:     break;
     43:   }
     44:   PetscFunctionReturn(PETSC_SUCCESS);
     45: }
     46: static PetscErrorCode ComputeDamping(TAO_BRGN *gn)
     47: {
     48:   const PetscScalar *diag_ary;
     49:   PetscScalar       *damping_ary;
     50:   PetscInt           i, n;

     52:   PetscFunctionBegin;
     53:   /* update damping */
     54:   PetscCall(VecGetArray(gn->damping, &damping_ary));
     55:   PetscCall(VecGetArrayRead(gn->diag, &diag_ary));
     56:   PetscCall(VecGetLocalSize(gn->damping, &n));
     57:   for (i = 0; i < n; i++) damping_ary[i] = PetscClipInterval(diag_ary[i], PETSC_SQRT_MACHINE_EPSILON, PetscSqrtReal(PETSC_MAX_REAL));
     58:   PetscCall(VecScale(gn->damping, gn->lambda));
     59:   PetscCall(VecRestoreArray(gn->damping, &damping_ary));
     60:   PetscCall(VecRestoreArrayRead(gn->diag, &diag_ary));
     61:   PetscFunctionReturn(PETSC_SUCCESS);
     62: }
     63: /*@
     64:   TaoBRGNGetDampingVector - Get the damping vector $\mathrm{diag}(J^T J)$ from a `TAOBRGN` with `TAOBRGN_REGULARIZATION_LM` regularization

     66:   Collective

     68:   Input Parameter:
     69: . tao - a `Tao` of type `TAOBRGN` with `TAOBRGN_REGULARIZATION_LM` regularization

     71:   Output Parameter:
     72: . d - the damping vector

     74:   Level: developer

     76: .seealso: [](ch_tao), `Tao`, `TAOBRGN`, `TaoBRGNRegularzationTypes`
     77: @*/
     78: PetscErrorCode TaoBRGNGetDampingVector(Tao tao, Vec *d)
     79: {
     80:   PetscFunctionBegin;
     82:   PetscAssertPointer(d, 2);
     83:   PetscUseMethod((PetscObject)tao, "TaoBRGNGetDampingVector_C", (Tao, Vec *), (tao, d));
     84:   PetscFunctionReturn(PETSC_SUCCESS);
     85: }

     87: static PetscErrorCode TaoBRGNGetDampingVector_BRGN(Tao tao, Vec *d)
     88: {
     89:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

     91:   PetscFunctionBegin;
     92:   PetscCheck(gn->reg_type == TAOBRGN_REGULARIZATION_LM, PetscObjectComm((PetscObject)tao), PETSC_ERR_SUP, "Damping vector is only available if regularization type is lm.");
     93:   *d = gn->damping;
     94:   PetscFunctionReturn(PETSC_SUCCESS);
     95: }

     97: static PetscErrorCode GNObjectiveGradientEval(Tao tao, Vec X, PetscReal *fcn, Vec G, void *ptr)
     98: {
     99:   TAO_BRGN   *gn = (TAO_BRGN *)ptr;
    100:   PetscInt    K; /* dimension of D*X */
    101:   PetscScalar yESum;
    102:   PetscReal   f_reg;

    104:   PetscFunctionBegin;
    105:   /* compute objective *fcn*/
    106:   /* compute first term 0.5*||ls_res||_2^2 */
    107:   PetscCall(TaoComputeResidual(tao, X, tao->ls_res));
    108:   PetscCall(VecDot(tao->ls_res, tao->ls_res, fcn));
    109:   *fcn *= 0.5;
    110:   /* compute gradient G */
    111:   PetscCall(TaoComputeResidualJacobian(tao, X, tao->ls_jac, tao->ls_jac_pre));
    112:   PetscCall(MatMultTranspose(tao->ls_jac, tao->ls_res, G));
    113:   /* add the regularization contribution */
    114:   switch (gn->reg_type) {
    115:   case TAOBRGN_REGULARIZATION_USER:
    116:     PetscCall((*gn->regularizerobjandgrad)(tao, X, &f_reg, gn->x_work, gn->reg_obj_ctx));
    117:     *fcn += gn->lambda * f_reg;
    118:     PetscCall(VecAXPY(G, gn->lambda, gn->x_work));
    119:     break;
    120:   case TAOBRGN_REGULARIZATION_L2PURE:
    121:     /* compute f = f + lambda*0.5*xk'*xk */
    122:     PetscCall(VecDot(X, X, &f_reg));
    123:     *fcn += gn->lambda * 0.5 * f_reg;
    124:     /* compute G = G + lambda*xk */
    125:     PetscCall(VecAXPY(G, gn->lambda, X));
    126:     break;
    127:   case TAOBRGN_REGULARIZATION_L2PROX:
    128:     /* compute f = f + lambda*0.5*(xk - xkm1)'*(xk - xkm1) */
    129:     PetscCall(VecAXPBYPCZ(gn->x_work, 1.0, -1.0, 0.0, X, gn->x_old));
    130:     PetscCall(VecDot(gn->x_work, gn->x_work, &f_reg));
    131:     *fcn += gn->lambda * 0.5 * f_reg;
    132:     /* compute G = G + lambda*(xk - xkm1) */
    133:     PetscCall(VecAXPBYPCZ(G, gn->lambda, -gn->lambda, 1.0, X, gn->x_old));
    134:     break;
    135:   case TAOBRGN_REGULARIZATION_L1DICT:
    136:     /* compute f = f + lambda*sum(sqrt(y.^2+epsilon^2) - epsilon), where y = D*x*/
    137:     if (gn->D) {
    138:       PetscCall(MatMult(gn->D, X, gn->y)); /* y = D*x */
    139:     } else {
    140:       PetscCall(VecCopy(X, gn->y));
    141:     }
    142:     PetscCall(VecPointwiseMult(gn->y_work, gn->y, gn->y));
    143:     PetscCall(VecShift(gn->y_work, gn->epsilon * gn->epsilon));
    144:     PetscCall(VecSqrtAbs(gn->y_work)); /* gn->y_work = sqrt(y.^2+epsilon^2) */
    145:     PetscCall(VecSum(gn->y_work, &yESum));
    146:     PetscCall(VecGetSize(gn->y, &K));
    147:     *fcn += gn->lambda * (yESum - K * gn->epsilon);
    148:     /* compute G = G + lambda*D'*(y./sqrt(y.^2+epsilon^2)),where y = D*x */
    149:     PetscCall(VecPointwiseDivide(gn->y_work, gn->y, gn->y_work)); /* reuse y_work = y./sqrt(y.^2+epsilon^2) */
    150:     if (gn->D) {
    151:       PetscCall(MatMultTranspose(gn->D, gn->y_work, gn->x_work));
    152:     } else {
    153:       PetscCall(VecCopy(gn->y_work, gn->x_work));
    154:     }
    155:     PetscCall(VecAXPY(G, gn->lambda, gn->x_work));
    156:     break;
    157:   case TAOBRGN_REGULARIZATION_LM:
    158:     break;
    159:   default:
    160:     break;
    161:   }
    162:   PetscFunctionReturn(PETSC_SUCCESS);
    163: }

    165: static PetscErrorCode GNComputeHessian(Tao tao, Vec X, Mat H, Mat Hpre, void *ptr)
    166: {
    167:   TAO_BRGN    *gn = (TAO_BRGN *)ptr;
    168:   PetscInt     i, n, cstart, cend;
    169:   PetscScalar *cnorms, *diag_ary;

    171:   PetscFunctionBegin;
    172:   PetscCall(TaoComputeResidualJacobian(tao, X, tao->ls_jac, tao->ls_jac_pre));
    173:   if (gn->mat_explicit) PetscCall(MatTransposeMatMult(tao->ls_jac, tao->ls_jac, MAT_REUSE_MATRIX, PETSC_DETERMINE, &gn->H));

    175:   switch (gn->reg_type) {
    176:   case TAOBRGN_REGULARIZATION_USER:
    177:     PetscCall((*gn->regularizerhessian)(tao, X, gn->Hreg, gn->reg_hess_ctx));
    178:     if (gn->mat_explicit) PetscCall(MatAXPY(gn->H, 1.0, gn->Hreg, DIFFERENT_NONZERO_PATTERN));
    179:     break;
    180:   case TAOBRGN_REGULARIZATION_L2PURE:
    181:     if (gn->mat_explicit) PetscCall(MatShift(gn->H, gn->lambda));
    182:     break;
    183:   case TAOBRGN_REGULARIZATION_L2PROX:
    184:     if (gn->mat_explicit) PetscCall(MatShift(gn->H, gn->lambda));
    185:     break;
    186:   case TAOBRGN_REGULARIZATION_L1DICT:
    187:     /* calculate and store diagonal matrix as a vector: diag = epsilon^2 ./ sqrt(x.^2+epsilon^2).^3* --> diag = epsilon^2 ./ sqrt(y.^2+epsilon^2).^3,where y = D*x */
    188:     if (gn->D) {
    189:       PetscCall(MatMult(gn->D, X, gn->y)); /* y = D*x */
    190:     } else {
    191:       PetscCall(VecCopy(X, gn->y));
    192:     }
    193:     PetscCall(VecPointwiseMult(gn->y_work, gn->y, gn->y));
    194:     PetscCall(VecShift(gn->y_work, gn->epsilon * gn->epsilon));
    195:     PetscCall(VecCopy(gn->y_work, gn->diag));                    /* gn->diag = y.^2+epsilon^2 */
    196:     PetscCall(VecSqrtAbs(gn->y_work));                           /* gn->y_work = sqrt(y.^2+epsilon^2) */
    197:     PetscCall(VecPointwiseMult(gn->diag, gn->y_work, gn->diag)); /* gn->diag = sqrt(y.^2+epsilon^2).^3 */
    198:     PetscCall(VecReciprocal(gn->diag));
    199:     PetscCall(VecScale(gn->diag, gn->epsilon * gn->epsilon));
    200:     if (gn->mat_explicit) PetscCall(MatDiagonalSet(gn->H, gn->diag, ADD_VALUES));
    201:     break;
    202:   case TAOBRGN_REGULARIZATION_LM:
    203:     /* compute diagonal of J^T J */
    204:     PetscCall(MatGetSize(gn->parent->ls_jac, NULL, &n));
    205:     PetscCall(PetscMalloc1(n, &cnorms));
    206:     PetscCall(MatGetColumnNorms(gn->parent->ls_jac, NORM_2, cnorms));
    207:     PetscCall(MatGetOwnershipRangeColumn(gn->parent->ls_jac, &cstart, &cend));
    208:     PetscCall(VecGetArray(gn->diag, &diag_ary));
    209:     for (i = 0; i < cend - cstart; i++) diag_ary[i] = cnorms[cstart + i] * cnorms[cstart + i];
    210:     PetscCall(VecRestoreArray(gn->diag, &diag_ary));
    211:     PetscCall(PetscFree(cnorms));
    212:     PetscCall(ComputeDamping(gn));
    213:     if (gn->mat_explicit) PetscCall(MatDiagonalSet(gn->H, gn->damping, ADD_VALUES));
    214:     break;
    215:   default:
    216:     break;
    217:   }
    218:   PetscFunctionReturn(PETSC_SUCCESS);
    219: }

    221: static PetscErrorCode GNHookFunction(Tao tao, PetscInt iter, PetscCtx ctx)
    222: {
    223:   TAO_BRGN *gn = (TAO_BRGN *)ctx;

    225:   PetscFunctionBegin;
    226:   /* Update basic tao information from the subsolver */
    227:   {
    228:     gn->parent->objective_term.term->nobj     = tao->objective_term.term->nobj;
    229:     gn->parent->objective_term.term->ngrad    = tao->objective_term.term->ngrad;
    230:     gn->parent->objective_term.term->nobjgrad = tao->objective_term.term->nobjgrad;
    231:     gn->parent->objective_term.term->nhess    = tao->objective_term.term->nhess;
    232:   }
    233:   gn->parent->nres        = tao->nres;
    234:   gn->parent->niter       = tao->niter;
    235:   gn->parent->ksp_its     = tao->ksp_its;
    236:   gn->parent->ksp_tot_its = tao->ksp_tot_its;
    237:   gn->parent->fc          = tao->fc;
    238:   PetscCall(TaoGetConvergedReason(tao, &gn->parent->reason));
    239:   /* Update the solution vectors */
    240:   if (iter == 0) {
    241:     PetscCall(VecSet(gn->x_old, 0.0));
    242:   } else {
    243:     PetscCall(VecCopy(tao->solution, gn->x_old));
    244:     PetscCall(VecCopy(tao->solution, gn->parent->solution));
    245:   }
    246:   /* Update the gradient */
    247:   PetscCall(VecCopy(tao->gradient, gn->parent->gradient));

    249:   /* Update damping parameter for LM */
    250:   if (gn->reg_type == TAOBRGN_REGULARIZATION_LM) {
    251:     if (iter > 0) {
    252:       if (gn->fc_old > tao->fc) {
    253:         gn->lambda = gn->lambda * gn->downhill_lambda_change;
    254:       } else {
    255:         /* uphill step */
    256:         gn->lambda = gn->lambda * gn->uphill_lambda_change;
    257:       }
    258:     }
    259:     gn->fc_old = tao->fc;
    260:   }

    262:   /* Call general purpose update function */
    263:   if (gn->parent->ops->update) PetscCall((*gn->parent->ops->update)(gn->parent, gn->parent->niter, gn->parent->user_update));
    264:   PetscFunctionReturn(PETSC_SUCCESS);
    265: }

    267: static PetscErrorCode TaoBRGNGetRegularizationType_BRGN(Tao tao, TaoBRGNRegularizationType *type)
    268: {
    269:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    271:   PetscFunctionBegin;
    272:   *type = gn->reg_type;
    273:   PetscFunctionReturn(PETSC_SUCCESS);
    274: }

    276: /*@
    277:   TaoBRGNGetRegularizationType - Get the `TaoBRGNRegularizationType` of a `TAOBRGN`

    279:   Not collective

    281:   Input Parameter:
    282: . tao - a `Tao` of type `TAOBRGN`

    284:   Output Parameter:
    285: . type - the `TaoBRGNRegularizationType`

    287:   Level: advanced

    289: .seealso: [](ch_tao), `Tao`, `TAOBRGN`, `TaoBRGNRegularizationType`, `TaoBRGNSetRegularizationType()`
    290: @*/
    291: PetscErrorCode TaoBRGNGetRegularizationType(Tao tao, TaoBRGNRegularizationType *type)
    292: {
    293:   PetscFunctionBegin;
    295:   PetscAssertPointer(type, 2);
    296:   PetscUseMethod((PetscObject)tao, "TaoBRGNGetRegularizationType_C", (Tao, TaoBRGNRegularizationType *), (tao, type));
    297:   PetscFunctionReturn(PETSC_SUCCESS);
    298: }

    300: static PetscErrorCode TaoBRGNSetRegularizationType_BRGN(Tao tao, TaoBRGNRegularizationType type)
    301: {
    302:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    304:   PetscFunctionBegin;
    305:   gn->reg_type = type;
    306:   PetscFunctionReturn(PETSC_SUCCESS);
    307: }

    309: /*@
    310:   TaoBRGNSetRegularizationType - Set the `TaoBRGNRegularizationType` of a `TAOBRGN`

    312:   Logically collective

    314:   Input Parameters:
    315: + tao  - a `Tao` of type `TAOBRGN`
    316: - type - the `TaoBRGNRegularizationType`

    318:   Level: advanced

    320: .seealso: [](ch_tao), `Tao`, `TAOBRGN`, `TaoBRGNRegularizationType`, `TaoBRGNGetRegularizationType`
    321: @*/
    322: PetscErrorCode TaoBRGNSetRegularizationType(Tao tao, TaoBRGNRegularizationType type)
    323: {
    324:   PetscFunctionBegin;
    327:   PetscTryMethod((PetscObject)tao, "TaoBRGNSetRegularizationType_C", (Tao, TaoBRGNRegularizationType), (tao, type));
    328:   PetscFunctionReturn(PETSC_SUCCESS);
    329: }

    331: static PetscErrorCode TaoSolve_BRGN(Tao tao)
    332: {
    333:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    335:   PetscFunctionBegin;
    336:   PetscCall(TaoSolve(gn->subsolver));
    337:   /* Update basic tao information from the subsolver */
    338:   /* NOTE: Due to subsolver nature of BRGN, TaoTerm cannot properly track counts */
    339:   tao->objective_term.term->nobj     = gn->subsolver->objective_term.term->nobj;
    340:   tao->objective_term.term->ngrad    = gn->subsolver->objective_term.term->ngrad;
    341:   tao->objective_term.term->nobjgrad = gn->subsolver->objective_term.term->nobjgrad;
    342:   tao->objective_term.term->nhess    = gn->subsolver->objective_term.term->nhess;

    344:   tao->nres        = gn->subsolver->nres;
    345:   tao->niter       = gn->subsolver->niter;
    346:   tao->ksp_its     = gn->subsolver->ksp_its;
    347:   tao->ksp_tot_its = gn->subsolver->ksp_tot_its;
    348:   PetscCall(TaoGetConvergedReason(gn->subsolver, &tao->reason));
    349:   /* Update vectors */
    350:   PetscCall(VecCopy(gn->subsolver->solution, tao->solution));
    351:   PetscCall(VecCopy(gn->subsolver->gradient, tao->gradient));
    352:   PetscFunctionReturn(PETSC_SUCCESS);
    353: }

    355: static PetscErrorCode TaoSetFromOptions_BRGN(Tao tao, PetscOptionItems PetscOptionsObject)
    356: {
    357:   TAO_BRGN     *gn = (TAO_BRGN *)tao->data;
    358:   TaoLineSearch ls;

    360:   PetscFunctionBegin;
    361:   PetscOptionsHeadBegin(PetscOptionsObject, "least-squares problems with regularizer: ||f(x)||^2 + lambda*g(x), g(x) = ||xk-xkm1||^2 or ||Dx||_1 or user defined function.");
    362:   PetscCall(PetscOptionsBool("-tao_brgn_mat_explicit", "switches the Hessian construction to be an explicit matrix rather than MATSHELL", "", gn->mat_explicit, &gn->mat_explicit, NULL));
    363:   PetscCall(PetscOptionsReal("-tao_brgn_regularizer_weight", "regularizer weight (default 1e-4)", "", gn->lambda, &gn->lambda, NULL));
    364:   PetscCall(PetscOptionsReal("-tao_brgn_l1_smooth_epsilon", "L1-norm smooth approximation parameter: ||x||_1 = sum(sqrt(x.^2+epsilon^2)-epsilon) (default 1e-6)", "", gn->epsilon, &gn->epsilon, NULL));
    365:   PetscCall(PetscOptionsReal("-tao_brgn_lm_downhill_lambda_change", "Factor to decrease trust region by on downhill steps", "", gn->downhill_lambda_change, &gn->downhill_lambda_change, NULL));
    366:   PetscCall(PetscOptionsReal("-tao_brgn_lm_uphill_lambda_change", "Factor to increase trust region by on uphill steps", "", gn->uphill_lambda_change, &gn->uphill_lambda_change, NULL));
    367:   PetscCall(PetscOptionsEnum("-tao_brgn_regularization_type", "regularization type", "", TaoBRGNRegularizationTypes, (PetscEnum)gn->reg_type, (PetscEnum *)&gn->reg_type, NULL));
    368:   PetscOptionsHeadEnd();
    369:   /* set unit line search direction as the default when using the lm regularizer */
    370:   if (gn->reg_type == TAOBRGN_REGULARIZATION_LM) {
    371:     PetscCall(TaoGetLineSearch(gn->subsolver, &ls));
    372:     PetscCall(TaoLineSearchSetType(ls, TAOLINESEARCHUNIT));
    373:   }
    374:   PetscCall(TaoSetFromOptions(gn->subsolver));
    375:   PetscFunctionReturn(PETSC_SUCCESS);
    376: }

    378: static PetscErrorCode TaoView_BRGN(Tao tao, PetscViewer viewer)
    379: {
    380:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;
    381:   PetscBool isascii;

    383:   PetscFunctionBegin;
    384:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    385:   if (isascii) {
    386:     PetscCall(PetscViewerASCIIPushTab(viewer));
    387:     PetscCall(PetscViewerASCIIPrintf(viewer, "Regularizer weight: %g\n", (double)gn->lambda));
    388:     PetscCall(PetscViewerASCIIPrintf(viewer, "BRGN Regularization Type: %s\n", TaoBRGNRegularizationTypes[gn->reg_type]));
    389:     switch (gn->reg_type) {
    390:     case TAOBRGN_REGULARIZATION_L1DICT:
    391:       PetscCall(PetscViewerASCIIPrintf(viewer, "L1 smooth epsilon: %g\n", (double)gn->epsilon));
    392:       break;
    393:     case TAOBRGN_REGULARIZATION_LM:
    394:       PetscCall(PetscViewerASCIIPrintf(viewer, "Downhill trust region decrease factor:: %g\n", (double)gn->downhill_lambda_change));
    395:       PetscCall(PetscViewerASCIIPrintf(viewer, "Uphill trust region increase factor:: %g\n", (double)gn->uphill_lambda_change));
    396:       break;
    397:     case TAOBRGN_REGULARIZATION_L2PROX:
    398:     case TAOBRGN_REGULARIZATION_L2PURE:
    399:     case TAOBRGN_REGULARIZATION_USER:
    400:     default:
    401:       break;
    402:     }
    403:     PetscCall(PetscViewerASCIIPopTab(viewer));
    404:   }
    405:   PetscCall(PetscViewerASCIIPushTab(viewer));
    406:   PetscCall(TaoView(gn->subsolver, viewer));
    407:   PetscCall(PetscViewerASCIIPopTab(viewer));
    408:   PetscFunctionReturn(PETSC_SUCCESS);
    409: }

    411: static PetscErrorCode TaoSetUp_BRGN(Tao tao)
    412: {
    413:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;
    414:   PetscBool is_bnls, is_bntr, is_bntl;
    415:   PetscInt  n, N, K; /* dict has size K*N*/

    417:   PetscFunctionBegin;
    418:   PetscCheck(tao->ls_res, PetscObjectComm((PetscObject)tao), PETSC_ERR_ORDER, "TaoSetResidualRoutine() must be called before setup!");
    419:   PetscCall(PetscObjectTypeCompare((PetscObject)gn->subsolver, TAOBNLS, &is_bnls));
    420:   PetscCall(PetscObjectTypeCompare((PetscObject)gn->subsolver, TAOBNTR, &is_bntr));
    421:   PetscCall(PetscObjectTypeCompare((PetscObject)gn->subsolver, TAOBNTL, &is_bntl));
    422:   PetscCheck((!is_bnls && !is_bntr && !is_bntl) || tao->ls_jac, PetscObjectComm((PetscObject)tao), PETSC_ERR_ORDER, "TaoSetResidualJacobianRoutine() must be called before setup!");
    423:   if (!tao->gradient) PetscCall(VecDuplicate(tao->solution, &tao->gradient));
    424:   if (!gn->x_work) PetscCall(VecDuplicate(tao->solution, &gn->x_work));
    425:   if (!gn->r_work) PetscCall(VecDuplicate(tao->ls_res, &gn->r_work));
    426:   if (!gn->x_old) {
    427:     PetscCall(VecDuplicate(tao->solution, &gn->x_old));
    428:     PetscCall(VecSet(gn->x_old, 0.0));
    429:   }

    431:   if (TAOBRGN_REGULARIZATION_L1DICT == gn->reg_type) {
    432:     if (!gn->y) {
    433:       if (gn->D) {
    434:         PetscCall(MatGetSize(gn->D, &K, &N)); /* Shell matrices still must have sizes defined. K = N for identity matrix, K=N-1 or N for gradient matrix */
    435:         PetscCall(MatCreateVecs(gn->D, NULL, &gn->y));
    436:       } else {
    437:         PetscCall(VecDuplicate(tao->solution, &gn->y)); /* If user does not setup dict matrix, use identity matrix, K=N */
    438:       }
    439:       PetscCall(VecSet(gn->y, 0.0));
    440:     }
    441:     if (!gn->y_work) PetscCall(VecDuplicate(gn->y, &gn->y_work));
    442:     if (!gn->diag) {
    443:       PetscCall(VecDuplicate(gn->y, &gn->diag));
    444:       PetscCall(VecSet(gn->diag, 0.0));
    445:     }
    446:   }
    447:   if (TAOBRGN_REGULARIZATION_LM == gn->reg_type) {
    448:     if (!gn->diag) PetscCall(MatCreateVecs(tao->ls_jac, &gn->diag, NULL));
    449:     if (!gn->damping) PetscCall(MatCreateVecs(tao->ls_jac, &gn->damping, NULL));
    450:   }

    452:   if (!tao->setupcalled) {
    453:     /* Hessian setup */
    454:     if (gn->mat_explicit) {
    455:       PetscCall(TaoComputeResidualJacobian(tao, tao->solution, tao->ls_jac, tao->ls_jac_pre));
    456:       PetscCall(MatTransposeMatMult(tao->ls_jac, tao->ls_jac, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &gn->H));
    457:     } else {
    458:       PetscCall(VecGetLocalSize(tao->solution, &n));
    459:       PetscCall(VecGetSize(tao->solution, &N));
    460:       PetscCall(MatCreate(PetscObjectComm((PetscObject)tao), &gn->H));
    461:       PetscCall(MatSetSizes(gn->H, n, n, N, N));
    462:       PetscCall(MatSetType(gn->H, MATSHELL));
    463:       PetscCall(MatSetOption(gn->H, MAT_SYMMETRIC, PETSC_TRUE));
    464:       PetscCall(MatShellSetOperation(gn->H, MATOP_MULT, (PetscErrorCodeFn *)GNHessianProd));
    465:       PetscCall(MatShellSetContext(gn->H, gn));
    466:     }
    467:     PetscCall(MatSetUp(gn->H));
    468:     /* Subsolver setup,include initial vector and dictionary D */
    469:     PetscCall(TaoSetUpdate(gn->subsolver, GNHookFunction, gn));
    470:     PetscCall(TaoSetSolution(gn->subsolver, tao->solution));
    471:     if (tao->bounded) PetscCall(TaoSetVariableBounds(gn->subsolver, tao->XL, tao->XU));
    472:     PetscCall(TaoSetResidualRoutine(gn->subsolver, tao->ls_res, tao->ops->computeresidual, tao->user_lsresP));
    473:     PetscCall(TaoSetJacobianResidualRoutine(gn->subsolver, tao->ls_jac, tao->ls_jac, tao->ops->computeresidualjacobian, tao->user_lsjacP));
    474:     PetscCall(TaoSetObjectiveAndGradient(gn->subsolver, NULL, GNObjectiveGradientEval, gn));
    475:     PetscCall(TaoSetHessian(gn->subsolver, gn->H, gn->H, GNComputeHessian, gn));
    476:     /* Propagate some options down */
    477:     PetscCall(TaoSetTolerances(gn->subsolver, tao->gatol, tao->grtol, tao->gttol));
    478:     PetscCall(TaoSetMaximumIterations(gn->subsolver, tao->max_it));
    479:     PetscCall(TaoSetMaximumFunctionEvaluations(gn->subsolver, tao->max_funcs));
    480:     PetscCall(TaoSetUp(gn->subsolver));
    481:   }
    482:   PetscFunctionReturn(PETSC_SUCCESS);
    483: }

    485: static PetscErrorCode TaoDestroy_BRGN(Tao tao)
    486: {
    487:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    489:   PetscFunctionBegin;
    490:   if (tao->setupcalled) {
    491:     PetscCall(VecDestroy(&tao->gradient));
    492:     PetscCall(VecDestroy(&gn->x_work));
    493:     PetscCall(VecDestroy(&gn->r_work));
    494:     PetscCall(VecDestroy(&gn->x_old));
    495:     PetscCall(VecDestroy(&gn->diag));
    496:     PetscCall(VecDestroy(&gn->y));
    497:     PetscCall(VecDestroy(&gn->y_work));
    498:   }
    499:   PetscCall(VecDestroy(&gn->damping));
    500:   PetscCall(VecDestroy(&gn->diag));
    501:   PetscCall(MatDestroy(&gn->H));
    502:   PetscCall(MatDestroy(&gn->D));
    503:   PetscCall(MatDestroy(&gn->Hreg));
    504:   PetscCall(TaoDestroy(&gn->subsolver));
    505:   gn->parent = NULL;
    506:   PetscCall(PetscFree(tao->data));
    507:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNGetRegularizationType_C", NULL));
    508:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizationType_C", NULL));
    509:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNGetDampingVector_C", NULL));
    510:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetDictionaryMatrix_C", NULL));
    511:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNGetSubsolver_C", NULL));
    512:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizerWeight_C", NULL));
    513:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetL1SmoothEpsilon_C", NULL));
    514:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizerObjectiveAndGradientRoutine_C", NULL));
    515:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizerHessianRoutine_C", NULL));
    516:   PetscFunctionReturn(PETSC_SUCCESS);
    517: }

    519: /*@
    520:   TaoBRGNGetSubsolver - Get the pointer to the subsolver inside a `TAOBRGN`

    522:   Collective

    524:   Input Parameters:
    525: + tao       - the Tao solver context
    526: - subsolver - the `Tao` sub-solver context

    528:   Level: advanced

    530: .seealso: `Tao`, `Mat`, `TAOBRGN`
    531: @*/
    532: PetscErrorCode TaoBRGNGetSubsolver(Tao tao, Tao *subsolver)
    533: {
    534:   PetscFunctionBegin;
    536:   PetscUseMethod((PetscObject)tao, "TaoBRGNGetSubsolver_C", (Tao, Tao *), (tao, subsolver));
    537:   PetscFunctionReturn(PETSC_SUCCESS);
    538: }

    540: static PetscErrorCode TaoBRGNGetSubsolver_BRGN(Tao tao, Tao *subsolver)
    541: {
    542:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    544:   PetscFunctionBegin;
    545:   *subsolver = gn->subsolver;
    546:   PetscFunctionReturn(PETSC_SUCCESS);
    547: }

    549: /*@
    550:   TaoBRGNSetRegularizerWeight - Set the regularizer weight for the Gauss-Newton least-squares algorithm

    552:   Collective

    554:   Input Parameters:
    555: + tao    - the `Tao` solver context
    556: - lambda - L1-norm regularizer weight

    558:   Level: beginner

    560: .seealso: `Tao`, `Mat`, `TAOBRGN`
    561: @*/
    562: PetscErrorCode TaoBRGNSetRegularizerWeight(Tao tao, PetscReal lambda)
    563: {
    564:   PetscFunctionBegin;
    567:   PetscTryMethod((PetscObject)tao, "TaoBRGNSetRegularizerWeight_C", (Tao, PetscReal), (tao, lambda));
    568:   PetscFunctionReturn(PETSC_SUCCESS);
    569: }

    571: static PetscErrorCode TaoBRGNSetRegularizerWeight_BRGN(Tao tao, PetscReal lambda)
    572: {
    573:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    575:   PetscFunctionBegin;
    576:   gn->lambda = lambda;
    577:   PetscFunctionReturn(PETSC_SUCCESS);
    578: }

    580: /*@
    581:   TaoBRGNSetL1SmoothEpsilon - Set the L1-norm smooth approximation parameter for L1-regularized least-squares algorithm

    583:   Collective

    585:   Input Parameters:
    586: + tao     - the `Tao` solver context
    587: - epsilon - L1-norm smooth approximation parameter

    589:   Level: advanced

    591: .seealso: `Tao`, `Mat`, `TAOBRGN`
    592: @*/
    593: PetscErrorCode TaoBRGNSetL1SmoothEpsilon(Tao tao, PetscReal epsilon)
    594: {
    595:   PetscFunctionBegin;
    598:   PetscTryMethod((PetscObject)tao, "TaoBRGNSetL1SmoothEpsilon_C", (Tao, PetscReal), (tao, epsilon));
    599:   PetscFunctionReturn(PETSC_SUCCESS);
    600: }

    602: static PetscErrorCode TaoBRGNSetL1SmoothEpsilon_BRGN(Tao tao, PetscReal epsilon)
    603: {
    604:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    606:   PetscFunctionBegin;
    607:   gn->epsilon = epsilon;
    608:   PetscFunctionReturn(PETSC_SUCCESS);
    609: }

    611: /*@
    612:   TaoBRGNSetDictionaryMatrix - bind the dictionary matrix from user application context to gn->D, for compressed sensing (with least-squares problem)

    614:   Input Parameters:
    615: + tao  - the `Tao` context
    616: - dict - the user specified dictionary matrix.  We allow to set a `NULL` dictionary, which means identity matrix by default

    618:   Level: advanced

    620: .seealso: `Tao`, `Mat`, `TAOBRGN`
    621: @*/
    622: PetscErrorCode TaoBRGNSetDictionaryMatrix(Tao tao, Mat dict)
    623: {
    624:   PetscFunctionBegin;
    626:   PetscTryMethod((PetscObject)tao, "TaoBRGNSetDictionaryMatrix_C", (Tao, Mat), (tao, dict));
    627:   PetscFunctionReturn(PETSC_SUCCESS);
    628: }

    630: static PetscErrorCode TaoBRGNSetDictionaryMatrix_BRGN(Tao tao, Mat dict)
    631: {
    632:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    634:   PetscFunctionBegin;
    635:   if (dict) {
    637:     PetscCheckSameComm(tao, 1, dict, 2);
    638:     PetscCall(PetscObjectReference((PetscObject)dict));
    639:   }
    640:   PetscCall(MatDestroy(&gn->D));
    641:   gn->D = dict;
    642:   PetscFunctionReturn(PETSC_SUCCESS);
    643: }

    645: /*@C
    646:   TaoBRGNSetRegularizerObjectiveAndGradientRoutine - Sets the user-defined regularizer call-back
    647:   function into the algorithm.

    649:   Input Parameters:
    650: + tao  - the Tao context
    651: . func - function pointer for the regularizer value and gradient evaluation
    652: - ctx  - user context for the regularizer

    654:   Calling sequence:
    655: + tao - the `Tao` context
    656: . u   - the location at which to compute the objective and gradient
    657: . val - location to store objective function value
    658: . g   - location to store gradient
    659: - ctx - user context for the regularizer Hessian

    661:   Level: advanced

    663: .seealso: `Tao`, `Mat`, `TAOBRGN`
    664: @*/
    665: PetscErrorCode TaoBRGNSetRegularizerObjectiveAndGradientRoutine(Tao tao, PetscErrorCode (*func)(Tao tao, Vec u, PetscReal *val, Vec g, PetscCtx ctx), PetscCtx ctx)
    666: {
    667:   PetscFunctionBegin;
    669:   PetscTryMethod((PetscObject)tao, "TaoBRGNSetRegularizerObjectiveAndGradientRoutine_C", (Tao, PetscErrorCode (*)(Tao, Vec, PetscReal *, Vec, void *), void *), (tao, func, ctx));
    670:   PetscFunctionReturn(PETSC_SUCCESS);
    671: }

    673: static PetscErrorCode TaoBRGNSetRegularizerObjectiveAndGradientRoutine_BRGN(Tao tao, PetscErrorCode (*func)(Tao tao, Vec u, PetscReal *val, Vec g, PetscCtx ctx), PetscCtx ctx)
    674: {
    675:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    677:   PetscFunctionBegin;
    678:   if (ctx) gn->reg_obj_ctx = ctx;
    679:   if (func) gn->regularizerobjandgrad = func;
    680:   PetscFunctionReturn(PETSC_SUCCESS);
    681: }

    683: /*@C
    684:   TaoBRGNSetRegularizerHessianRoutine - Sets the user-defined regularizer call-back
    685:   function into the algorithm.

    687:   Input Parameters:
    688: + tao  - the `Tao` context
    689: . Hreg - user-created matrix for the Hessian of the regularization term
    690: . func - function pointer for the regularizer Hessian evaluation
    691: - ctx  - user context for the regularizer Hessian

    693:   Calling sequence:
    694: + tao  - the `Tao` context
    695: . u    - the location at which to compute the Hessian
    696: . Hreg - user-created matrix for the Hessian of the regularization term
    697: - ctx  - user context for the regularizer Hessian

    699:   Level: advanced

    701: .seealso: `Tao`, `Mat`, `TAOBRGN`
    702: @*/
    703: PetscErrorCode TaoBRGNSetRegularizerHessianRoutine(Tao tao, Mat Hreg, PetscErrorCode (*func)(Tao tao, Vec u, Mat Hreg, PetscCtx ctx), PetscCtx ctx)
    704: {
    705:   PetscFunctionBegin;
    707:   PetscTryMethod((PetscObject)tao, "TaoBRGNSetRegularizerHessianRoutine_C", (Tao, Mat, PetscErrorCode (*)(Tao, Vec, Mat, void *), void *), (tao, Hreg, func, ctx));
    708:   PetscFunctionReturn(PETSC_SUCCESS);
    709: }

    711: static PetscErrorCode TaoBRGNSetRegularizerHessianRoutine_BRGN(Tao tao, Mat Hreg, PetscErrorCode (*func)(Tao tao, Vec u, Mat Hreg, PetscCtx ctx), PetscCtx ctx)
    712: {
    713:   TAO_BRGN *gn = (TAO_BRGN *)tao->data;

    715:   PetscFunctionBegin;
    716:   if (Hreg) {
    718:     PetscCheckSameComm(tao, 1, Hreg, 2);
    719:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_WRONG, "NULL Hessian detected! User must provide valid Hessian for the regularizer.");
    720:   if (ctx) gn->reg_hess_ctx = ctx;
    721:   if (func) gn->regularizerhessian = func;
    722:   if (Hreg) {
    723:     PetscCall(PetscObjectReference((PetscObject)Hreg));
    724:     PetscCall(MatDestroy(&gn->Hreg));
    725:     gn->Hreg = Hreg;
    726:   }
    727:   PetscFunctionReturn(PETSC_SUCCESS);
    728: }

    730: /*MC
    731:   TAOBRGN - Bounded Regularized Gauss-Newton method for solving nonlinear least-squares
    732:             problems with bound constraints. This algorithm is a thin wrapper around `TAOBNTL`
    733:             that constructs the Gauss-Newton problem with the user-provided least-squares
    734:             residual and Jacobian. The algorithm offers an L2-norm (`l2pure`), L2-norm proximal point (`l2prox`)
    735:             regularizer, and L1-norm dictionary regularizer (`l1dict`), where we approximate the
    736:             L1-norm $\|x\|_1$ by $\sum_i{\sqrt{x_i^2+\epsilon^2}-\epsilon}$ with a small positive number $\epsilon$.
    737:             Also offered is the `lm` regularizer which uses a scaled diagonal of $J^T J$.
    738:             With the `lm` regularizer, `TAOBRGN` is a Levenberg-Marquardt optimizer.
    739:             The user can also provide their own regularization function.

    741:   Options Database Keys:
    742: + -tao_brgn_regularization_type (user|l2prox|l2pure|l1dict|lm) - regularization type, default `l2prox`
    743: . -tao_brgn_regularizer_weight                                 - regularizer weight (default 1e-4)
    744: - -tao_brgn_l1_smooth_epsilon                                  - L1-norm smooth approximation parameter: $\|x\|_1 = \sum_i{\sqrt{x_i^2+\epsilon^2}-\epsilon}$ (default 1e-6)

    746:   Level: beginner

    748: .seealso: `Tao`, `TaoBRGNGetSubsolver()`, `TaoBRGNSetRegularizerWeight()`, `TaoBRGNSetL1SmoothEpsilon()`, `TaoBRGNSetDictionaryMatrix()`,
    749:           `TaoBRGNSetRegularizerObjectiveAndGradientRoutine()`, `TaoBRGNSetRegularizerHessianRoutine()`
    750: M*/
    751: PETSC_EXTERN PetscErrorCode TaoCreate_BRGN(Tao tao)
    752: {
    753:   TAO_BRGN   *gn;
    754:   const char *prefix;

    756:   PetscFunctionBegin;
    757:   PetscCall(PetscNew(&gn));

    759:   tao->ops->destroy        = TaoDestroy_BRGN;
    760:   tao->ops->setup          = TaoSetUp_BRGN;
    761:   tao->ops->setfromoptions = TaoSetFromOptions_BRGN;
    762:   tao->ops->view           = TaoView_BRGN;
    763:   tao->ops->solve          = TaoSolve_BRGN;
    764:   tao->uses_gradient       = PETSC_TRUE;

    766:   PetscCall(TaoParametersInitialize(tao));

    768:   tao->data                  = gn;
    769:   gn->reg_type               = TAOBRGN_REGULARIZATION_L2PROX;
    770:   gn->lambda                 = 1e-4;
    771:   gn->epsilon                = 1e-6;
    772:   gn->downhill_lambda_change = 1. / 5.;
    773:   gn->uphill_lambda_change   = 1.5;
    774:   gn->parent                 = tao;

    776:   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)tao, &prefix));
    777:   PetscCall(TaoCreate(PetscObjectComm((PetscObject)tao), &gn->subsolver));
    778:   PetscCall(TaoSetType(gn->subsolver, TAOBNLS));
    779:   PetscCall(TaoSetOptionsPrefix(gn->subsolver, prefix));
    780:   PetscCall(TaoAppendOptionsPrefix(gn->subsolver, "tao_brgn_subsolver_"));
    781:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNGetRegularizationType_C", TaoBRGNGetRegularizationType_BRGN));
    782:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizationType_C", TaoBRGNSetRegularizationType_BRGN));
    783:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNGetDampingVector_C", TaoBRGNGetDampingVector_BRGN));
    784:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetDictionaryMatrix_C", TaoBRGNSetDictionaryMatrix_BRGN));
    785:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNGetSubsolver_C", TaoBRGNGetSubsolver_BRGN));
    786:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizerWeight_C", TaoBRGNSetRegularizerWeight_BRGN));
    787:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetL1SmoothEpsilon_C", TaoBRGNSetL1SmoothEpsilon_BRGN));
    788:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizerObjectiveAndGradientRoutine_C", TaoBRGNSetRegularizerObjectiveAndGradientRoutine_BRGN));
    789:   PetscCall(PetscObjectComposeFunction((PetscObject)tao, "TaoBRGNSetRegularizerHessianRoutine_C", TaoBRGNSetRegularizerHessianRoutine_BRGN));
    790:   PetscFunctionReturn(PETSC_SUCCESS);
    791: }