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;    /* Numer 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   taylor;   /* Flag to determine whether to run Taylor test or not */
 34:   PetscBool   use_admm; /* Flag to determine whether to run Taylor test or not */
 35: } *UserCtx;

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

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

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

 65:   /* Set matrix elements in  2-D five point stencil format. */
 66:   if (!(ctx->matops)) {
 68:     gridN = (PetscInt)PetscSqrtReal((PetscReal)ctx->m);
 70:     for (Ii = Istart; Ii < Iend; Ii++) {
 71:       i   = Ii / gridN;
 72:       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)) MatSetOption(ctx->F, MAT_SYMMETRIC, PETSC_TRUE);
 93:   MatTransposeMatMult(ctx->F, ctx->F, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &(ctx->W));
 94:   /* Setup Hessian Workspace in same shape as W */
 95:   MatDuplicate(ctx->W, MAT_DO_NOT_COPY_VALUES, &(ctx->Hm));
 96:   MatDuplicate(ctx->W, MAT_DO_NOT_COPY_VALUES, &(ctx->Hr));
 97:   return 0;
 98: }

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

104:   MatCreateVecs(ctx->F, &ctx->workLeft[0], &ctx->workRight[0]);
105:   for (i = 1; i < NWORKLEFT; i++) VecDuplicate(ctx->workLeft[0], &(ctx->workLeft[i]));
106:   for (i = 1; i < NWORKRIGHT; i++) VecDuplicate(ctx->workRight[0], &(ctx->workRight[i]));
107:   return 0;
108: }

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

153: static PetscErrorCode DestroyContext(UserCtx *ctx)
154: {
155:   PetscInt i;

157:   MatDestroy(&((*ctx)->F));
158:   MatDestroy(&((*ctx)->W));
159:   MatDestroy(&((*ctx)->Hm));
160:   MatDestroy(&((*ctx)->Hr));
161:   VecDestroy(&((*ctx)->d));
162:   for (i = 0; i < NWORKLEFT; i++) VecDestroy(&((*ctx)->workLeft[i]));
163:   for (i = 0; i < NWORKRIGHT; i++) VecDestroy(&((*ctx)->workRight[i]));
164:   PetscRandomDestroy(&((*ctx)->rctx));
165:   PetscFree(*ctx);
166:   return 0;
167: }

169: /* compute (1/2) * ||F x - d||^2 */
170: static PetscErrorCode ObjectiveMisfit(Tao tao, Vec x, PetscReal *J, void *_ctx)
171: {
172:   UserCtx ctx = (UserCtx)_ctx;
173:   Vec     y;

175:   y = ctx->workLeft[0];
176:   MatMult(ctx->F, x, y);
177:   VecAXPY(y, -1., ctx->d);
178:   VecDot(y, y, J);
179:   *J *= 0.5;
180:   return 0;
181: }

183: /* compute V = FTFx - FTd */
184: static PetscErrorCode GradientMisfit(Tao tao, Vec x, Vec V, void *_ctx)
185: {
186:   UserCtx ctx = (UserCtx)_ctx;
187:   Vec     FTFx, FTd;

189:   /* work1 is A^T Ax, work2 is Ab, W is A^T A*/
190:   FTFx = ctx->workRight[0];
191:   FTd  = ctx->workRight[1];
192:   MatMult(ctx->W, x, FTFx);
193:   MatMultTranspose(ctx->F, ctx->d, FTd);
194:   VecWAXPY(V, -1., FTd, FTFx);
195:   return 0;
196: }

198: /* returns FTF */
199: static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
200: {
201:   UserCtx ctx = (UserCtx)_ctx;

203:   if (H != ctx->W) MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);
204:   if (Hpre != ctx->W) MatCopy(ctx->W, Hpre, DIFFERENT_NONZERO_PATTERN);
205:   return 0;
206: }

