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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    640: /*TEST

    642:   build:
    643:     requires: !complex

    645:   test:
    646:     suffix: 0
    647:     args:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    705:   test:
    706:     suffix: soft
    707:     args: -taylor 0 -soft 1
    708:     output_file: output/empty.out

    710:   test:
    711:     suffix: soft_view_10
    712:     args: -taylor 0 -soft 1 -tao_view -tao_max_funcs 10

    714: TEST*/