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, ®, _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*/