208: /* computes augment Lagrangian objective (with scaled dual):
209:  * 0.5 * ||F x - d||^2  + 0.5 * mu ||x - z + u||^2 */
210: static PetscErrorCode ObjectiveMisfitADMM(Tao tao, Vec x, PetscReal *J, void *_ctx)
211: {
212:   UserCtx   ctx = (UserCtx)_ctx;
213:   PetscReal mu, workNorm, misfit;
214:   Vec       z, u, temp;

216:   mu   = ctx->mu;
217:   z    = ctx->workRight[5];
218:   u    = ctx->workRight[6];
219:   temp = ctx->workRight[10];
220:   /* misfit = f(x) */
221:   ObjectiveMisfit(tao, x, &misfit, _ctx);
222:   VecCopy(x, temp);
223:   /* temp = x - z + u */
224:   VecAXPBYPCZ(temp, -1., 1., 1., z, u);
225:   /* workNorm = ||x - z + u||^2 */
226:   VecDot(temp, temp, &workNorm);
227:   /* augment Lagrangian objective (with scaled dual): f(x) + 0.5 * mu ||x -z + u||^2 */
228:   *J = misfit + 0.5 * mu * workNorm;
229:   return 0;
230: }

232: /* computes FTFx - FTd  mu*(x - z + u) */
233: static PetscErrorCode GradientMisfitADMM(Tao tao, Vec x, Vec V, void *_ctx)
234: {
235:   UserCtx   ctx = (UserCtx)_ctx;
236:   PetscReal mu;
237:   Vec       z, u, temp;

239:   mu   = ctx->mu;
240:   z    = ctx->workRight[5];
241:   u    = ctx->workRight[6];
242:   temp = ctx->workRight[10];
243:   GradientMisfit(tao, x, V, _ctx);
244:   VecCopy(x, temp);
245:   /* temp = x - z + u */
246:   VecAXPBYPCZ(temp, -1., 1., 1., z, u);
247:   /* V =  FTFx - FTd  mu*(x - z + u) */
248:   VecAXPY(V, mu, temp);
249:   return 0;
250: }

252: /* returns FTF + diag(mu) */
253: static PetscErrorCode HessianMisfitADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
254: {
255:   UserCtx ctx = (UserCtx)_ctx;

257:   MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);
258:   MatShift(H, ctx->mu);
259:   if (Hpre != H) MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);
260:   return 0;
261: }

263: /* computes || x ||_p (mult by 0.5 in case of NORM_2) */
264: static PetscErrorCode ObjectiveRegularization(Tao tao, Vec x, PetscReal *J, void *_ctx)
265: {
266:   UserCtx   ctx = (UserCtx)_ctx;
267:   PetscReal norm;

269:   *J = 0;
270:   VecNorm(x, ctx->p, &norm);
271:   if (ctx->p == NORM_2) norm = 0.5 * norm * norm;
272:   *J = ctx->alpha * norm;
273:   return 0;
274: }

