Date: 2022-03-31T04:49:08+00:00

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: {
     13:   PetscInt    m;                     /* The row dimension of F */
     14:   PetscInt    n;                     /* The column dimension of F */
     15:   PetscInt    matops;                /* Matrix format. 0 for stencil, 1 for random */
     16:   PetscInt    iter;                  /* Numer of iterations for ADMM */
     17:   PetscReal   hStart;                /* Starting point for Taylor test */
     18:   PetscReal   hFactor;               /* Taylor test step factor */
     19:   PetscReal   hMin;                  /* Taylor test end goal */
     20:   PetscReal   alpha;                 /* regularization constant applied to || x ||_p */
     21:   PetscReal   eps;                   /* small constant for approximating gradient of || x ||_1 */
     22:   PetscReal   mu;                    /* the augmented Lagrangian term in ADMM */
     23:   PetscReal   abstol;
     24:   PetscReal   reltol;
     25:   Mat         F;                     /* matrix in least squares component $(1/2) * || F x - d ||_2^2$ */
     26:   Mat         W;                     /* Workspace matrix. ATA */
     27:   Mat         Hm;                    /* Hessian Misfit*/
     28:   Mat         Hr;                    /* Hessian Reg*/
     29:   Vec         d;                     /* RHS in least squares component $(1/2) * || F x - d ||_2^2$ */
     30:   Vec         workLeft[NWORKLEFT];   /* Workspace for temporary vec */
     31:   Vec         workRight[NWORKRIGHT]; /* Workspace for temporary vec */
     32:   NormType    p;
     33:   PetscRandom rctx;
     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:   /* build the rhs d in ctx */
     41:   VecCreate(PETSC_COMM_WORLD,&(ctx->d));
     42:   VecSetSizes(ctx->d,PETSC_DECIDE,ctx->m);
     43:   VecSetFromOptions(ctx->d);
     44:   VecSetRandom(ctx->d,ctx->rctx);
     45:   return 0;
     46: }

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

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

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

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

    106:   MatCreateVecs(ctx->F, &ctx->workLeft[0], &ctx->workRight[0]);
    107:   for (i=1; i<NWORKLEFT; i++) {
    108:     VecDuplicate(ctx->workLeft[0], &(ctx->workLeft[i]));
    109:   }
    110:   for (i=1; i<NWORKRIGHT; i++) {
    111:     VecDuplicate(ctx->workRight[0], &(ctx->workRight[i]));
    112:   }
    113:   return 0;
    114: }

    116: static PetscErrorCode ConfigureContext(UserCtx ctx)
    117: {

    120:   ctx->m        = 16;
    121:   ctx->n        = 16;
    122:   ctx->eps      = 1.e-3;
    123:   ctx->abstol   = 1.e-4;
    124:   ctx->reltol   = 1.e-2;
    125:   ctx->hStart   = 1.;
    126:   ctx->hMin     = 1.e-3;
    127:   ctx->hFactor  = 0.5;
    128:   ctx->alpha    = 1.;
    129:   ctx->mu       = 1.0;
    130:   ctx->matops   = 0;
    131:   ctx->iter     = 10;
    132:   ctx->p        = NORM_2;
    133:   ctx->taylor   = PETSC_TRUE;
    134:   ctx->use_admm = PETSC_FALSE;
    135:   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Configure separable objection example", "ex4.c");
    136:   PetscOptionsInt("-m", "The row dimension of matrix F", "ex4.c", ctx->m, &(ctx->m), NULL);
    137:   PetscOptionsInt("-n", "The column dimension of matrix F", "ex4.c", ctx->n, &(ctx->n), NULL);
    138:   PetscOptionsInt("-matrix_format","Decide format of F matrix. 0 for stencil, 1 for random", "ex4.c", ctx->matops, &(ctx->matops), NULL);
    139:   PetscOptionsInt("-iter","Iteration number ADMM", "ex4.c", ctx->iter, &(ctx->iter), NULL);
    140:   PetscOptionsReal("-alpha", "The regularization multiplier. 1 default", "ex4.c", ctx->alpha, &(ctx->alpha), NULL);
    141:   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);
    142:   PetscOptionsReal("-mu", "The augmented lagrangian multiplier in ADMM", "ex4.c", ctx->mu, &(ctx->mu), NULL);
    143:   PetscOptionsReal("-hStart", "Taylor test starting point. 1 default.", "ex4.c", ctx->hStart, &(ctx->hStart), NULL);
    144:   PetscOptionsReal("-hFactor", "Taylor test multiplier factor. 0.5 default", "ex4.c", ctx->hFactor, &(ctx->hFactor), NULL);
    145:   PetscOptionsReal("-hMin", "Taylor test ending condition. 1.e-3 default", "ex4.c", ctx->hMin, &(ctx->hMin), NULL);
    146:   PetscOptionsReal("-abstol", "Absolute stopping criterion for ADMM", "ex4.c", ctx->abstol, &(ctx->abstol), NULL);
    147:   PetscOptionsReal("-reltol", "Relative stopping criterion for ADMM", "ex4.c", ctx->reltol, &(ctx->reltol), NULL);
    148:   PetscOptionsBool("-taylor","Flag for Taylor test. Default is true.", "ex4.c", ctx->taylor, &(ctx->taylor), NULL);
    149:   PetscOptionsBool("-use_admm","Use the ADMM solver in this example.", "ex4.c", ctx->use_admm, &(ctx->use_admm), NULL);
    150:   PetscOptionsEnum("-p","Norm type.", "ex4.c", NormTypes, (PetscEnum)ctx->p, (PetscEnum *) &(ctx->p), NULL);
    151:   PetscOptionsEnd();
    152:   /* Creating random ctx */
    153:   PetscRandomCreate(PETSC_COMM_WORLD,&(ctx->rctx));
    154:   PetscRandomSetFromOptions(ctx->rctx);
    155:   CreateMatrix(ctx);
    156:   CreateRHS(ctx);
    157:   SetupWorkspace(ctx);
    158:   return 0;
    159: }

    161: static PetscErrorCode DestroyContext(UserCtx *ctx)
    162: {
    163:   PetscInt       i;

    165:   MatDestroy(&((*ctx)->F));
    166:   MatDestroy(&((*ctx)->W));
    167:   MatDestroy(&((*ctx)->Hm));
    168:   MatDestroy(&((*ctx)->Hr));
    169:   VecDestroy(&((*ctx)->d));
    170:   for (i=0; i<NWORKLEFT; i++) {
    171:     VecDestroy(&((*ctx)->workLeft[i]));
    172:   }
    173:   for (i=0; i<NWORKRIGHT; i++) {
    174:     VecDestroy(&((*ctx)->workRight[i]));
    175:   }
    176:   PetscRandomDestroy(&((*ctx)->rctx));
    177:   PetscFree(*ctx);
    178:   return 0;
    179: }

    181: /* compute (1/2) * ||F x - d||^2 */
    182: static PetscErrorCode ObjectiveMisfit(Tao tao, Vec x, PetscReal *J, void *_ctx)
    183: {
    184:   UserCtx        ctx = (UserCtx) _ctx;
    185:   Vec            y;

    187:   y    = ctx->workLeft[0];
    188:   MatMult(ctx->F, x, y);
    189:   VecAXPY(y, -1., ctx->d);
    190:   VecDot(y, y, J);
    191:   *J  *= 0.5;
    192:   return 0;
    193: }

    195: /* compute V = FTFx - FTd */
    196: static PetscErrorCode GradientMisfit(Tao tao, Vec x, Vec V, void *_ctx)
    197: {
    198:   UserCtx        ctx = (UserCtx) _ctx;
    199:   Vec            FTFx, FTd;

    201:   /* work1 is A^T Ax, work2 is Ab, W is A^T A*/
    202:   FTFx = ctx->workRight[0];
    203:   FTd  = ctx->workRight[1];
    204:   MatMult(ctx->W,x,FTFx);
    205:   MatMultTranspose(ctx->F, ctx->d, FTd);
    206:   VecWAXPY(V, -1., FTd, FTFx);
    207:   return 0;
    208: }

    210: /* returns FTF */
    211: static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
    212: {
    213:   UserCtx        ctx = (UserCtx) _ctx;

    215:   if (H != ctx->W) MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);
    216:   if (Hpre != ctx->W) MatCopy(ctx->W, Hpre, DIFFERENT_NONZERO_PATTERN);
    217:   return 0;
    218: }

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

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

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

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

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

    269:   MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);
    270:   MatShift(H, ctx->mu);
    271:   if (Hpre != H) {
    272:     MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);
    273:   }
    274:   return 0;
    275: }

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

    283:   *J = 0;
    284:   VecNorm (x, ctx->p, &norm);
    285:   if (ctx->p == NORM_2) norm = 0.5 * norm * norm;
    286:   *J = ctx->alpha * norm;
    287:   return 0;
    288: }

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

    298:   if (ctx->p == NORM_2) {
    299:     VecCopy(x, V);
    300:   } else if (ctx->p == NORM_1) {
    301:     VecCopy(x, ctx->workRight[1]);
    302:     VecAbs(ctx->workRight[1]);
    303:     VecShift(ctx->workRight[1], eps);
    304:     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:   return 0;
    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:   if (ctx->p == NORM_2) {
    318:     /* Identity matrix scaled by mu */
    319:     MatZeroEntries(H);
    320:     MatShift(H,ctx->mu);
    321:     if (Hpre != H) {
    322:       MatZeroEntries(Hpre);
    323:       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:     VecCopy(x, copy1);
    332:     VecPow(copy1,2);
    333:     VecShift(copy1, eps);
    334:     VecSqrtAbs(copy1);
    335:     VecReciprocal(copy1);
    336:     /* copy2:  x_i^2.*/
    337:     VecCopy(x,copy2);
    338:     VecPow(copy2,2);
    339:     /* copy3: abs(x_i^2 + eps) */
    340:     VecCopy(x,copy3);
    341:     VecPow(copy3,2);
    342:     VecShift(copy3, eps);
    343:     VecAbs(copy3);
    344:     /* copy2: 1 - x_i^2/abs(x_i^2 + eps) */
    345:     VecPointwiseDivide(copy2, copy2,copy3);
    346:     VecScale(copy2, -1.);
    347:     VecShift(copy2, 1.);
    348:     VecAXPY(copy1,1.,copy2);
    349:     VecScale(copy1, ctx->mu);
    350:     MatZeroEntries(H);
    351:     MatDiagonalSet(H, copy1,INSERT_VALUES);
    352:     if (Hpre != H) {
    353:       MatZeroEntries(Hpre);
    354:       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:   return 0;
    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:   mu   = ctx->mu;
    369:   x    = ctx->workRight[4];
    370:   u    = ctx->workRight[6];
    371:   temp = ctx->workRight[10];
    372:   ObjectiveRegularization(tao, z, &reg, _ctx);
    373:   VecCopy(z,temp);
    374:   /* temp = x + u -z */
    375:   VecAXPBYPCZ(temp,1.,1.,-1.,x,u);
    376:   /* workNorm = ||x + u - z ||^2 */
    377:   VecDot(temp, temp, &workNorm);
    378:   *J   = reg + 0.5 * mu * workNorm;
    379:   return 0;
    380: }

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

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

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

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

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

    431:   ObjectiveMisfit(tao, x, &Jm, ctx);
    432:   ObjectiveRegularization(tao, x, &Jr, ctx);
    433:   *J   = Jm + Jr;
    434:   return 0;
    435: }

    437: /* NORM_2 Case: FTFx - FTd + x
    438:  * NORM_1 Case: FTFx - FTd + x/(|x| + eps) */
    439: static PetscErrorCode GradientComplete(Tao tao, Vec x, Vec V, void *ctx)
    440: {
    441:   UserCtx        cntx = (UserCtx) ctx;

    443:   GradientMisfit(tao, x, cntx->workRight[2], ctx);
    444:   GradientRegularization(tao, x, cntx->workRight[3], ctx);
    445:   VecWAXPY(V,1,cntx->workRight[2],cntx->workRight[3]);
    446:   return 0;
    447: }

    449: /* NORM_2 Case: diag(mu) + FTF
    450:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) + FTF  */
    451: static PetscErrorCode HessianComplete(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx)
    452: {
    453:   Mat            tempH;

    455:   MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH);
    456:   HessianMisfit(tao, x, H, H, ctx);
    457:   HessianRegularization(tao, x, tempH, tempH, ctx);
    458:   MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN);
    459:   if (Hpre != H) {
    460:     MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);
    461:   }
    462:   MatDestroy(&tempH);
    463:   return 0;
    464: }

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

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

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

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

    551:   VecDuplicate(x, &g);
    552:   VecDuplicate(x, &xhat);
    553:   /* choose a perturbation direction */
    554:   VecDuplicate(x, &dx);
    555:   VecSetRandom(dx,ctx->rctx);
    556:   /* evaluate objective at x: J(x) */
    557:   TaoComputeObjective(tao, x, &Jx);
    558:   /* evaluate gradient at x, save in vector g */
    559:   TaoComputeGradient(tao, x, g);
    560:   VecDot(g, dx, &gdotdx);

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

    589: int main(int argc, char ** argv)
    590: {
    591:   UserCtx ctx;
    592:   Tao     tao;
    593:   Vec     x;
    594:   Mat     H;

    596:   PetscInitialize(&argc, &argv, NULL,help);
    597:   PetscNew(&ctx);
    598:   ConfigureContext(ctx);
    599:   /* Define two functions that could pass as objectives to TaoSetObjective(): one
    600:    * for the misfit component, and one for the regularization component */
    601:   /* ObjectiveMisfit() and ObjectiveRegularization() */

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

    633: /*TEST

    635:   build:
    636:     requires: !complex

    638:   test:
    639:     suffix: 0
    640:     args:

    642:   test:
    643:     suffix: l1_1
    644:     args: -p 1 -tao_type lmvm -alpha 1. -epsilon 1.e-7 -m 64 -n 64 -view_sol -matrix_format 1

    646:   test:
    647:     suffix: hessian_1
    648:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nls

    650:   test:
    651:     suffix: hessian_2
    652:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nls

    654:   test:
    655:     suffix: nm_1
    656:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nm -tao_max_it 50

    658:   test:
    659:     suffix: nm_2
    660:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nm -tao_max_it 50

    662:   test:
    663:     suffix: lmvm_1
    664:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type lmvm -tao_max_it 40

    666:   test:
    667:     suffix: lmvm_2
    668:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type lmvm -tao_max_it 15

    670:   test:
    671:     suffix: soft_threshold_admm_1
    672:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm

    674:   test:
    675:     suffix: hessian_admm_1
    676:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nls -misfit_tao_type nls

    678:   test:
    679:     suffix: hessian_admm_2
    680:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nls -misfit_tao_type nls

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

    686:   test:
    687:     suffix: nm_admm_2
    688:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nm -misfit_tao_type nm -iter 7

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

    694:   test:
    695:     suffix: lmvm_admm_2
    696:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm

    698: TEST*/