Actual source code: ex4.c

petsc-3.14.0 2020-09-29
Report Typos and Errors
  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: {

 43:   /* build the rhs d in ctx */
 44:   VecCreate(PETSC_COMM_WORLD,&(ctx->d));
 45:   VecSetSizes(ctx->d,PETSC_DECIDE,ctx->m);
 46:   VecSetFromOptions(ctx->d);
 47:   VecSetRandom(ctx->d,ctx->rctx);
 48:   return(0);
 49: }

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

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

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

107: static PetscErrorCode SetupWorkspace(UserCtx ctx)
108: {
109:   PetscInt       i;

113:   MatCreateVecs(ctx->F, &ctx->workLeft[0], &ctx->workRight[0]);
114:   for (i=1; i<NWORKLEFT; i++) {
115:     VecDuplicate(ctx->workLeft[0], &(ctx->workLeft[i]));
116:   }
117:   for (i=1; i<NWORKRIGHT; i++) {
118:     VecDuplicate(ctx->workRight[0], &(ctx->workRight[i]));
119:   }
120:   return(0);
121: }

123: static PetscErrorCode ConfigureContext(UserCtx ctx)
124: {

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

169: static PetscErrorCode DestroyContext(UserCtx *ctx)
170: {
171:   PetscInt       i;

175:   MatDestroy(&((*ctx)->F));
176:   MatDestroy(&((*ctx)->W));
177:   MatDestroy(&((*ctx)->Hm));
178:   MatDestroy(&((*ctx)->Hr));
179:   VecDestroy(&((*ctx)->d));
180:   for (i=0; i<NWORKLEFT; i++) {
181:     VecDestroy(&((*ctx)->workLeft[i]));
182:   }
183:   for (i=0; i<NWORKRIGHT; i++) {
184:     VecDestroy(&((*ctx)->workRight[i]));
185:   }
186:   PetscRandomDestroy(&((*ctx)->rctx));
187:   PetscFree(*ctx);
188:   return(0);
189: }

191: /* compute (1/2) * ||F x - d||^2 */
192: static PetscErrorCode ObjectiveMisfit(Tao tao, Vec x, PetscReal *J, void *_ctx)
193: {
194:   UserCtx        ctx = (UserCtx) _ctx;
196:   Vec            y;

199:   y    = ctx->workLeft[0];
200:   MatMult(ctx->F, x, y);
201:   VecAXPY(y, -1., ctx->d);
202:   VecDot(y, y, J);
203:   *J  *= 0.5;
204:   return(0);
205: }

207: /* compute V = FTFx - FTd */
208: static PetscErrorCode GradientMisfit(Tao tao, Vec x, Vec V, void *_ctx)
209: {
210:   UserCtx        ctx = (UserCtx) _ctx;
212:   Vec            FTFx, FTd;

215:   /* work1 is A^T Ax, work2 is Ab, W is A^T A*/
216:   FTFx = ctx->workRight[0];
217:   FTd  = ctx->workRight[1];
218:   MatMult(ctx->W,x,FTFx);
219:   MatMultTranspose(ctx->F, ctx->d, FTd);
220:   VecWAXPY(V, -1., FTd, FTFx);
221:   return(0);
222: }

224: /* returns FTF */
225: static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
226: {
227:   UserCtx        ctx = (UserCtx) _ctx;

231:   if (H != ctx->W) {MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);}
232:   if (Hpre != ctx->W) {MatCopy(ctx->W, Hpre, DIFFERENT_NONZERO_PATTERN);}
233:   return(0);
234: }

