Actual source code: ex4.c

      1: static char help[] = "Simple example to test separable objective optimizers.\n";

      3: #include <petsc.h>
      4: #include <petsctao.h>
      5: #include <petscvec.h>
      6: #include <petscmath.h>

      8: #define NWORKLEFT  4
      9: #define NWORKRIGHT 12

     11: typedef struct _UserCtx {
     12:   PetscInt    m;       /* The row dimension of F */
     13:   PetscInt    n;       /* The column dimension of F */
     14:   PetscInt    matops;  /* Matrix format. 0 for stencil, 1 for random */
     15:   PetscInt    iter;    /* Number of iterations for ADMM */
     16:   PetscReal   hStart;  /* Starting point for Taylor test */
     17:   PetscReal   hFactor; /* Taylor test step factor */
     18:   PetscReal   hMin;    /* Taylor test end goal */
     19:   PetscReal   alpha;   /* regularization constant applied to || x ||_p */
     20:   PetscReal   eps;     /* small constant for approximating gradient of || x ||_1 */
     21:   PetscReal   mu;      /* the augmented Lagrangian term in ADMM */
     22:   PetscReal   abstol;
     23:   PetscReal   reltol;
     24:   Mat         F;                     /* matrix in least squares component $(1/2) * || F x - d ||_2^2$ */
     25:   Mat         W;                     /* Workspace matrix. ATA */
     26:   Mat         Hm;                    /* Hessian Misfit*/
     27:   Mat         Hr;                    /* Hessian Reg*/
     28:   Vec         d;                     /* RHS in least squares component $(1/2) * || F x - d ||_2^2$ */
     29:   Vec         workLeft[NWORKLEFT];   /* Workspace for temporary vec */
     30:   Vec         workRight[NWORKRIGHT]; /* Workspace for temporary vec */
     31:   NormType    p;
     32:   PetscRandom rctx;
     33:   PetscBool   soft;
     34:   PetscBool   taylor;   /* Flag to determine whether to run Taylor test or not */
     35:   PetscBool   use_admm; /* Flag to determine whether to run Taylor test or not */
     36: } *UserCtx;

     38: static PetscErrorCode CreateRHS(UserCtx ctx)
     39: {
     40:   PetscFunctionBegin;
     41:   /* build the rhs d in ctx */
     42:   PetscCall(VecCreate(PETSC_COMM_WORLD, &ctx->d));
     43:   PetscCall(VecSetSizes(ctx->d, PETSC_DECIDE, ctx->m));
     44:   PetscCall(VecSetFromOptions(ctx->d));
     45:   PetscCall(VecSetRandom(ctx->d, ctx->rctx));
     46:   PetscFunctionReturn(PETSC_SUCCESS);
     47: }

     49: static PetscErrorCode CreateMatrix(UserCtx ctx)
     50: {
     51:   PetscInt      Istart, Iend, i, j, Ii, gridN, I_n, I_s, I_e, I_w;
     52:   PetscLogStage stage;

     54:   PetscFunctionBegin;
     55:   /* build the matrix F in ctx */
     56:   PetscCall(MatCreate(PETSC_COMM_WORLD, &ctx->F));
     57:   PetscCall(MatSetSizes(ctx->F, PETSC_DECIDE, PETSC_DECIDE, ctx->m, ctx->n));
     58:   PetscCall(MatSetType(ctx->F, MATAIJ));                          /* TODO: Decide specific SetType other than dummy*/
     59:   PetscCall(MatMPIAIJSetPreallocation(ctx->F, 5, NULL, 5, NULL)); /*TODO: some number other than 5?*/
     60:   PetscCall(MatSeqAIJSetPreallocation(ctx->F, 5, NULL));
     61:   PetscCall(MatSetUp(ctx->F));
     62:   PetscCall(MatGetOwnershipRange(ctx->F, &Istart, &Iend));
     63:   PetscCall(PetscLogStageRegister("Assembly", &stage));
     64:   PetscCall(PetscLogStagePush(stage));

     66:   /* Set matrix elements in  2-D five point stencil format. */
     67:   if (!ctx->matops) {
     68:     PetscCheck(ctx->m == ctx->n, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "Stencil matrix must be square");
     69:     gridN = (PetscInt)PetscSqrtReal((PetscReal)ctx->m);
     70:     PetscCheck(gridN * gridN == ctx->m, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "Number of rows must be square");
     71:     for (Ii = Istart; Ii < Iend; Ii++) {
     72:       i   = Ii / gridN;
     73:       j   = Ii % gridN;
     74:       I_n = i * gridN + j + 1;
     75:       if (j + 1 >= gridN) I_n = -1;
     76:       I_s = i * gridN + j - 1;
     77:       if (j - 1 < 0) I_s = -1;
     78:       I_e = (i + 1) * gridN + j;
     79:       if (i + 1 >= gridN) I_e = -1;
     80:       I_w = (i - 1) * gridN + j;
     81:       if (i - 1 < 0) I_w = -1;
     82:       PetscCall(MatSetValue(ctx->F, Ii, Ii, 4., INSERT_VALUES));
     83:       PetscCall(MatSetValue(ctx->F, Ii, I_n, -1., INSERT_VALUES));
     84:       PetscCall(MatSetValue(ctx->F, Ii, I_s, -1., INSERT_VALUES));
     85:       PetscCall(MatSetValue(ctx->F, Ii, I_e, -1., INSERT_VALUES));
     86:       PetscCall(MatSetValue(ctx->F, Ii, I_w, -1., INSERT_VALUES));
     87:     }
     88:   } else PetscCall(MatSetRandom(ctx->F, ctx->rctx));
     89:   PetscCall(MatAssemblyBegin(ctx->F, MAT_FINAL_ASSEMBLY));
     90:   PetscCall(MatAssemblyEnd(ctx->F, MAT_FINAL_ASSEMBLY));
     91:   PetscCall(PetscLogStagePop());
     92:   /* Stencil matrix is symmetric. Setting symmetric flag for ICC/Cholesky preconditioner */
     93:   if (!ctx->matops) PetscCall(MatSetOption(ctx->F, MAT_SYMMETRIC, PETSC_TRUE));
     94:   PetscCall(MatTransposeMatMult(ctx->F, ctx->F, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &ctx->W));
     95:   /* Setup Hessian Workspace in same shape as W */
     96:   PetscCall(MatDuplicate(ctx->W, MAT_DO_NOT_COPY_VALUES, &ctx->Hm));
     97:   PetscCall(MatDuplicate(ctx->W, MAT_DO_NOT_COPY_VALUES, &ctx->Hr));
     98:   PetscFunctionReturn(PETSC_SUCCESS);
     99: }

    101: static PetscErrorCode SetupWorkspace(UserCtx ctx)
    102: {
    103:   PetscInt i;

    105:   PetscFunctionBegin;
    106:   PetscCall(MatCreateVecs(ctx->F, &ctx->workLeft[0], &ctx->workRight[0]));
    107:   for (i = 1; i < NWORKLEFT; i++) PetscCall(VecDuplicate(ctx->workLeft[0], &ctx->workLeft[i]));
    108:   for (i = 1; i < NWORKRIGHT; i++) PetscCall(VecDuplicate(ctx->workRight[0], &ctx->workRight[i]));
    109:   PetscFunctionReturn(PETSC_SUCCESS);
    110: }

    112: static PetscErrorCode ConfigureContext(UserCtx ctx)
    113: {
    114:   PetscFunctionBegin;
    115:   ctx->m        = 16;
    116:   ctx->n        = 16;
    117:   ctx->eps      = 1.e-3;
    118:   ctx->abstol   = 1.e-4;
    119:   ctx->reltol   = 1.e-2;
    120:   ctx->hStart   = 1.;
    121:   ctx->hMin     = 1.e-3;
    122:   ctx->hFactor  = 0.5;
    123:   ctx->alpha    = 1.;
    124:   ctx->mu       = 1.0;
    125:   ctx->matops   = 0;
    126:   ctx->iter     = 10;
    127:   ctx->p        = NORM_2;
    128:   ctx->soft     = PETSC_FALSE;
    129:   ctx->taylor   = PETSC_TRUE;
    130:   ctx->use_admm = PETSC_FALSE;
    131:   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Configure separable objection example", "ex4.c");
    132:   PetscCall(PetscOptionsInt("-m", "The row dimension of matrix F", "ex4.c", ctx->m, &ctx->m, NULL));
    133:   PetscCall(PetscOptionsInt("-n", "The column dimension of matrix F", "ex4.c", ctx->n, &ctx->n, NULL));
    134:   PetscCall(PetscOptionsInt("-matrix_format", "Decide format of F matrix. 0 for stencil, 1 for random", "ex4.c", ctx->matops, &ctx->matops, NULL));
    135:   PetscCall(PetscOptionsInt("-iter", "Iteration number ADMM", "ex4.c", ctx->iter, &ctx->iter, NULL));
    136:   PetscCall(PetscOptionsReal("-alpha", "The regularization multiplier. 1 default", "ex4.c", ctx->alpha, &ctx->alpha, NULL));
    137:   PetscCall(PetscOptionsReal("-epsilon", "The small constant added to |x_i| in the denominator to approximate the gradient of ||x||_1", "ex4.c", ctx->eps, &ctx->eps, NULL));
    138:   PetscCall(PetscOptionsReal("-mu", "The augmented lagrangian multiplier in ADMM", "ex4.c", ctx->mu, &ctx->mu, NULL));
    139:   PetscCall(PetscOptionsReal("-hStart", "Taylor test starting point. 1 default.", "ex4.c", ctx->hStart, &ctx->hStart, NULL));
    140:   PetscCall(PetscOptionsReal("-hFactor", "Taylor test multiplier factor. 0.5 default", "ex4.c", ctx->hFactor, &ctx->hFactor, NULL));
    141:   PetscCall(PetscOptionsReal("-hMin", "Taylor test ending condition. 1.e-3 default", "ex4.c", ctx->hMin, &ctx->hMin, NULL));
    142:   PetscCall(PetscOptionsReal("-abstol", "Absolute stopping criterion for ADMM", "ex4.c", ctx->abstol, &ctx->abstol, NULL));
    143:   PetscCall(PetscOptionsReal("-reltol", "Relative stopping criterion for ADMM", "ex4.c", ctx->reltol, &ctx->reltol, NULL));
    144:   PetscCall(PetscOptionsBool("-taylor", "Flag for Taylor test. Default is true.", "ex4.c", ctx->taylor, &ctx->taylor, NULL));
    145:   PetscCall(PetscOptionsBool("-soft", "Flag for testing soft threshold no-op case. Default is false.", "ex4.c", ctx->soft, &ctx->soft, NULL));
    146:   PetscCall(PetscOptionsBool("-use_admm", "Use the ADMM solver in this example.", "ex4.c", ctx->use_admm, &ctx->use_admm, NULL));
    147:   PetscCall(PetscOptionsEnum("-p", "Norm type.", "ex4.c", NormTypes, (PetscEnum)ctx->p, (PetscEnum *)&ctx->p, NULL));
    148:   PetscOptionsEnd();
    149:   /* Creating random ctx */
    150:   PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &ctx->rctx));
    151:   PetscCall(PetscRandomSetFromOptions(ctx->rctx));
    152:   PetscCall(CreateMatrix(ctx));
    153:   PetscCall(CreateRHS(ctx));
    154:   PetscCall(SetupWorkspace(ctx));
    155:   PetscFunctionReturn(PETSC_SUCCESS);
    156: }

    158: static PetscErrorCode DestroyContext(UserCtx *ctx)
    159: {
    160:   PetscInt i;

    162:   PetscFunctionBegin;
    163:   PetscCall(MatDestroy(&(*ctx)->F));
    164:   PetscCall(MatDestroy(&(*ctx)->W));
    165:   PetscCall(MatDestroy(&(*ctx)->Hm));
    166:   PetscCall(MatDestroy(&(*ctx)->Hr));
    167:   PetscCall(VecDestroy(&(*ctx)->d));
    168:   for (i = 0; i < NWORKLEFT; i++) PetscCall(VecDestroy(&(*ctx)->workLeft[i]));
    169:   for (i = 0; i < NWORKRIGHT; i++) PetscCall(VecDestroy(&(*ctx)->workRight[i]));
    170:   PetscCall(PetscRandomDestroy(&(*ctx)->rctx));
    171:   PetscCall(PetscFree(*ctx));
    172:   PetscFunctionReturn(PETSC_SUCCESS);
    173: }

    175: /* compute (1/2) * ||F x - d||^2 */
    176: static PetscErrorCode ObjectiveMisfit(Tao tao, Vec x, PetscReal *J, void *_ctx)
    177: {
    178:   UserCtx ctx = (UserCtx)_ctx;
    179:   Vec     y;

    181:   PetscFunctionBegin;
    182:   y = ctx->workLeft[0];
    183:   PetscCall(MatMult(ctx->F, x, y));
    184:   PetscCall(VecAXPY(y, -1., ctx->d));
    185:   PetscCall(VecDot(y, y, J));
    186:   *J *= 0.5;
    187:   PetscFunctionReturn(PETSC_SUCCESS);
    188: }

    190: /* compute V = FTFx - FTd */
    191: static PetscErrorCode GradientMisfit(Tao tao, Vec x, Vec V, void *_ctx)
    192: {
    193:   UserCtx ctx = (UserCtx)_ctx;
    194:   Vec     FTFx, FTd;

    196:   PetscFunctionBegin;
    197:   /* work1 is A^T Ax, work2 is Ab, W is A^T A*/
    198:   FTFx = ctx->workRight[0];
    199:   FTd  = ctx->workRight[1];
    200:   PetscCall(MatMult(ctx->W, x, FTFx));
    201:   PetscCall(MatMultTranspose(ctx->F, ctx->d, FTd));
    202:   PetscCall(VecWAXPY(V, -1., FTd, FTFx));
    203:   PetscFunctionReturn(PETSC_SUCCESS);
    204: }

    206: /* returns FTF */
    207: static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
    208: {
    209:   UserCtx ctx = (UserCtx)_ctx;

    211:   PetscFunctionBegin;
    212:   if (H != ctx->W) PetscCall(MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN));
    213:   if (Hpre != ctx->W) PetscCall(MatCopy(ctx->W, Hpre, DIFFERENT_NONZERO_PATTERN));
    214:   PetscFunctionReturn(PETSC_SUCCESS);
    215: }

    217: /* computes augment Lagrangian objective (with scaled dual):
    218:  * 0.5 * ||F x - d||^2  + 0.5 * mu ||x - z + u||^2 */
    219: static PetscErrorCode ObjectiveMisfitADMM(Tao tao, Vec x, PetscReal *J, void *_ctx)
    220: {
    221:   UserCtx   ctx = (UserCtx)_ctx;
    222:   PetscReal mu, workNorm, misfit;
    223:   Vec       z, u, temp;

    225:   PetscFunctionBegin;
    226:   mu   = ctx->mu;
    227:   z    = ctx->workRight[5];
    228:   u    = ctx->workRight[6];
    229:   temp = ctx->workRight[10];
    230:   /* misfit = f(x) */
    231:   PetscCall(ObjectiveMisfit(tao, x, &misfit, _ctx));
    232:   PetscCall(VecCopy(x, temp));
    233:   /* temp = x - z + u */
    234:   PetscCall(VecAXPBYPCZ(temp, -1., 1., 1., z, u));
    235:   /* workNorm = ||x - z + u||^2 */
    236:   PetscCall(VecDot(temp, temp, &workNorm));
    237:   /* augment Lagrangian objective (with scaled dual): f(x) + 0.5 * mu ||x -z + u||^2 */
    238:   *J = misfit + 0.5 * mu * workNorm;
    239:   PetscFunctionReturn(PETSC_SUCCESS);
    240: }

    242: /* computes FTFx - FTd  mu*(x - z + u) */
    243: static PetscErrorCode GradientMisfitADMM(Tao tao, Vec x, Vec V, void *_ctx)
    244: {
    245:   UserCtx   ctx = (UserCtx)_ctx;
    246:   PetscReal mu;
    247:   Vec       z, u, temp;

    249:   PetscFunctionBegin;
    250:   mu   = ctx->mu;
    251:   z    = ctx->workRight[5];
    252:   u    = ctx->workRight[6];
    253:   temp = ctx->workRight[10];
    254:   PetscCall(GradientMisfit(tao, x, V, _ctx));
    255:   PetscCall(VecCopy(x, temp));
    256:   /* temp = x - z + u */
    257:   PetscCall(VecAXPBYPCZ(temp, -1., 1., 1., z, u));
    258:   /* V =  FTFx - FTd  mu*(x - z + u) */
    259:   PetscCall(VecAXPY(V, mu, temp));
    260:   PetscFunctionReturn(PETSC_SUCCESS);
    261: }

    263: /* returns FTF + diag(mu) */
    264: static PetscErrorCode HessianMisfitADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
    265: {
    266:   UserCtx ctx = (UserCtx)_ctx;

    268:   PetscFunctionBegin;
    269:   PetscCall(MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN));
    270:   PetscCall(MatShift(H, ctx->mu));
    271:   if (Hpre != H) PetscCall(MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN));
    272:   PetscFunctionReturn(PETSC_SUCCESS);
    273: }

    275: /* computes || x ||_p (mult by 0.5 in case of NORM_2) */
    276: static PetscErrorCode ObjectiveRegularization(Tao tao, Vec x, PetscReal *J, void *_ctx)
    277: {
    278:   UserCtx   ctx = (UserCtx)_ctx;
    279:   PetscReal norm;

    281:   PetscFunctionBegin;
    282:   *J = 0;
    283:   PetscCall(VecNorm(x, ctx->p, &norm));
    284:   if (ctx->p == NORM_2) norm = 0.5 * norm * norm;
    285:   *J = ctx->alpha * norm;
    286:   PetscFunctionReturn(PETSC_SUCCESS);
    287: }

    289: /* NORM_2 Case: return x
    290:  * NORM_1 Case: x/(|x| + eps)
    291:  * Else: TODO */
    292: static PetscErrorCode GradientRegularization(Tao tao, Vec x, Vec V, void *_ctx)
    293: {
    294:   UserCtx   ctx = (UserCtx)_ctx;
    295:   PetscReal eps = ctx->eps;

    297:   PetscFunctionBegin;
    298:   if (ctx->p == NORM_2) {
    299:     PetscCall(VecCopy(x, V));
    300:   } else if (ctx->p == NORM_1) {
    301:     PetscCall(VecCopy(x, ctx->workRight[1]));
    302:     PetscCall(VecAbs(ctx->workRight[1]));
    303:     PetscCall(VecShift(ctx->workRight[1], eps));
    304:     PetscCall(VecPointwiseDivide(V, x, ctx->workRight[1]));
    305:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
    306:   PetscFunctionReturn(PETSC_SUCCESS);
    307: }

    309: /* NORM_2 Case: returns diag(mu)
    310:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)))  */
    311: static PetscErrorCode HessianRegularization(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
    312: {
    313:   UserCtx   ctx = (UserCtx)_ctx;
    314:   PetscReal eps = ctx->eps;
    315:   Vec       copy1, copy2, copy3;

    317:   PetscFunctionBegin;
    318:   if (ctx->p == NORM_2) {
    319:     /* Identity matrix scaled by mu */
    320:     PetscCall(MatZeroEntries(H));
    321:     PetscCall(MatShift(H, ctx->mu));
    322:     if (Hpre != H) {
    323:       PetscCall(MatZeroEntries(Hpre));
    324:       PetscCall(MatShift(Hpre, ctx->mu));
    325:     }
    326:   } else if (ctx->p == NORM_1) {
    327:     /* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)) */
    328:     copy1 = ctx->workRight[1];
    329:     copy2 = ctx->workRight[2];
    330:     copy3 = ctx->workRight[3];
    331:     /* copy1 : 1/sqrt(x_i^2 + eps) */
    332:     PetscCall(VecCopy(x, copy1));
    333:     PetscCall(VecPow(copy1, 2));
    334:     PetscCall(VecShift(copy1, eps));
    335:     PetscCall(VecSqrtAbs(copy1));
    336:     PetscCall(VecReciprocal(copy1));
    337:     /* copy2:  x_i^2.*/
    338:     PetscCall(VecCopy(x, copy2));
    339:     PetscCall(VecPow(copy2, 2));
    340:     /* copy3: abs(x_i^2 + eps) */
    341:     PetscCall(VecCopy(x, copy3));
    342:     PetscCall(VecPow(copy3, 2));
    343:     PetscCall(VecShift(copy3, eps));
    344:     PetscCall(VecAbs(copy3));
    345:     /* copy2: 1 - x_i^2/abs(x_i^2 + eps) */
    346:     PetscCall(VecPointwiseDivide(copy2, copy2, copy3));
    347:     PetscCall(VecScale(copy2, -1.));
    348:     PetscCall(VecShift(copy2, 1.));
    349:     PetscCall(VecAXPY(copy1, 1., copy2));
    350:     PetscCall(VecScale(copy1, ctx->mu));
    351:     PetscCall(MatZeroEntries(H));
    352:     PetscCall(MatDiagonalSet(H, copy1, INSERT_VALUES));
    353:     if (Hpre != H) {
    354:       PetscCall(MatZeroEntries(Hpre));
    355:       PetscCall(MatDiagonalSet(Hpre, copy1, INSERT_VALUES));
    356:     }
    357:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
    358:   PetscFunctionReturn(PETSC_SUCCESS);
    359: }

    361: /* NORM_2 Case: 0.5 || x ||_2 + 0.5 * mu * ||x + u - z||^2
    362:  * Else : || x ||_2 + 0.5 * mu * ||x + u - z||^2 */
    363: static PetscErrorCode ObjectiveRegularizationADMM(Tao tao, Vec z, PetscReal *J, void *_ctx)
    364: {
    365:   UserCtx   ctx = (UserCtx)_ctx;
    366:   PetscReal mu, workNorm, reg;
    367:   Vec       x, u, temp;

    369:   PetscFunctionBegin;
    370:   mu   = ctx->mu;
    371:   x    = ctx->workRight[4];
    372:   u    = ctx->workRight[6];
    373:   temp = ctx->workRight[10];
    374:   PetscCall(ObjectiveRegularization(tao, z, &reg, _ctx));
    375:   PetscCall(VecCopy(z, temp));
    376:   /* temp = x + u -z */
    377:   PetscCall(VecAXPBYPCZ(temp, 1., 1., -1., x, u));
    378:   /* workNorm = ||x + u - z ||^2 */
    379:   PetscCall(VecDot(temp, temp, &workNorm));
    380:   *J = reg + 0.5 * mu * workNorm;
    381:   PetscFunctionReturn(PETSC_SUCCESS);
    382: }

    384: /* NORM_2 Case: x - mu*(x + u - z)
    385:  * NORM_1 Case: x/(|x| + eps) - mu*(x + u - z)
    386:  * Else: TODO */
    387: static PetscErrorCode GradientRegularizationADMM(Tao tao, Vec z, Vec V, void *_ctx)
    388: {
    389:   UserCtx   ctx = (UserCtx)_ctx;
    390:   PetscReal mu;
    391:   Vec       x, u, temp;

    393:   PetscFunctionBegin;
    394:   mu   = ctx->mu;
    395:   x    = ctx->workRight[4];
    396:   u    = ctx->workRight[6];
    397:   temp = ctx->workRight[10];
    398:   PetscCall(GradientRegularization(tao, z, V, _ctx));
    399:   PetscCall(VecCopy(z, temp));
    400:   /* temp = x + u -z */
    401:   PetscCall(VecAXPBYPCZ(temp, 1., 1., -1., x, u));
    402:   PetscCall(VecAXPY(V, -mu, temp));
    403:   PetscFunctionReturn(PETSC_SUCCESS);
    404: }

    406: /* NORM_2 Case: returns diag(mu)
    407:  * NORM_1 Case: FTF + diag(mu) */
    408: static PetscErrorCode HessianRegularizationADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
    409: {
    410:   UserCtx ctx = (UserCtx)_ctx;

    412:   PetscFunctionBegin;
    413:   if (ctx->p == NORM_2) {
    414:     /* Identity matrix scaled by mu */
    415:     PetscCall(MatZeroEntries(H));
    416:     PetscCall(MatShift(H, ctx->mu));
    417:     if (Hpre != H) {
    418:       PetscCall(MatZeroEntries(Hpre));
    419:       PetscCall(MatShift(Hpre, ctx->mu));
    420:     }
    421:   } else if (ctx->p == NORM_1) {
    422:     PetscCall(HessianMisfit(tao, x, H, Hpre, (void *)ctx));
    423:     PetscCall(MatShift(H, ctx->mu));
    424:     if (Hpre != H) PetscCall(MatShift(Hpre, ctx->mu));
    425:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
    426:   PetscFunctionReturn(PETSC_SUCCESS);
    427: }

    429: /* NORM_2 Case : (1/2) * ||F x - d||^2 + 0.5 * || x ||_p
    430: *  NORM_1 Case : (1/2) * ||F x - d||^2 + || x ||_p */
    431: static PetscErrorCode ObjectiveComplete(Tao tao, Vec x, PetscReal *J, void *ctx)
    432: {
    433:   PetscReal Jm, Jr;

    435:   PetscFunctionBegin;
    436:   PetscCall(ObjectiveMisfit(tao, x, &Jm, ctx));
    437:   PetscCall(ObjectiveRegularization(tao, x, &Jr, ctx));
    438:   *J = Jm + Jr;
    439:   PetscFunctionReturn(PETSC_SUCCESS);
    440: }

    442: /* NORM_2 Case: FTFx - FTd + x
    443:  * NORM_1 Case: FTFx - FTd + x/(|x| + eps) */
    444: static PetscErrorCode GradientComplete(Tao tao, Vec x, Vec V, void *ctx)
    445: {
    446:   UserCtx cntx = (UserCtx)ctx;

    448:   PetscFunctionBegin;
    449:   PetscCall(GradientMisfit(tao, x, cntx->workRight[2], ctx));
    450:   PetscCall(GradientRegularization(tao, x, cntx->workRight[3], ctx));
    451:   PetscCall(VecWAXPY(V, 1, cntx->workRight[2], cntx->workRight[3]));
    452:   PetscFunctionReturn(PETSC_SUCCESS);
    453: }

    455: /* NORM_2 Case: diag(mu) + FTF
    456:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) + FTF  */
    457: static PetscErrorCode HessianComplete(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx)
    458: {
    459:   Mat tempH;

    461:   PetscFunctionBegin;
    462:   PetscCall(MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH));
    463:   PetscCall(HessianMisfit(tao, x, H, H, ctx));
    464:   PetscCall(HessianRegularization(tao, x, tempH, tempH, ctx));
    465:   PetscCall(MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN));
    466:   if (Hpre != H) PetscCall(MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN));
    467:   PetscCall(MatDestroy(&tempH));
    468:   PetscFunctionReturn(PETSC_SUCCESS);
    469: }

    471: static PetscErrorCode TaoSolveADMM(UserCtx ctx, Vec x)
    472: {
    473:   PetscInt  i;
    474:   PetscReal u_norm, r_norm, s_norm, primal, dual, x_norm, z_norm;
    475:   Tao       tao1, tao2;
    476:   Vec       xk, z, u, diff, zold, zdiff, temp;
    477:   PetscReal mu;

    479:   PetscFunctionBegin;
    480:   xk    = ctx->workRight[4];
    481:   z     = ctx->workRight[5];
    482:   u     = ctx->workRight[6];
    483:   diff  = ctx->workRight[7];
    484:   zold  = ctx->workRight[8];
    485:   zdiff = ctx->workRight[9];
    486:   temp  = ctx->workRight[11];
    487:   mu    = ctx->mu;
    488:   PetscCall(VecSet(u, 0.));
    489:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao1));
    490:   PetscCall(TaoSetType(tao1, TAONLS));
    491:   PetscCall(TaoSetObjective(tao1, ObjectiveMisfitADMM, (void *)ctx));
    492:   PetscCall(TaoSetGradient(tao1, NULL, GradientMisfitADMM, (void *)ctx));
    493:   PetscCall(TaoSetHessian(tao1, ctx->Hm, ctx->Hm, HessianMisfitADMM, (void *)ctx));
    494:   PetscCall(VecSet(xk, 0.));
    495:   PetscCall(TaoSetSolution(tao1, xk));
    496:   PetscCall(TaoSetOptionsPrefix(tao1, "misfit_"));
    497:   PetscCall(TaoSetFromOptions(tao1));
    498:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao2));
    499:   if (ctx->p == NORM_2) {
    500:     PetscCall(TaoSetType(tao2, TAONLS));
    501:     PetscCall(TaoSetObjective(tao2, ObjectiveRegularizationADMM, (void *)ctx));
    502:     PetscCall(TaoSetGradient(tao2, NULL, GradientRegularizationADMM, (void *)ctx));
    503:     PetscCall(TaoSetHessian(tao2, ctx->Hr, ctx->Hr, HessianRegularizationADMM, (void *)ctx));
    504:   }
    505:   PetscCall(VecSet(z, 0.));
    506:   PetscCall(TaoSetSolution(tao2, z));
    507:   PetscCall(TaoSetOptionsPrefix(tao2, "reg_"));
    508:   PetscCall(TaoSetFromOptions(tao2));

    510:   for (i = 0; i < ctx->iter; i++) {
    511:     PetscCall(VecCopy(z, zold));
    512:     PetscCall(TaoSolve(tao1)); /* Updates xk */
    513:     if (ctx->p == NORM_1) {
    514:       PetscCall(VecWAXPY(temp, 1., xk, u));
    515:       PetscCall(TaoSoftThreshold(temp, -ctx->alpha / mu, ctx->alpha / mu, z));
    516:     } else {
    517:       PetscCall(TaoSolve(tao2)); /* Update zk */
    518:     }
    519:     /* u = u + xk -z */
    520:     PetscCall(VecAXPBYPCZ(u, 1., -1., 1., xk, z));
    521:     /* r_norm : norm(x-z) */
    522:     PetscCall(VecWAXPY(diff, -1., z, xk));
    523:     PetscCall(VecNorm(diff, NORM_2, &r_norm));
    524:     /* s_norm : norm(-mu(z-zold)) */
    525:     PetscCall(VecWAXPY(zdiff, -1., zold, z));
    526:     PetscCall(VecNorm(zdiff, NORM_2, &s_norm));
    527:     s_norm = s_norm * mu;
    528:     /* primal : sqrt(n)*ABSTOL + RELTOL*max(norm(x), norm(-z))*/
    529:     PetscCall(VecNorm(xk, NORM_2, &x_norm));
    530:     PetscCall(VecNorm(z, NORM_2, &z_norm));
    531:     primal = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * PetscMax(x_norm, z_norm);
    532:     /* Duality : sqrt(n)*ABSTOL + RELTOL*norm(mu*u)*/
    533:     PetscCall(VecNorm(u, NORM_2, &u_norm));
    534:     dual = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * u_norm * mu;
    535:     PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao1), "Iter %" PetscInt_FMT " : ||x-z||: %g, mu*||z-zold||: %g\n", i, (double)r_norm, (double)s_norm));
    536:     if (r_norm < primal && s_norm < dual) break;
    537:   }
    538:   PetscCall(VecCopy(xk, x));
    539:   PetscCall(TaoDestroy(&tao1));
    540:   PetscCall(TaoDestroy(&tao2));
    541:   PetscFunctionReturn(PETSC_SUCCESS);
    542: }

    544: /* Second order Taylor remainder convergence test */
    545: static PetscErrorCode TaylorTest(UserCtx ctx, Tao tao, Vec x, PetscReal *C)
    546: {
    547:   PetscReal  h, J, temp;
    548:   PetscInt   i, j;
    549:   PetscInt   numValues;
    550:   PetscReal  Jx, Jxhat_comp, Jxhat_pred;
    551:   PetscReal *Js, *hs;
    552:   PetscReal  gdotdx;
    553:   PetscReal  minrate = PETSC_MAX_REAL;
    554:   MPI_Comm   comm    = PetscObjectComm((PetscObject)x);
    555:   Vec        g, dx, xhat;

    557:   PetscFunctionBegin;
    558:   PetscCall(VecDuplicate(x, &g));
    559:   PetscCall(VecDuplicate(x, &xhat));
    560:   /* choose a perturbation direction */
    561:   PetscCall(VecDuplicate(x, &dx));
    562:   PetscCall(VecSetRandom(dx, ctx->rctx));
    563:   /* evaluate objective at x: J(x) */
    564:   PetscCall(TaoComputeObjective(tao, x, &Jx));
    565:   /* evaluate gradient at x, save in vector g */
    566:   PetscCall(TaoComputeGradient(tao, x, g));
    567:   PetscCall(VecDot(g, dx, &gdotdx));

    569:   for (numValues = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor) numValues++;
    570:   PetscCall(PetscCalloc2(numValues, &Js, numValues, &hs));
    571:   for (i = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor, i++) {
    572:     PetscCall(VecWAXPY(xhat, h, dx, x));
    573:     PetscCall(TaoComputeObjective(tao, xhat, &Jxhat_comp));
    574:     /* J(\hat(x)) \approx J(x) + g^T (xhat - x) = J(x) + h * g^T dx */
    575:     Jxhat_pred = Jx + h * gdotdx;
    576:     /* Vector to dJdm scalar? Dot?*/
    577:     J = PetscAbsReal(Jxhat_comp - Jxhat_pred);
    578:     PetscCall(PetscPrintf(comm, "J(xhat): %g, predicted: %g, diff %g\n", (double)Jxhat_comp, (double)Jxhat_pred, (double)J));
    579:     Js[i] = J;
    580:     hs[i] = h;
    581:   }
    582:   for (j = 1; j < numValues; j++) {
    583:     temp = PetscLogReal(Js[j] / Js[j - 1]) / PetscLogReal(hs[j] / hs[j - 1]);
    584:     PetscCall(PetscPrintf(comm, "Convergence rate step %" PetscInt_FMT ": %g\n", j - 1, (double)temp));
    585:     minrate = PetscMin(minrate, temp);
    586:   }
    587:   /* If O is not ~2, then the test is wrong */
    588:   PetscCall(PetscFree2(Js, hs));
    589:   *C = minrate;
    590:   PetscCall(VecDestroy(&dx));
    591:   PetscCall(VecDestroy(&xhat));
    592:   PetscCall(VecDestroy(&g));
    593:   PetscFunctionReturn(PETSC_SUCCESS);
    594: }

    596: int main(int argc, char **argv)
    597: {
    598:   UserCtx ctx;
    599:   Tao     tao;
    600:   Vec     x;
    601:   Mat     H;

    603:   PetscFunctionBeginUser;
    604:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
    605:   PetscCall(PetscNew(&ctx));
    606:   PetscCall(ConfigureContext(ctx));
    607:   /* Define two functions that could pass as objectives to TaoSetObjective(): one
    608:    * for the misfit component, and one for the regularization component */
    609:   /* ObjectiveMisfit() and ObjectiveRegularization() */

    611:   /* Define a single function that calls both components adds them together: the complete objective,
    612:    * in the absence of a Tao implementation that handles separability */
    613:   /* ObjectiveComplete() */
    614:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
    615:   PetscCall(TaoSetType(tao, TAONM));
    616:   PetscCall(TaoSetObjective(tao, ObjectiveComplete, (void *)ctx));
    617:   PetscCall(TaoSetGradient(tao, NULL, GradientComplete, (void *)ctx));
    618:   PetscCall(MatDuplicate(ctx->W, MAT_SHARE_NONZERO_PATTERN, &H));
    619:   PetscCall(TaoSetHessian(tao, H, H, HessianComplete, (void *)ctx));
    620:   PetscCall(MatCreateVecs(ctx->F, NULL, &x));
    621:   PetscCall(VecSet(x, 0.));
    622:   PetscCall(TaoSetSolution(tao, x));
    623:   PetscCall(TaoSetFromOptions(tao));
    624:   if (ctx->use_admm) PetscCall(TaoSolveADMM(ctx, x));
    625:   else PetscCall(TaoSolve(tao));
    626:   /* examine solution */
    627:   PetscCall(VecViewFromOptions(x, NULL, "-view_sol"));
    628:   if (ctx->taylor) {
    629:     PetscReal rate;
    630:     PetscCall(TaylorTest(ctx, tao, x, &rate));
    631:   }
    632:   if (ctx->soft) { PetscCall(TaoSoftThreshold(x, 0., 0., x)); }
    633:   PetscCall(MatDestroy(&H));
    634:   PetscCall(TaoDestroy(&tao));
    635:   PetscCall(VecDestroy(&x));
    636:   PetscCall(DestroyContext(&ctx));
    637:   PetscCall(PetscFinalize());
    638:   return 0;
    639: }

    641: /*TEST

    643:   build:
    644:     requires: !complex

    646:   test:
    647:     suffix: 0
    648:     args:

    650:   test:
    651:     suffix: l1_1
    652:     args: -p 1 -tao_type lmvm -alpha 1. -epsilon 1.e-7 -m 64 -n 64 -view_sol -matrix_format 1

    654:   test:
    655:     suffix: hessian_1
    656:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nls

    658:   test:
    659:     suffix: hessian_2
    660:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nls

    662:   test:
    663:     suffix: nm_1
    664:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nm -tao_max_it 50

    666:   test:
    667:     suffix: nm_2
    668:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nm -tao_max_it 50

    670:   test:
    671:     suffix: lmvm_1
    672:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type lmvm -tao_max_it 40

    674:   test:
    675:     suffix: lmvm_2
    676:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type lmvm -tao_max_it 15

    678:   test:
    679:     suffix: soft_threshold_admm_1
    680:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm

    682:   test:
    683:     suffix: hessian_admm_1
    684:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nls -misfit_tao_type nls

    686:   test:
    687:     suffix: hessian_admm_2
    688:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nls -misfit_tao_type nls

    690:   test:
    691:     suffix: nm_admm_1
    692:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nm -misfit_tao_type nm

    694:   test:
    695:     suffix: nm_admm_2
    696:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nm -misfit_tao_type nm -iter 7

    698:   test:
    699:     suffix: lmvm_admm_1
    700:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm

    702:   test:
    703:     suffix: lmvm_admm_2
    704:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm

    706:   test:
    707:     suffix: soft
    708:     args: -taylor 0 -soft 1

    710: TEST*/