276: /* NORM_2 Case: return x
277:  * NORM_1 Case: x/(|x| + eps)
278:  * Else: TODO */
279: static PetscErrorCode GradientRegularization(Tao tao, Vec x, Vec V, void *_ctx)
280: {
281:   UserCtx   ctx = (UserCtx)_ctx;
282:   PetscReal eps = ctx->eps;

284:   if (ctx->p == NORM_2) {
285:     VecCopy(x, V);
286:   } else if (ctx->p == NORM_1) {
287:     VecCopy(x, ctx->workRight[1]);
288:     VecAbs(ctx->workRight[1]);
289:     VecShift(ctx->workRight[1], eps);
290:     VecPointwiseDivide(V, x, ctx->workRight[1]);
291:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
292:   return 0;
293: }

295: /* NORM_2 Case: returns diag(mu)
296:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)))  */
297: static PetscErrorCode HessianRegularization(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
298: {
299:   UserCtx   ctx = (UserCtx)_ctx;
300:   PetscReal eps = ctx->eps;
301:   Vec       copy1, copy2, copy3;

303:   if (ctx->p == NORM_2) {
304:     /* Identity matrix scaled by mu */
305:     MatZeroEntries(H);
306:     MatShift(H, ctx->mu);
307:     if (Hpre != H) {
308:       MatZeroEntries(Hpre);
309:       MatShift(Hpre, ctx->mu);
310:     }
311:   } else if (ctx->p == NORM_1) {
312:     /* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)) */
313:     copy1 = ctx->workRight[1];
314:     copy2 = ctx->workRight[2];
315:     copy3 = ctx->workRight[3];
316:     /* copy1 : 1/sqrt(x_i^2 + eps) */
317:     VecCopy(x, copy1);
318:     VecPow(copy1, 2);
319:     VecShift(copy1, eps);
320:     VecSqrtAbs(copy1);
321:     VecReciprocal(copy1);
322:     /* copy2:  x_i^2.*/
323:     VecCopy(x, copy2);
324:     VecPow(copy2, 2);
325:     /* copy3: abs(x_i^2 + eps) */
326:     VecCopy(x, copy3);
327:     VecPow(copy3, 2);
328:     VecShift(copy3, eps);
329:     VecAbs(copy3);
330:     /* copy2: 1 - x_i^2/abs(x_i^2 + eps) */
331:     VecPointwiseDivide(copy2, copy2, copy3);
332:     VecScale(copy2, -1.);
333:     VecShift(copy2, 1.);
334:     VecAXPY(copy1, 1., copy2);
335:     VecScale(copy1, ctx->mu);
336:     MatZeroEntries(H);
337:     MatDiagonalSet(H, copy1, INSERT_VALUES);
338:     if (Hpre != H) {
339:       MatZeroEntries(Hpre);
340:       MatDiagonalSet(Hpre, copy1, INSERT_VALUES);
341:     }
342:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
343:   return 0;
344: }

346: /* NORM_2 Case: 0.5 || x ||_2 + 0.5 * mu * ||x + u - z||^2
347:  * Else : || x ||_2 + 0.5 * mu * ||x + u - z||^2 */
348: static PetscErrorCode ObjectiveRegularizationADMM(Tao tao, Vec z, PetscReal *J, void *_ctx)
349: {
350:   UserCtx   ctx = (UserCtx)_ctx;
351:   PetscReal mu, workNorm, reg;
352:   Vec       x, u, temp;

354:   mu   = ctx->mu;
355:   x    = ctx->workRight[4];
356:   u    = ctx->workRight[6];
357:   temp = ctx->workRight[10];
358:   ObjectiveRegularization(tao, z, &reg, _ctx);
359:   VecCopy(z, temp);
360:   /* temp = x + u -z */
361:   VecAXPBYPCZ(temp, 1., 1., -1., x, u);
362:   /* workNorm = ||x + u - z ||^2 */
363:   VecDot(temp, temp, &workNorm);
364:   *J = reg + 0.5 * mu * workNorm;
365:   return 0;
366: }

368: /* NORM_2 Case: x - mu*(x + u - z)
369:  * NORM_1 Case: x/(|x| + eps) - mu*(x + u - z)
370:  * Else: TODO */
371: static PetscErrorCode GradientRegularizationADMM(Tao tao, Vec z, Vec V, void *_ctx)
372: {
373:   UserCtx   ctx = (UserCtx)_ctx;
374:   PetscReal mu;
375:   Vec       x, u, temp;

377:   mu   = ctx->mu;
378:   x    = ctx->workRight[4];
379:   u    = ctx->workRight[6];
380:   temp = ctx->workRight[10];
381:   GradientRegularization(tao, z, V, _ctx);
382:   VecCopy(z, temp);
383:   /* temp = x + u -z */
384:   VecAXPBYPCZ(temp, 1., 1., -1., x, u);
385:   VecAXPY(V, -mu, temp);
386:   return 0;
387: }

389: /* NORM_2 Case: returns diag(mu)
390:  * NORM_1 Case: FTF + diag(mu) */
391: static PetscErrorCode HessianRegularizationADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
392: {
393:   UserCtx ctx = (UserCtx)_ctx;

395:   if (ctx->p == NORM_2) {
396:     /* Identity matrix scaled by mu */
397:     MatZeroEntries(H);
398:     MatShift(H, ctx->mu);
399:     if (Hpre != H) {
400:       MatZeroEntries(Hpre);
401:       MatShift(Hpre, ctx->mu);
402:     }
403:   } else if (ctx->p == NORM_1) {
404:     HessianMisfit(tao, x, H, Hpre, (void *)ctx);
405:     MatShift(H, ctx->mu);
406:     if (Hpre != H) MatShift(Hpre, ctx->mu);
407:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
408:   return 0;
409: }

411: /* NORM_2 Case : (1/2) * ||F x - d||^2 + 0.5 * || x ||_p
412: *  NORM_1 Case : (1/2) * ||F x - d||^2 + || x ||_p */
413: static PetscErrorCode ObjectiveComplete(Tao tao, Vec x, PetscReal *J, void *ctx)
414: {
415:   PetscReal Jm, Jr;

417:   ObjectiveMisfit(tao, x, &Jm, ctx);
418:   ObjectiveRegularization(tao, x, &Jr, ctx);
419:   *J = Jm + Jr;
420:   return 0;
421: }

423: /* NORM_2 Case: FTFx - FTd + x
424:  * NORM_1 Case: FTFx - FTd + x/(|x| + eps) */
425: static PetscErrorCode GradientComplete(Tao tao, Vec x, Vec V, void *ctx)
426: {
427:   UserCtx cntx = (UserCtx)ctx;

429:   GradientMisfit(tao, x, cntx->workRight[2], ctx);
430:   GradientRegularization(tao, x, cntx->workRight[3], ctx);
431:   VecWAXPY(V, 1, cntx->workRight[2], cntx->workRight[3]);
432:   return 0;
433: }

435: /* NORM_2 Case: diag(mu) + FTF
436:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) + FTF  */
437: static PetscErrorCode HessianComplete(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx)
438: {
439:   Mat tempH;

441:   MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH);
442:   HessianMisfit(tao, x, H, H, ctx);
443:   HessianRegularization(tao, x, tempH, tempH, ctx);
444:   MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN);
445:   if (Hpre != H) MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);
446:   MatDestroy(&tempH);
447:   return 0;
448: }

450: static PetscErrorCode TaoSolveADMM(UserCtx ctx, Vec x)
451: {
452:   PetscInt  i;
453:   PetscReal u_norm, r_norm, s_norm, primal, dual, x_norm, z_norm;
454:   Tao       tao1, tao2;
455:   Vec       xk, z, u, diff, zold, zdiff, temp;
456:   PetscReal mu;

458:   xk    = ctx->workRight[4];
459:   z     = ctx->workRight[5];
460:   u     = ctx->workRight[6];
461:   diff  = ctx->workRight[7];
462:   zold  = ctx->workRight[8];
463:   zdiff = ctx->workRight[9];
464:   temp  = ctx->workRight[11];
465:   mu    = ctx->mu;
466:   VecSet(u, 0.);
467:   TaoCreate(PETSC_COMM_WORLD, &tao1);
468:   TaoSetType(tao1, TAONLS);
469:   TaoSetObjective(tao1, ObjectiveMisfitADMM, (void *)ctx);
470:   TaoSetGradient(tao1, NULL, GradientMisfitADMM, (void *)ctx);
471:   TaoSetHessian(tao1, ctx->Hm, ctx->Hm, HessianMisfitADMM, (void *)ctx);
472:   VecSet(xk, 0.);
473:   TaoSetSolution(tao1, xk);
474:   TaoSetOptionsPrefix(tao1, "misfit_");
475:   TaoSetFromOptions(tao1);
476:   TaoCreate(PETSC_COMM_WORLD, &tao2);
477:   if (ctx->p == NORM_2) {
478:     TaoSetType(tao2, TAONLS);
479:     TaoSetObjective(tao2, ObjectiveRegularizationADMM, (void *)ctx);
480:     TaoSetGradient(tao2, NULL, GradientRegularizationADMM, (void *)ctx);
481:     TaoSetHessian(tao2, ctx->Hr, ctx->Hr, HessianRegularizationADMM, (void *)ctx);
482:   }
483:   VecSet(z, 0.);
484:   TaoSetSolution(tao2, z);
485:   TaoSetOptionsPrefix(tao2, "reg_");
486:   TaoSetFromOptions(tao2);

488:   for (i = 0; i < ctx->iter; i++) {
489:     VecCopy(z, zold);
490:     TaoSolve(tao1); /* Updates xk */
491:     if (ctx->p == NORM_1) {
492:       VecWAXPY(temp, 1., xk, u);
493:       TaoSoftThreshold(temp, -ctx->alpha / mu, ctx->alpha / mu, z);
494:     } else {
495:       TaoSolve(tao2); /* Update zk */
496:     }
497:     /* u = u + xk -z */
498:     VecAXPBYPCZ(u, 1., -1., 1., xk, z);
499:     /* r_norm : norm(x-z) */
500:     VecWAXPY(diff, -1., z, xk);
501:     VecNorm(diff, NORM_2, &r_norm);
502:     /* s_norm : norm(-mu(z-zold)) */
503:     VecWAXPY(zdiff, -1., zold, z);
504:     VecNorm(zdiff, NORM_2, &s_norm);
505:     s_norm = s_norm * mu;
506:     /* primal : sqrt(n)*ABSTOL + RELTOL*max(norm(x), norm(-z))*/
507:     VecNorm(xk, NORM_2, &x_norm);
508:     VecNorm(z, NORM_2, &z_norm);
509:     primal = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * PetscMax(x_norm, z_norm);
510:     /* Duality : sqrt(n)*ABSTOL + RELTOL*norm(mu*u)*/
511:     VecNorm(u, NORM_2, &u_norm);
512:     dual = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * u_norm * mu;
513:     PetscPrintf(PetscObjectComm((PetscObject)tao1), "Iter %" PetscInt_FMT " : ||x-z||: %g, mu*||z-zold||: %g\n", i, (double)r_norm, (double)s_norm);
514:     if (r_norm < primal && s_norm < dual) break;
515:   }
516:   VecCopy(xk, x);
517:   TaoDestroy(&tao1);
518:   TaoDestroy(&tao2);
519:   return 0;
520: }

522: /* Second order Taylor remainder convergence test */
523: static PetscErrorCode TaylorTest(UserCtx ctx, Tao tao, Vec x, PetscReal *C)
524: {
525:   PetscReal  h, J, temp;
526:   PetscInt   i, j;
527:   PetscInt   numValues;
528:   PetscReal  Jx, Jxhat_comp, Jxhat_pred;
529:   PetscReal *Js, *hs;
530:   PetscReal  gdotdx;
531:   PetscReal  minrate = PETSC_MAX_REAL;
532:   MPI_Comm   comm    = PetscObjectComm((PetscObject)x);
533:   Vec        g, dx, xhat;

535:   VecDuplicate(x, &g);
536:   VecDuplicate(x, &xhat);
537:   /* choose a perturbation direction */
538:   VecDuplicate(x, &dx);
539:   VecSetRandom(dx, ctx->rctx);
540:   /* evaluate objective at x: J(x) */
541:   TaoComputeObjective(tao, x, &Jx);
542:   /* evaluate gradient at x, save in vector g */
543:   TaoComputeGradient(tao, x, g);
544:   VecDot(g, dx, &gdotdx);

546:   for (numValues = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor) numValues++;
547:   PetscCalloc2(numValues, &Js, numValues, &hs);
548:   for (i = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor, i++) {
549:     VecWAXPY(xhat, h, dx, x);
550:     TaoComputeObjective(tao, xhat, &Jxhat_comp);
551:     /* J(\hat(x)) \approx J(x) + g^T (xhat - x) = J(x) + h * g^T dx */
552:     Jxhat_pred = Jx + h * gdotdx;
553:     /* Vector to dJdm scalar? Dot?*/
554:     J = PetscAbsReal(Jxhat_comp - Jxhat_pred);
555:     PetscPrintf(comm, "J(xhat): %g, predicted: %g, diff %g\n", (double)Jxhat_comp, (double)Jxhat_pred, (double)J);
556:     Js[i] = J;
557:     hs[i] = h;
558:   }
559:   for (j = 1; j < numValues; j++) {
560:     temp = PetscLogReal(Js[j] / Js[j - 1]) / PetscLogReal(hs[j] / hs[j - 1]);
561:     PetscPrintf(comm, "Convergence rate step %" PetscInt_FMT ": %g\n", j - 1, (double)temp);
562:     minrate = PetscMin(minrate, temp);
563:   }
564:   /* If O is not ~2, then the test is wrong */
565:   PetscFree2(Js, hs);
566:   *C = minrate;
567:   VecDestroy(&dx);
568:   VecDestroy(&xhat);
569:   VecDestroy(&g);
570:   return 0;
571: }

573: int main(int argc, char **argv)
574: {
575:   UserCtx ctx;
576:   Tao     tao;
577:   Vec     x;
578:   Mat     H;

581:   PetscInitialize(&argc, &argv, NULL, help);
582:   PetscNew(&ctx);
583:   ConfigureContext(ctx);
584:   /* Define two functions that could pass as objectives to TaoSetObjective(): one
585:    * for the misfit component, and one for the regularization component */
586:   /* ObjectiveMisfit() and ObjectiveRegularization() */

588:   /* Define a single function that calls both components adds them together: the complete objective,
589:    * in the absence of a Tao implementation that handles separability */
590:   /* ObjectiveComplete() */
591:   TaoCreate(PETSC_COMM_WORLD, &tao);
592:   TaoSetType(tao, TAONM);
593:   TaoSetObjective(tao, ObjectiveComplete, (void *)ctx);
594:   TaoSetGradient(tao, NULL, GradientComplete, (void *)ctx);
595:   MatDuplicate(ctx->W, MAT_SHARE_NONZERO_PATTERN, &H);
596:   TaoSetHessian(tao, H, H, HessianComplete, (void *)ctx);
597:   MatCreateVecs(ctx->F, NULL, &x);
598:   VecSet(x, 0.);
599:   TaoSetSolution(tao, x);
600:   TaoSetFromOptions(tao);
601:   if (ctx->use_admm) TaoSolveADMM(ctx, x);
602:   else TaoSolve(tao);
603:   /* examine solution */
604:   VecViewFromOptions(x, NULL, "-view_sol");
605:   if (ctx->taylor) {
606:     PetscReal rate;
607:     TaylorTest(ctx, tao, x, &rate);
608:   }
609:   MatDestroy(&H);
610:   TaoDestroy(&tao);
611:   VecDestroy(&x);
612:   DestroyContext(&ctx);
613:   PetscFinalize();
614:   return 0;
615: }

617: /*TEST

619:   build:
620:     requires: !complex

622:   test:
623:     suffix: 0
624:     args:

626:   test:
627:     suffix: l1_1
628:     args: -p 1 -tao_type lmvm -alpha 1. -epsilon 1.e-7 -m 64 -n 64 -view_sol -matrix_format 1

630:   test:
631:     suffix: hessian_1
632:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nls

634:   test:
635:     suffix: hessian_2
636:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nls

638:   test:
639:     suffix: nm_1
640:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nm -tao_max_it 50

642:   test:
643:     suffix: nm_2
644:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nm -tao_max_it 50

646:   test:
647:     suffix: lmvm_1
648:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type lmvm -tao_max_it 40

650:   test:
651:     suffix: lmvm_2
652:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type lmvm -tao_max_it 15

654:   test:
655:     suffix: soft_threshold_admm_1
656:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm

658:   test:
659:     suffix: hessian_admm_1
660:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nls -misfit_tao_type nls

662:   test:
663:     suffix: hessian_admm_2
664:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nls -misfit_tao_type nls

666:   test:
667:     suffix: nm_admm_1
668:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nm -misfit_tao_type nm

670:   test:
671:     suffix: nm_admm_2
672:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nm -misfit_tao_type nm -iter 7

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

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

682: TEST*/