236: /* computes augment Lagrangian objective (with scaled dual):
237:  * 0.5 * ||F x - d||^2  + 0.5 * mu ||x - z + u||^2 */
238: static PetscErrorCode ObjectiveMisfitADMM(Tao tao, Vec x, PetscReal *J, void *_ctx)
239: {
240:   UserCtx        ctx = (UserCtx) _ctx;
241:   PetscReal      mu, workNorm, misfit;
242:   Vec            z, u, temp;

246:   mu   = ctx->mu;
247:   z    = ctx->workRight[5];
248:   u    = ctx->workRight[6];
249:   temp = ctx->workRight[10];
250:   /* misfit = f(x) */
251:   ObjectiveMisfit(tao, x, &misfit, _ctx);
252:   VecCopy(x,temp);
253:   /* temp = x - z + u */
254:   VecAXPBYPCZ(temp,-1.,1.,1.,z,u);
255:   /* workNorm = ||x - z + u||^2 */
256:   VecDot(temp, temp, &workNorm);
257:   /* augment Lagrangian objective (with scaled dual): f(x) + 0.5 * mu ||x -z + u||^2 */
258:   *J = misfit + 0.5 * mu * workNorm;
259:   return(0);
260: }

262: /* computes FTFx - FTd  mu*(x - z + u) */
263: static PetscErrorCode GradientMisfitADMM(Tao tao, Vec x, Vec V, void *_ctx)
264: {
265:   UserCtx        ctx = (UserCtx) _ctx;
266:   PetscReal      mu;
267:   Vec            z, u, temp;

271:   mu   = ctx->mu;
272:   z    = ctx->workRight[5];
273:   u    = ctx->workRight[6];
274:   temp = ctx->workRight[10];
275:   GradientMisfit(tao, x, V, _ctx);
276:   VecCopy(x, temp);
277:   /* temp = x - z + u */
278:   VecAXPBYPCZ(temp,-1.,1.,1.,z,u);
279:   /* V =  FTFx - FTd  mu*(x - z + u) */
280:   VecAXPY(V, mu, temp);
281:   return(0);
282: }

284: /* returns FTF + diag(mu) */
285: static PetscErrorCode HessianMisfitADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
286: {
287:   UserCtx        ctx = (UserCtx) _ctx;

291:   MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);
292:   MatShift(H, ctx->mu);
293:   if (Hpre != H) {
294:     MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);
295:   }
296:   return(0);
297: }

299: /* computes || x ||_p (mult by 0.5 in case of NORM_2) */
300: static PetscErrorCode ObjectiveRegularization(Tao tao, Vec x, PetscReal *J, void *_ctx)
301: {
302:   UserCtx        ctx = (UserCtx) _ctx;
303:   PetscReal      norm;

307:   *J = 0;
308:   VecNorm (x, ctx->p, &norm);
309:   if (ctx->p == NORM_2) norm = 0.5 * norm * norm;
310:   *J = ctx->alpha * norm;
311:   return(0);
312: }

