Actual source code: ex4.c
petsc-3.14.6 2021-03-30
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, ®, _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*/