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>
7: #define NWORKLEFT 4
8: #define NWORKRIGHT 12
10: typedef struct _UserCtx {
11: PetscInt m; /* The row dimension of F */
12: PetscInt n; /* The column dimension of F */
13: PetscInt matops; /* Matrix format. 0 for stencil, 1 for random */
14: PetscInt iter; /* Number of iterations for ADMM */
15: PetscReal hStart; /* Starting point for Taylor test */
16: PetscReal hFactor; /* Taylor test step factor */
17: PetscReal hMin; /* Taylor test end goal */
18: PetscReal alpha; /* regularization constant applied to || x ||_p */
19: PetscReal eps; /* small constant for approximating gradient of || x ||_1 */
20: PetscReal mu; /* the augmented Lagrangian term in ADMM */
21: PetscReal abstol;
22: PetscReal reltol;
23: Mat F; /* matrix in least squares component $(1/2) \| F x - d \|_2^2$ */
24: Mat W; /* Workspace matrix. ATA */
25: Mat Hm; /* Hessian Misfit*/
26: Mat Hr; /* Hessian Reg*/
27: Vec d; /* RHS in least squares component $(1/2) \| F x - d \|_2^2$ */
28: Vec workLeft[NWORKLEFT]; /* Workspace for temporary vec */
29: Vec workRight[NWORKRIGHT]; /* Workspace for temporary vec */
30: NormType p;
31: PetscRandom rctx;
32: PetscBool soft;
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: PetscFunctionBegin;
40: /* build the rhs d in ctx */
41: PetscCall(VecCreate(PETSC_COMM_WORLD, &ctx->d));
42: PetscCall(VecSetSizes(ctx->d, PETSC_DECIDE, ctx->m));
43: PetscCall(VecSetFromOptions(ctx->d));
44: PetscCall(VecSetRandom(ctx->d, ctx->rctx));
45: PetscFunctionReturn(PETSC_SUCCESS);
46: }
48: static PetscErrorCode CreateMatrix(UserCtx ctx)
49: {
50: PetscInt Istart, Iend, i, j, Ii, gridN, I_n, I_s, I_e, I_w;
51: PetscLogStage stage;
53: PetscFunctionBegin;
54: /* build the matrix F in ctx */
55: PetscCall(MatCreate(PETSC_COMM_WORLD, &ctx->F));
56: PetscCall(MatSetSizes(ctx->F, PETSC_DECIDE, PETSC_DECIDE, ctx->m, ctx->n));
57: PetscCall(MatSetType(ctx->F, MATAIJ)); /* TODO: Decide specific SetType other than dummy*/
58: PetscCall(MatMPIAIJSetPreallocation(ctx->F, 5, NULL, 5, NULL)); /*TODO: some number other than 5?*/
59: PetscCall(MatSeqAIJSetPreallocation(ctx->F, 5, NULL));
60: PetscCall(MatSetUp(ctx->F));
61: PetscCall(MatGetOwnershipRange(ctx->F, &Istart, &Iend));
62: PetscCall(PetscLogStageRegister("Assembly", &stage));
63: PetscCall(PetscLogStagePush(stage));
65: /* Set matrix elements in 2-D five point stencil format. */
66: if (!ctx->matops) {
67: PetscCheck(ctx->m == ctx->n, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "Stencil matrix must be square");
68: gridN = (PetscInt)PetscSqrtReal((PetscReal)ctx->m);
69: PetscCheck(gridN * gridN == ctx->m, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "Number of rows must be square");
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: PetscCall(MatSetValue(ctx->F, Ii, Ii, 4., INSERT_VALUES));
82: PetscCall(MatSetValue(ctx->F, Ii, I_n, -1., INSERT_VALUES));
83: PetscCall(MatSetValue(ctx->F, Ii, I_s, -1., INSERT_VALUES));
84: PetscCall(MatSetValue(ctx->F, Ii, I_e, -1., INSERT_VALUES));
85: PetscCall(MatSetValue(ctx->F, Ii, I_w, -1., INSERT_VALUES));
86: }
87: } else PetscCall(MatSetRandom(ctx->F, ctx->rctx));
88: PetscCall(MatAssemblyBegin(ctx->F, MAT_FINAL_ASSEMBLY));
89: PetscCall(MatAssemblyEnd(ctx->F, MAT_FINAL_ASSEMBLY));
90: PetscCall(PetscLogStagePop());
91: /* Stencil matrix is symmetric. Setting symmetric flag for ICC/Cholesky preconditioner */
92: if (!ctx->matops) PetscCall(MatSetOption(ctx->F, MAT_SYMMETRIC, PETSC_TRUE));
93: PetscCall(MatTransposeMatMult(ctx->F, ctx->F, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &ctx->W));
94: /* Setup Hessian Workspace in same shape as W */
95: PetscCall(MatDuplicate(ctx->W, MAT_DO_NOT_COPY_VALUES, &ctx->Hm));
96: PetscCall(MatDuplicate(ctx->W, MAT_DO_NOT_COPY_VALUES, &ctx->Hr));
97: PetscFunctionReturn(PETSC_SUCCESS);
98: }
100: static PetscErrorCode SetupWorkspace(UserCtx ctx)
101: {
102: PetscFunctionBegin;
103: PetscCall(MatCreateVecs(ctx->F, &ctx->workLeft[0], &ctx->workRight[0]));
104: for (PetscInt i = 1; i < NWORKLEFT; i++) PetscCall(VecDuplicate(ctx->workLeft[0], &ctx->workLeft[i]));
105: for (PetscInt i = 1; i < NWORKRIGHT; i++) PetscCall(VecDuplicate(ctx->workRight[0], &ctx->workRight[i]));
106: PetscFunctionReturn(PETSC_SUCCESS);
107: }
109: static PetscErrorCode ConfigureContext(UserCtx ctx)
110: {
111: PetscFunctionBegin;
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->soft = PETSC_FALSE;
126: ctx->taylor = PETSC_TRUE;
127: ctx->use_admm = PETSC_FALSE;
128: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Configure separable objective example", "ex4.c");
129: PetscCall(PetscOptionsInt("-m", "The row dimension of matrix F", "ex4.c", ctx->m, &ctx->m, NULL));
130: PetscCall(PetscOptionsInt("-n", "The column dimension of matrix F", "ex4.c", ctx->n, &ctx->n, NULL));
131: PetscCall(PetscOptionsInt("-matrix_format", "Decide format of F matrix. 0 for stencil, 1 for random", "ex4.c", ctx->matops, &ctx->matops, NULL));
132: PetscCall(PetscOptionsInt("-iter", "Iteration number ADMM", "ex4.c", ctx->iter, &ctx->iter, NULL));
133: PetscCall(PetscOptionsReal("-alpha", "The regularization multiplier. 1 default", "ex4.c", ctx->alpha, &ctx->alpha, NULL));
134: PetscCall(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));
135: PetscCall(PetscOptionsReal("-mu", "The augmented lagrangian multiplier in ADMM", "ex4.c", ctx->mu, &ctx->mu, NULL));
136: PetscCall(PetscOptionsReal("-hStart", "Taylor test starting point. 1 default.", "ex4.c", ctx->hStart, &ctx->hStart, NULL));
137: PetscCall(PetscOptionsReal("-hFactor", "Taylor test multiplier factor. 0.5 default", "ex4.c", ctx->hFactor, &ctx->hFactor, NULL));
138: PetscCall(PetscOptionsReal("-hMin", "Taylor test ending condition. 1.e-3 default", "ex4.c", ctx->hMin, &ctx->hMin, NULL));
139: PetscCall(PetscOptionsReal("-abstol", "Absolute stopping criterion for ADMM", "ex4.c", ctx->abstol, &ctx->abstol, NULL));
140: PetscCall(PetscOptionsReal("-reltol", "Relative stopping criterion for ADMM", "ex4.c", ctx->reltol, &ctx->reltol, NULL));
141: PetscCall(PetscOptionsBool("-taylor", "Flag for Taylor test. Default is true.", "ex4.c", ctx->taylor, &ctx->taylor, NULL));
142: PetscCall(PetscOptionsBool("-soft", "Flag for testing soft threshold no-op case. Default is false.", "ex4.c", ctx->soft, &ctx->soft, NULL));
143: PetscCall(PetscOptionsBool("-use_admm", "Use the ADMM solver in this example.", "ex4.c", ctx->use_admm, &ctx->use_admm, NULL));
144: PetscCall(PetscOptionsEnum("-p", "Norm type.", "ex4.c", NormTypes, (PetscEnum)ctx->p, (PetscEnum *)&ctx->p, NULL));
145: PetscOptionsEnd();
146: /* Creating random ctx */
147: PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &ctx->rctx));
148: PetscCall(PetscRandomSetFromOptions(ctx->rctx));
149: PetscCall(CreateMatrix(ctx));
150: PetscCall(CreateRHS(ctx));
151: PetscCall(SetupWorkspace(ctx));
152: PetscFunctionReturn(PETSC_SUCCESS);
153: }
155: static PetscErrorCode DestroyContext(UserCtx *ctx)
156: {
157: PetscFunctionBegin;
158: PetscCall(MatDestroy(&(*ctx)->F));
159: PetscCall(MatDestroy(&(*ctx)->W));
160: PetscCall(MatDestroy(&(*ctx)->Hm));
161: PetscCall(MatDestroy(&(*ctx)->Hr));
162: PetscCall(VecDestroy(&(*ctx)->d));
163: for (PetscInt i = 0; i < NWORKLEFT; i++) PetscCall(VecDestroy(&(*ctx)->workLeft[i]));
164: for (PetscInt i = 0; i < NWORKRIGHT; i++) PetscCall(VecDestroy(&(*ctx)->workRight[i]));
165: PetscCall(PetscRandomDestroy(&(*ctx)->rctx));
166: PetscCall(PetscFree(*ctx));
167: PetscFunctionReturn(PETSC_SUCCESS);
168: }
170: /* compute (1/2) * ||F x - d||^2 */
171: static PetscErrorCode ObjectiveMisfit(Tao tao, Vec x, PetscReal *J, void *_ctx)
172: {
173: UserCtx ctx = (UserCtx)_ctx;
174: Vec y;
176: PetscFunctionBegin;
177: y = ctx->workLeft[0];
178: PetscCall(MatMult(ctx->F, x, y));
179: PetscCall(VecAXPY(y, -1., ctx->d));
180: PetscCall(VecDot(y, y, J));
181: *J *= 0.5;
182: PetscFunctionReturn(PETSC_SUCCESS);
183: }
185: /* compute V = FTFx - FTd */
186: static PetscErrorCode GradientMisfit(Tao tao, Vec x, Vec V, void *_ctx)
187: {
188: UserCtx ctx = (UserCtx)_ctx;
189: Vec FTFx, FTd;
191: PetscFunctionBegin;
192: /* work1 is A^T Ax, work2 is Ab, W is A^T A*/
193: FTFx = ctx->workRight[0];
194: FTd = ctx->workRight[1];
195: PetscCall(MatMult(ctx->W, x, FTFx));
196: PetscCall(MatMultTranspose(ctx->F, ctx->d, FTd));
197: PetscCall(VecWAXPY(V, -1., FTd, FTFx));
198: PetscFunctionReturn(PETSC_SUCCESS);
199: }
201: /* returns FTF */
202: static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
203: {
204: UserCtx ctx = (UserCtx)_ctx;
206: PetscFunctionBegin;
207: if (H != ctx->W) PetscCall(MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN));
208: if (Hpre != ctx->W) PetscCall(MatCopy(ctx->W, Hpre, DIFFERENT_NONZERO_PATTERN));
209: PetscFunctionReturn(PETSC_SUCCESS);
210: }
212: /* computes augment Lagrangian objective (with scaled dual):
213: * 0.5 * ||F x - d||^2 + 0.5 * mu ||x - z + u||^2 */
214: static PetscErrorCode ObjectiveMisfitADMM(Tao tao, Vec x, PetscReal *J, void *_ctx)
215: {
216: UserCtx ctx = (UserCtx)_ctx;
217: PetscReal mu, workNorm, misfit;
218: Vec z, u, temp;
220: PetscFunctionBegin;
221: mu = ctx->mu;
222: z = ctx->workRight[5];
223: u = ctx->workRight[6];
224: temp = ctx->workRight[10];
225: /* misfit = f(x) */
226: PetscCall(ObjectiveMisfit(tao, x, &misfit, _ctx));
227: PetscCall(VecCopy(x, temp));
228: /* temp = x - z + u */
229: PetscCall(VecAXPBYPCZ(temp, -1., 1., 1., z, u));
230: /* workNorm = ||x - z + u||^2 */
231: PetscCall(VecDot(temp, temp, &workNorm));
232: /* augment Lagrangian objective (with scaled dual): f(x) + 0.5 * mu ||x -z + u||^2 */
233: *J = misfit + 0.5 * mu * workNorm;
234: PetscFunctionReturn(PETSC_SUCCESS);
235: }
237: /* computes FTFx - FTd mu*(x - z + u) */
238: static PetscErrorCode GradientMisfitADMM(Tao tao, Vec x, Vec V, void *_ctx)
239: {
240: UserCtx ctx = (UserCtx)_ctx;
241: PetscReal mu;
242: Vec z, u, temp;
244: PetscFunctionBegin;
245: mu = ctx->mu;
246: z = ctx->workRight[5];
247: u = ctx->workRight[6];
248: temp = ctx->workRight[10];
249: PetscCall(GradientMisfit(tao, x, V, _ctx));
250: PetscCall(VecCopy(x, temp));
251: /* temp = x - z + u */
252: PetscCall(VecAXPBYPCZ(temp, -1., 1., 1., z, u));
253: /* V = FTFx - FTd mu*(x - z + u) */
254: PetscCall(VecAXPY(V, mu, temp));
255: PetscFunctionReturn(PETSC_SUCCESS);
256: }
258: /* returns FTF + diag(mu) */
259: static PetscErrorCode HessianMisfitADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
260: {
261: UserCtx ctx = (UserCtx)_ctx;
263: PetscFunctionBegin;
264: PetscCall(MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN));
265: PetscCall(MatShift(H, ctx->mu));
266: if (Hpre != H) PetscCall(MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN));
267: PetscFunctionReturn(PETSC_SUCCESS);
268: }
270: /* computes || x ||_p (mult by 0.5 in case of NORM_2) */
271: static PetscErrorCode ObjectiveRegularization(Tao tao, Vec x, PetscReal *J, void *_ctx)
272: {
273: UserCtx ctx = (UserCtx)_ctx;
274: PetscReal norm;
276: PetscFunctionBegin;
277: *J = 0;
278: PetscCall(VecNorm(x, ctx->p, &norm));
279: if (ctx->p == NORM_2) norm = 0.5 * norm * norm;
280: *J = ctx->alpha * norm;
281: PetscFunctionReturn(PETSC_SUCCESS);
282: }
284: /* NORM_2 Case: return x
285: * NORM_1 Case: x/(|x| + eps)
286: * Else: TODO */
287: static PetscErrorCode GradientRegularization(Tao tao, Vec x, Vec V, void *_ctx)
288: {
289: UserCtx ctx = (UserCtx)_ctx;
290: PetscReal eps = ctx->eps;
292: PetscFunctionBegin;
293: if (ctx->p == NORM_2) PetscCall(VecCopy(x, V));
294: else if (ctx->p == NORM_1) {
295: PetscCall(VecCopy(x, ctx->workRight[1]));
296: PetscCall(VecAbs(ctx->workRight[1]));
297: PetscCall(VecShift(ctx->workRight[1], eps));
298: PetscCall(VecPointwiseDivide(V, x, ctx->workRight[1]));
299: } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
300: PetscFunctionReturn(PETSC_SUCCESS);
301: }
303: /* NORM_2 Case: returns diag(mu)
304: * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) */
305: static PetscErrorCode HessianRegularization(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
306: {
307: UserCtx ctx = (UserCtx)_ctx;
308: PetscReal eps = ctx->eps;
309: Vec copy1, copy2, copy3;
311: PetscFunctionBegin;
312: if (ctx->p == NORM_2) {
313: /* Identity matrix scaled by mu */
314: PetscCall(MatZeroEntries(H));
315: PetscCall(MatShift(H, ctx->mu));
316: if (Hpre != H) {
317: PetscCall(MatZeroEntries(Hpre));
318: PetscCall(MatShift(Hpre, ctx->mu));
319: }
320: } else if (ctx->p == NORM_1) {
321: /* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)) */
322: copy1 = ctx->workRight[1];
323: copy2 = ctx->workRight[2];
324: copy3 = ctx->workRight[3];
325: /* copy1 : 1/sqrt(x_i^2 + eps) */
326: PetscCall(VecCopy(x, copy1));
327: PetscCall(VecPow(copy1, 2));
328: PetscCall(VecShift(copy1, eps));
329: PetscCall(VecSqrtAbs(copy1));
330: PetscCall(VecReciprocal(copy1));
331: /* copy2: x_i^2.*/
332: PetscCall(VecCopy(x, copy2));
333: PetscCall(VecPow(copy2, 2));
334: /* copy3: abs(x_i^2 + eps) */
335: PetscCall(VecCopy(x, copy3));
336: PetscCall(VecPow(copy3, 2));
337: PetscCall(VecShift(copy3, eps));
338: PetscCall(VecAbs(copy3));
339: /* copy2: 1 - x_i^2/abs(x_i^2 + eps) */
340: PetscCall(VecPointwiseDivide(copy2, copy2, copy3));
341: PetscCall(VecScale(copy2, -1.));
342: PetscCall(VecShift(copy2, 1.));
343: PetscCall(VecAXPY(copy1, 1., copy2));
344: PetscCall(VecScale(copy1, ctx->mu));
345: PetscCall(MatZeroEntries(H));
346: PetscCall(MatDiagonalSet(H, copy1, INSERT_VALUES));
347: if (Hpre != H) {
348: PetscCall(MatZeroEntries(Hpre));
349: PetscCall(MatDiagonalSet(Hpre, copy1, INSERT_VALUES));
350: }
351: } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
352: PetscFunctionReturn(PETSC_SUCCESS);
353: }
355: /* NORM_2 Case: 0.5 || x ||_2 + 0.5 * mu * ||x + u - z||^2
356: * Else : || x ||_2 + 0.5 * mu * ||x + u - z||^2 */
357: static PetscErrorCode ObjectiveRegularizationADMM(Tao tao, Vec z, PetscReal *J, void *_ctx)
358: {
359: UserCtx ctx = (UserCtx)_ctx;
360: PetscReal mu, workNorm, reg;
361: Vec x, u, temp;
363: PetscFunctionBegin;
364: mu = ctx->mu;
365: x = ctx->workRight[4];
366: u = ctx->workRight[6];
367: temp = ctx->workRight[10];
368: PetscCall(ObjectiveRegularization(tao, z, ®, _ctx));
369: PetscCall(VecCopy(z, temp));
370: /* temp = x + u -z */
371: PetscCall(VecAXPBYPCZ(temp, 1., 1., -1., x, u));
372: /* workNorm = ||x + u - z ||^2 */
373: PetscCall(VecDot(temp, temp, &workNorm));
374: *J = reg + 0.5 * mu * workNorm;
375: PetscFunctionReturn(PETSC_SUCCESS);
376: }
378: /* NORM_2 Case: x - mu*(x + u - z)
379: * NORM_1 Case: x/(|x| + eps) - mu*(x + u - z)
380: * Else: TODO */
381: static PetscErrorCode GradientRegularizationADMM(Tao tao, Vec z, Vec V, void *_ctx)
382: {
383: UserCtx ctx = (UserCtx)_ctx;
384: PetscReal mu;
385: Vec x, u, temp;
387: PetscFunctionBegin;
388: mu = ctx->mu;
389: x = ctx->workRight[4];
390: u = ctx->workRight[6];
391: temp = ctx->workRight[10];
392: PetscCall(GradientRegularization(tao, z, V, _ctx));
393: PetscCall(VecCopy(z, temp));
394: /* temp = x + u -z */
395: PetscCall(VecAXPBYPCZ(temp, 1., 1., -1., x, u));
396: PetscCall(VecAXPY(V, -mu, temp));
397: PetscFunctionReturn(PETSC_SUCCESS);
398: }
400: /* NORM_2 Case: returns diag(mu)
401: * NORM_1 Case: FTF + diag(mu) */
402: static PetscErrorCode HessianRegularizationADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
403: {
404: UserCtx ctx = (UserCtx)_ctx;
406: PetscFunctionBegin;
407: if (ctx->p == NORM_2) {
408: /* Identity matrix scaled by mu */
409: PetscCall(MatZeroEntries(H));
410: PetscCall(MatShift(H, ctx->mu));
411: if (Hpre != H) {
412: PetscCall(MatZeroEntries(Hpre));
413: PetscCall(MatShift(Hpre, ctx->mu));
414: }
415: } else if (ctx->p == NORM_1) {
416: PetscCall(HessianMisfit(tao, x, H, Hpre, (void *)ctx));
417: PetscCall(MatShift(H, ctx->mu));
418: if (Hpre != H) PetscCall(MatShift(Hpre, ctx->mu));
419: } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
420: PetscFunctionReturn(PETSC_SUCCESS);
421: }
423: /* NORM_2 Case : (1/2) * ||F x - d||^2 + 0.5 * || x ||_p
424: * NORM_1 Case : (1/2) * ||F x - d||^2 + || x ||_p */
425: static PetscErrorCode ObjectiveComplete(Tao tao, Vec x, PetscReal *J, PetscCtx ctx)
426: {
427: PetscReal Jm, Jr;
429: PetscFunctionBegin;
430: PetscCall(ObjectiveMisfit(tao, x, &Jm, ctx));
431: PetscCall(ObjectiveRegularization(tao, x, &Jr, ctx));
432: *J = Jm + Jr;
433: PetscFunctionReturn(PETSC_SUCCESS);
434: }
436: /* NORM_2 Case: FTFx - FTd + x
437: * NORM_1 Case: FTFx - FTd + x/(|x| + eps) */
438: static PetscErrorCode GradientComplete(Tao tao, Vec x, Vec V, PetscCtx ctx)
439: {
440: UserCtx cntx = (UserCtx)ctx;
442: PetscFunctionBegin;
443: PetscCall(GradientMisfit(tao, x, cntx->workRight[2], ctx));
444: PetscCall(GradientRegularization(tao, x, cntx->workRight[3], ctx));
445: PetscCall(VecWAXPY(V, 1, cntx->workRight[2], cntx->workRight[3]));
446: PetscFunctionReturn(PETSC_SUCCESS);
447: }
449: /* NORM_2 Case: diag(mu) + FTF
450: * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) + FTF */
451: static PetscErrorCode HessianComplete(Tao tao, Vec x, Mat H, Mat Hpre, PetscCtx ctx)
452: {
453: Mat tempH;
455: PetscFunctionBegin;
456: PetscCall(MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH));
457: PetscCall(HessianMisfit(tao, x, H, H, ctx));
458: PetscCall(HessianRegularization(tao, x, tempH, tempH, ctx));
459: PetscCall(MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN));
460: if (Hpre != H) PetscCall(MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN));
461: PetscCall(MatDestroy(&tempH));
462: PetscFunctionReturn(PETSC_SUCCESS);
463: }
465: static PetscErrorCode TaoSolveADMM(UserCtx ctx, Vec x)
466: {
467: PetscInt i;
468: PetscReal u_norm, r_norm, s_norm, primal, dual, x_norm, z_norm;
469: Tao tao1, tao2;
470: Vec xk, z, u, diff, zold, zdiff, temp;
471: PetscReal mu;
473: PetscFunctionBegin;
474: xk = ctx->workRight[4];
475: z = ctx->workRight[5];
476: u = ctx->workRight[6];
477: diff = ctx->workRight[7];
478: zold = ctx->workRight[8];
479: zdiff = ctx->workRight[9];
480: temp = ctx->workRight[11];
481: mu = ctx->mu;
482: PetscCall(VecSet(u, 0.));
483: PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao1));
484: PetscCall(TaoSetType(tao1, TAONLS));
485: PetscCall(TaoSetObjective(tao1, ObjectiveMisfitADMM, (void *)ctx));
486: PetscCall(TaoSetGradient(tao1, NULL, GradientMisfitADMM, (void *)ctx));
487: PetscCall(TaoSetHessian(tao1, ctx->Hm, ctx->Hm, HessianMisfitADMM, (void *)ctx));
488: PetscCall(VecSet(xk, 0.));
489: PetscCall(TaoSetSolution(tao1, xk));
490: PetscCall(TaoSetOptionsPrefix(tao1, "misfit_"));
491: PetscCall(TaoSetFromOptions(tao1));
492: PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao2));
493: if (ctx->p == NORM_2) {
494: PetscCall(TaoSetType(tao2, TAONLS));
495: PetscCall(TaoSetObjective(tao2, ObjectiveRegularizationADMM, (void *)ctx));
496: PetscCall(TaoSetGradient(tao2, NULL, GradientRegularizationADMM, (void *)ctx));
497: PetscCall(TaoSetHessian(tao2, ctx->Hr, ctx->Hr, HessianRegularizationADMM, (void *)ctx));
498: }
499: PetscCall(VecSet(z, 0.));
500: PetscCall(TaoSetSolution(tao2, z));
501: PetscCall(TaoSetOptionsPrefix(tao2, "reg_"));
502: PetscCall(TaoSetFromOptions(tao2));
504: for (i = 0; i < ctx->iter; i++) {
505: PetscCall(VecCopy(z, zold));
506: PetscCall(TaoSolve(tao1)); /* Updates xk */
507: if (ctx->p == NORM_1) {
508: PetscCall(VecWAXPY(temp, 1., xk, u));
509: PetscCall(TaoSoftThreshold(temp, -ctx->alpha / mu, ctx->alpha / mu, z));
510: } else {
511: PetscCall(TaoSolve(tao2)); /* Update zk */
512: }
513: /* u = u + xk -z */
514: PetscCall(VecAXPBYPCZ(u, 1., -1., 1., xk, z));
515: /* r_norm : norm(x-z) */
516: PetscCall(VecWAXPY(diff, -1., z, xk));
517: PetscCall(VecNorm(diff, NORM_2, &r_norm));
518: /* s_norm : norm(-mu(z-zold)) */
519: PetscCall(VecWAXPY(zdiff, -1., zold, z));
520: PetscCall(VecNorm(zdiff, NORM_2, &s_norm));
521: s_norm = s_norm * mu;
522: /* primal : sqrt(n)*ABSTOL + RELTOL*max(norm(x), norm(-z))*/
523: PetscCall(VecNorm(xk, NORM_2, &x_norm));
524: PetscCall(VecNorm(z, NORM_2, &z_norm));
525: primal = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * PetscMax(x_norm, z_norm);
526: /* Duality : sqrt(n)*ABSTOL + RELTOL*norm(mu*u)*/
527: PetscCall(VecNorm(u, NORM_2, &u_norm));
528: dual = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * u_norm * mu;
529: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao1), "Iter %" PetscInt_FMT " : ||x-z||: %g, mu*||z-zold||: %g\n", i, (double)r_norm, (double)s_norm));
530: if (r_norm < primal && s_norm < dual) break;
531: }
532: PetscCall(VecCopy(xk, x));
533: PetscCall(TaoDestroy(&tao1));
534: PetscCall(TaoDestroy(&tao2));
535: PetscFunctionReturn(PETSC_SUCCESS);
536: }
538: /* Second order Taylor remainder convergence test */
539: static PetscErrorCode TaylorTest(UserCtx ctx, Tao tao, Vec x, PetscReal *C)
540: {
541: PetscReal h, J, temp;
542: PetscInt i;
543: PetscInt numValues;
544: PetscReal Jx, Jxhat_comp, Jxhat_pred;
545: PetscReal *Js, *hs;
546: PetscReal gdotdx;
547: PetscReal minrate = PETSC_MAX_REAL;
548: MPI_Comm comm = PetscObjectComm((PetscObject)x);
549: Vec g, dx, xhat;
551: PetscFunctionBegin;
552: PetscCall(VecDuplicate(x, &g));
553: PetscCall(VecDuplicate(x, &xhat));
554: /* choose a perturbation direction */
555: PetscCall(VecDuplicate(x, &dx));
556: PetscCall(VecSetRandom(dx, ctx->rctx));
557: /* evaluate objective at x: J(x) */
558: PetscCall(TaoComputeObjective(tao, x, &Jx));
559: /* evaluate gradient at x, save in vector g */
560: PetscCall(TaoComputeGradient(tao, x, g));
561: PetscCall(VecDot(g, dx, &gdotdx));
563: for (numValues = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor) numValues++;
564: PetscCall(PetscCalloc2(numValues, &Js, numValues, &hs));
565: for (i = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor, i++) {
566: PetscCall(VecWAXPY(xhat, h, dx, x));
567: PetscCall(TaoComputeObjective(tao, xhat, &Jxhat_comp));
568: /* J(\hat(x)) \approx J(x) + g^T (xhat - x) = J(x) + h * g^T dx */
569: Jxhat_pred = Jx + h * gdotdx;
570: /* Vector to dJdm scalar? Dot?*/
571: J = PetscAbsReal(Jxhat_comp - Jxhat_pred);
572: PetscCall(PetscPrintf(comm, "J(xhat): %g, predicted: %g, diff %g\n", (double)Jxhat_comp, (double)Jxhat_pred, (double)J));
573: Js[i] = J;
574: hs[i] = h;
575: }
576: for (PetscInt j = 1; j < numValues; j++) {
577: temp = PetscLogReal(Js[j] / Js[j - 1]) / PetscLogReal(hs[j] / hs[j - 1]);
578: PetscCall(PetscPrintf(comm, "Convergence rate step %" PetscInt_FMT ": %g\n", j - 1, (double)temp));
579: minrate = PetscMin(minrate, temp);
580: }
581: /* If O is not ~2, then the test is wrong */
582: PetscCall(PetscFree2(Js, hs));
583: *C = minrate;
584: PetscCall(VecDestroy(&dx));
585: PetscCall(VecDestroy(&xhat));
586: PetscCall(VecDestroy(&g));
587: PetscFunctionReturn(PETSC_SUCCESS);
588: }
590: int main(int argc, char **argv)
591: {
592: UserCtx ctx;
593: Tao tao;
594: Vec x;
595: Mat H;
597: PetscFunctionBeginUser;
598: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
599: PetscCall(PetscNew(&ctx));
600: PetscCall(ConfigureContext(ctx));
601: /* Define two functions that could pass as objectives to TaoSetObjective(): one
602: * for the misfit component, and one for the regularization component */
603: /* ObjectiveMisfit() and ObjectiveRegularization() */
605: /* Define a single function that calls both components adds them together: the complete objective,
606: * in the absence of a Tao implementation that handles separability */
607: /* ObjectiveComplete() */
608: PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
609: PetscCall(TaoSetType(tao, TAONM));
610: PetscCall(TaoSetObjective(tao, ObjectiveComplete, (void *)ctx));
611: PetscCall(TaoSetGradient(tao, NULL, GradientComplete, (void *)ctx));
612: PetscCall(MatDuplicate(ctx->W, MAT_SHARE_NONZERO_PATTERN, &H));
613: PetscCall(TaoSetHessian(tao, H, H, HessianComplete, (void *)ctx));
614: PetscCall(MatCreateVecs(ctx->F, NULL, &x));
615: PetscCall(TaoSetSolution(tao, x));
616: PetscCall(TaoSetFromOptions(tao));
617: if (ctx->use_admm) PetscCall(TaoSolveADMM(ctx, x));
618: else PetscCall(TaoSolve(tao));
619: /* examine solution */
620: PetscCall(VecViewFromOptions(x, NULL, "-view_sol"));
621: if (ctx->taylor) {
622: PetscReal rate;
623: PetscCall(TaylorTest(ctx, tao, x, &rate));
624: }
625: if (ctx->soft) PetscCall(TaoSoftThreshold(x, 0., 0., x));
626: PetscCall(MatDestroy(&H));
627: PetscCall(TaoDestroy(&tao));
628: PetscCall(VecDestroy(&x));
629: PetscCall(DestroyContext(&ctx));
630: PetscCall(PetscFinalize());
631: return 0;
632: }
634: /*TEST
636: build:
637: requires: !complex
639: test:
640: suffix: 0
641: args:
643: test:
644: suffix: l1_1
645: args: -p 1 -tao_type lmvm -alpha 1. -epsilon 1.e-7 -m 64 -n 64 -view_sol -matrix_format 1
647: test:
648: suffix: hessian_1
649: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nls
651: test:
652: suffix: hessian_2
653: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nls
655: test:
656: suffix: nm_1
657: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nm -tao_max_it 50
659: test:
660: suffix: nm_2
661: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nm -tao_max_it 50
663: test:
664: suffix: lmvm_1
665: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type lmvm -tao_max_it 40
667: test:
668: suffix: lmvm_2
669: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type lmvm -tao_max_it 15
671: test:
672: suffix: soft_threshold_admm_1
673: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm
675: test:
676: suffix: hessian_admm_1
677: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nls -misfit_tao_type nls
679: test:
680: suffix: hessian_admm_2
681: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nls -misfit_tao_type nls
683: test:
684: suffix: nm_admm_1
685: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nm -misfit_tao_type nm
687: test:
688: suffix: nm_admm_2
689: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nm -misfit_tao_type nm -iter 7
691: test:
692: suffix: lmvm_admm_1
693: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm
695: test:
696: suffix: lmvm_admm_2
697: args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm
699: test:
700: suffix: soft
701: args: -taylor 0 -soft 1
702: output_file: output/empty.out
704: test:
705: suffix: soft_view_10
706: args: -taylor 0 -soft 1 -tao_view -tao_max_funcs 10
708: TEST*/