314: /* NORM_2 Case: return x
315:  * NORM_1 Case: x/(|x| + eps)
316:  * Else: TODO */
317: static PetscErrorCode GradientRegularization(Tao tao, Vec x, Vec V, void *_ctx)
318: {
319:   UserCtx        ctx = (UserCtx) _ctx;
321:   PetscReal      eps = ctx->eps;

324:   if (ctx->p == NORM_2) {
325:     VecCopy(x, V);
326:   } else if (ctx->p == NORM_1) {
327:     VecCopy(x, ctx->workRight[1]);
328:     VecAbs(ctx->workRight[1]);
329:     VecShift(ctx->workRight[1], eps);
330:     VecPointwiseDivide(V, x, ctx->workRight[1]);
331:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
332:   return(0);
333: }

335: /* NORM_2 Case: returns diag(mu)
336:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)))  */
337: static PetscErrorCode HessianRegularization(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
338: {
339:   UserCtx        ctx = (UserCtx) _ctx;
340:   PetscReal      eps = ctx->eps;
341:   Vec            copy1,copy2,copy3;

345:   if (ctx->p == NORM_2) {
346:     /* Identity matrix scaled by mu */
347:     MatZeroEntries(H);
348:     MatShift(H,ctx->mu);
349:     if (Hpre != H) {
350:       MatZeroEntries(Hpre);
351:       MatShift(Hpre,ctx->mu);
352:     }
353:   } else if (ctx->p == NORM_1) {
354:     /* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)) */
355:     copy1 = ctx->workRight[1];
356:     copy2 = ctx->workRight[2];
357:     copy3 = ctx->workRight[3];
358:     /* copy1 : 1/sqrt(x_i^2 + eps) */
359:     VecCopy(x, copy1);
360:     VecPow(copy1,2);
361:     VecShift(copy1, eps);
362:     VecSqrtAbs(copy1);
363:     VecReciprocal(copy1);
364:     /* copy2:  x_i^2.*/
365:     VecCopy(x,copy2);
366:     VecPow(copy2,2);
367:     /* copy3: abs(x_i^2 + eps) */
368:     VecCopy(x,copy3);
369:     VecPow(copy3,2);
370:     VecShift(copy3, eps);
371:     VecAbs(copy3);
372:     /* copy2: 1 - x_i^2/abs(x_i^2 + eps) */
373:     VecPointwiseDivide(copy2, copy2,copy3);
374:     VecScale(copy2, -1.);
375:     VecShift(copy2, 1.);
376:     VecAXPY(copy1,1.,copy2);
377:     VecScale(copy1, ctx->mu);
378:     MatZeroEntries(H);
379:     MatDiagonalSet(H, copy1,INSERT_VALUES);
380:     if (Hpre != H) {
381:       MatZeroEntries(Hpre);
382:       MatDiagonalSet(Hpre, copy1,INSERT_VALUES);
383:     }
384:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
385:   return(0);
386: }

388: /* NORM_2 Case: 0.5 || x ||_2 + 0.5 * mu * ||x + u - z||^2
389:  * Else : || x ||_2 + 0.5 * mu * ||x + u - z||^2 */
390: static PetscErrorCode ObjectiveRegularizationADMM(Tao tao, Vec z, PetscReal *J, void *_ctx)
391: {
392:   UserCtx        ctx = (UserCtx) _ctx;
393:   PetscReal      mu, workNorm, reg;
394:   Vec            x, u, temp;

398:   mu   = ctx->mu;
399:   x    = ctx->workRight[4];
400:   u    = ctx->workRight[6];
401:   temp = ctx->workRight[10];
402:   ObjectiveRegularization(tao, z, &reg, _ctx);
403:   VecCopy(z,temp);
404:   /* temp = x + u -z */
405:   VecAXPBYPCZ(temp,1.,1.,-1.,x,u);
406:   /* workNorm = ||x + u - z ||^2 */
407:   VecDot(temp, temp, &workNorm);
408:   *J   = reg + 0.5 * mu * workNorm;
409:   return(0);
410: }


413: /* NORM_2 Case: x - mu*(x + u - z)
414:  * NORM_1 Case: x/(|x| + eps) - mu*(x + u - z)
415:  * Else: TODO */
416: static PetscErrorCode GradientRegularizationADMM(Tao tao, Vec z, Vec V, void *_ctx)
417: {
418:   UserCtx        ctx = (UserCtx) _ctx;
419:   PetscReal      mu;
420:   Vec            x, u, temp;

424:   mu   = ctx->mu;
425:   x    = ctx->workRight[4];
426:   u    = ctx->workRight[6];
427:   temp = ctx->workRight[10];
428:   GradientRegularization(tao, z, V, _ctx);
429:   VecCopy(z, temp);
430:   /* temp = x + u -z */
431:   VecAXPBYPCZ(temp,1.,1.,-1.,x,u);
432:   VecAXPY(V, -mu, temp);
433:   return(0);
434: }

436: /* NORM_2 Case: returns diag(mu)
437:  * NORM_1 Case: FTF + diag(mu) */
438: static PetscErrorCode HessianRegularizationADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
439: {
440:   UserCtx        ctx = (UserCtx) _ctx;

444:   if (ctx->p == NORM_2) {
445:     /* Identity matrix scaled by mu */
446:     MatZeroEntries(H);
447:     MatShift(H,ctx->mu);
448:     if (Hpre != H) {
449:       MatZeroEntries(Hpre);
450:       MatShift(Hpre,ctx->mu);
451:     }
452:   } else if (ctx->p == NORM_1) {
453:     HessianMisfit(tao, x, H, Hpre, (void*) ctx);
454:     MatShift(H, ctx->mu);
455:     if (Hpre != H) {MatShift(Hpre, ctx->mu);}
456:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
457:   return(0);
458: }

460: /* NORM_2 Case : (1/2) * ||F x - d||^2 + 0.5 * || x ||_p
461: *  NORM_1 Case : (1/2) * ||F x - d||^2 + || x ||_p */
462: static PetscErrorCode ObjectiveComplete(Tao tao, Vec x, PetscReal *J, void *ctx)
463: {
464:   PetscReal      Jm, Jr;

468:   ObjectiveMisfit(tao, x, &Jm, ctx);
469:   ObjectiveRegularization(tao, x, &Jr, ctx);
470:   *J   = Jm + Jr;
471:   return(0);
472: }

474: /* NORM_2 Case: FTFx - FTd + x
475:  * NORM_1 Case: FTFx - FTd + x/(|x| + eps) */
476: static PetscErrorCode GradientComplete(Tao tao, Vec x, Vec V, void *ctx)
477: {
478:   UserCtx        cntx = (UserCtx) ctx;

482:   GradientMisfit(tao, x, cntx->workRight[2], ctx);
483:   GradientRegularization(tao, x, cntx->workRight[3], ctx);
484:   VecWAXPY(V,1,cntx->workRight[2],cntx->workRight[3]);
485:   return(0);
486: }

488: /* NORM_2 Case: diag(mu) + FTF
489:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) + FTF  */
490: static PetscErrorCode HessianComplete(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx)
491: {
492:   Mat            tempH;

496:   MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH);
497:   HessianMisfit(tao, x, H, H, ctx);
498:   HessianRegularization(tao, x, tempH, tempH, ctx);
499:   MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN);
500:   if (Hpre != H) {
501:     MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);
502:   }
503:   MatDestroy(&tempH);
504:   return(0);
505: }

507: static PetscErrorCode TaoSolveADMM(UserCtx ctx,  Vec x)
508: {
510:   PetscInt       i;
511:   PetscReal      u_norm, r_norm, s_norm, primal, dual, x_norm, z_norm;
512:   Tao            tao1,tao2;
513:   Vec            xk,z,u,diff,zold,zdiff,temp;
514:   PetscReal      mu;

517:   xk    = ctx->workRight[4];
518:   z     = ctx->workRight[5];
519:   u     = ctx->workRight[6];
520:   diff  = ctx->workRight[7];
521:   zold  = ctx->workRight[8];
522:   zdiff = ctx->workRight[9];
523:   temp  = ctx->workRight[11];
524:   mu    = ctx->mu;
525:   VecSet(u, 0.);
526:   TaoCreate(PETSC_COMM_WORLD, &tao1);
527:   TaoSetType(tao1,TAONLS);
528:   TaoSetObjectiveRoutine(tao1, ObjectiveMisfitADMM, (void*) ctx);
529:   TaoSetGradientRoutine(tao1, GradientMisfitADMM, (void*) ctx);
530:   TaoSetHessianRoutine(tao1, ctx->Hm, ctx->Hm, HessianMisfitADMM, (void*) ctx);
531:   VecSet(xk, 0.);
532:   TaoSetInitialVector(tao1, xk);
533:   TaoSetOptionsPrefix(tao1, "misfit_");
534:   TaoSetFromOptions(tao1);
535:   TaoCreate(PETSC_COMM_WORLD, &tao2);
536:   if (ctx->p == NORM_2) {
537:     TaoSetType(tao2,TAONLS);
538:     TaoSetObjectiveRoutine(tao2, ObjectiveRegularizationADMM, (void*) ctx);
539:     TaoSetGradientRoutine(tao2, GradientRegularizationADMM, (void*) ctx);
540:     TaoSetHessianRoutine(tao2, ctx->Hr, ctx->Hr, HessianRegularizationADMM, (void*) ctx);
541:   }
542:   VecSet(z, 0.);
543:   TaoSetInitialVector(tao2, z);
544:   TaoSetOptionsPrefix(tao2, "reg_");
545:   TaoSetFromOptions(tao2);

547:   for (i=0; i<ctx->iter; i++) {
548:     VecCopy(z,zold);
549:     TaoSolve(tao1); /* Updates xk */
550:     if (ctx->p == NORM_1){
551:       VecWAXPY(temp,1.,xk,u);
552:       TaoSoftThreshold(temp,-ctx->alpha/mu,ctx->alpha/mu,z);
553:     } else {
554:       TaoSolve(tao2); /* Update zk */
555:     }
556:     /* u = u + xk -z */
557:     VecAXPBYPCZ(u,1.,-1.,1.,xk,z);
558:     /* r_norm : norm(x-z) */
559:     VecWAXPY(diff,-1.,z,xk);
560:     VecNorm(diff,NORM_2,&r_norm);
561:     /* s_norm : norm(-mu(z-zold)) */
562:     VecWAXPY(zdiff, -1.,zold,z);
563:     VecNorm(zdiff,NORM_2,&s_norm);
564:     s_norm = s_norm * mu;
565:     /* primal : sqrt(n)*ABSTOL + RELTOL*max(norm(x), norm(-z))*/
566:     VecNorm(xk,NORM_2,&x_norm);
567:     VecNorm(z,NORM_2,&z_norm);
568:     primal = PetscSqrtReal(ctx->n)*ctx->abstol + ctx->reltol*PetscMax(x_norm,z_norm);
569:     /* Duality : sqrt(n)*ABSTOL + RELTOL*norm(mu*u)*/
570:     VecNorm(u,NORM_2,&u_norm);
571:     dual = PetscSqrtReal(ctx->n)*ctx->abstol + ctx->reltol*u_norm*mu;
572:     PetscPrintf(PetscObjectComm((PetscObject)tao1),"Iter %D : ||x-z||: %g, mu*||z-zold||: %g\n", i, (double) r_norm, (double) s_norm);
573:     if (r_norm < primal && s_norm < dual) break;
574:   }
575:   VecCopy(xk, x);
576:   TaoDestroy(&tao1);
577:   TaoDestroy(&tao2);
578:   return(0);
579: }

581: /* Second order Taylor remainder convergence test */
582: static PetscErrorCode TaylorTest(UserCtx ctx, Tao tao, Vec x, PetscReal *C)
583: {
584:   PetscReal      h,J,temp;
585:   PetscInt       i,j;
586:   PetscInt       numValues;
587:   PetscReal      Jx,Jxhat_comp,Jxhat_pred;
588:   PetscReal      *Js, *hs;
589:   PetscReal      gdotdx;
590:   PetscReal      minrate = PETSC_MAX_REAL;
591:   MPI_Comm       comm = PetscObjectComm((PetscObject)x);
592:   Vec            g, dx, xhat;

596:   VecDuplicate(x, &g);
597:   VecDuplicate(x, &xhat);
598:   /* choose a perturbation direction */
599:   VecDuplicate(x, &dx);
600:   VecSetRandom(dx,ctx->rctx);
601:   /* evaluate objective at x: J(x) */
602:   TaoComputeObjective(tao, x, &Jx);
603:   /* evaluate gradient at x, save in vector g */
604:   TaoComputeGradient(tao, x, g);
605:   VecDot(g, dx, &gdotdx);

607:   for (numValues=0, h=ctx->hStart; h>=ctx->hMin; h*=ctx->hFactor) numValues++;
608:   PetscCalloc2(numValues, &Js, numValues, &hs);
609:   for (i=0, h=ctx->hStart; h>=ctx->hMin; h*=ctx->hFactor, i++) {
610:     VecWAXPY(xhat, h, dx, x);
611:     TaoComputeObjective(tao, xhat, &Jxhat_comp);
612:     /* J(\hat(x)) \approx J(x) + g^T (xhat - x) = J(x) + h * g^T dx */
613:     Jxhat_pred = Jx + h * gdotdx;
614:     /* Vector to dJdm scalar? Dot?*/
615:     J     = PetscAbsReal(Jxhat_comp - Jxhat_pred);
616:     PetscPrintf (comm, "J(xhat): %g, predicted: %g, diff %g\n", (double) Jxhat_comp,(double) Jxhat_pred, (double) J);
617:     Js[i] = J;
618:     hs[i] = h;
619:   }
620:   for (j=1; j<numValues; j++) {
621:     temp    = PetscLogReal(Js[j]/Js[j-1]) / PetscLogReal (hs[j]/hs[j-1]);
622:     PetscPrintf (comm, "Convergence rate step %D: %g\n", j-1, (double) temp);
623:     minrate = PetscMin(minrate, temp);
624:   }
625:   /* If O is not ~2, then the test is wrong */
626:   PetscFree2(Js, hs);
627:   *C   = minrate;
628:   VecDestroy(&dx);
629:   VecDestroy(&xhat);
630:   VecDestroy(&g);
631:   return(0);
632: }

634: int main(int argc, char ** argv)
635: {
636:   UserCtx        ctx;
637:   Tao            tao;
638:   Vec            x;
639:   Mat            H;

642:   PetscInitialize(&argc, &argv, NULL,help);if (ierr) return ierr;
643:   PetscNew(&ctx);
644:   ConfigureContext(ctx);
645:   /* Define two functions that could pass as objectives to TaoSetObjectiveRoutine(): one
646:    * for the misfit component, and one for the regularization component */
647:   /* ObjectiveMisfit() and ObjectiveRegularization() */

649:   /* Define a single function that calls both components adds them together: the complete objective,
650:    * in the absence of a Tao implementation that handles separability */
651:   /* ObjectiveComplete() */
652:   TaoCreate(PETSC_COMM_WORLD, &tao);
653:   TaoSetType(tao,TAONM);
654:   TaoSetObjectiveRoutine(tao, ObjectiveComplete, (void*) ctx);
655:   TaoSetGradientRoutine(tao, GradientComplete, (void*) ctx);
656:   MatDuplicate(ctx->W, MAT_SHARE_NONZERO_PATTERN, &H);
657:   TaoSetHessianRoutine(tao, H, H, HessianComplete, (void*) ctx);
658:   MatCreateVecs(ctx->F, NULL, &x);
659:   VecSet(x, 0.);
660:   TaoSetInitialVector(tao, x);
661:   TaoSetFromOptions(tao);
662:   if (ctx->use_admm) {
663:     TaoSolveADMM(ctx,x);
664:   } else {TaoSolve(tao);}
665:   /* examine solution */
666:   VecViewFromOptions(x, NULL, "-view_sol");
667:   if (ctx->taylor) {
668:     PetscReal rate;
669:     TaylorTest(ctx, tao, x, &rate);
670:   }
671:   MatDestroy(&H);
672:   TaoDestroy(&tao);
673:   VecDestroy(&x);
674:   DestroyContext(&ctx);
675:   PetscFinalize();
676:   return ierr;
677: }

679: /*TEST

681:   build:
682:     requires: !complex

684:   test:
685:     suffix: 0
686:     args:

688:   test:
689:     suffix: l1_1
690:     args: -p 1 -tao_type lmvm -alpha 1. -epsilon 1.e-7 -m 64 -n 64 -view_sol -matrix_format 1

692:   test:
693:     suffix: hessian_1
694:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nls -tao_nls_ksp_monitor

696:   test:
697:     suffix: hessian_2
698:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nls -tao_nls_ksp_monitor

700:   test:
701:     suffix: nm_1
702:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nm -tao_max_it 50

704:   test:
705:     suffix: nm_2
706:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nm -tao_max_it 50

708:   test:
709:     suffix: lmvm_1
710:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type lmvm -tao_max_it 40

712:   test:
713:     suffix: lmvm_2
714:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type lmvm -tao_max_it 15

716:   test:
717:     suffix: soft_threshold_admm_1
718:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm

720:   test:
721:     suffix: hessian_admm_1
722:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nls -misfit_tao_type nls

724:   test:
725:     suffix: hessian_admm_2
726:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nls -misfit_tao_type nls

728:   test:
729:     suffix: nm_admm_1
730:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nm -misfit_tao_type nm

732:   test:
733:     suffix: nm_admm_2
734:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nm -misfit_tao_type nm -iter 7

736:   test:
737:     suffix: lmvm_admm_1
738:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm

740:   test:
741:     suffix: lmvm_admm_2
742:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm

744: TEST*/