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;    /* Number 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   soft;
 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: {
 40:   PetscFunctionBegin;
 41:   /* build the rhs d in ctx */
 42:   PetscCall(VecCreate(PETSC_COMM_WORLD, &ctx->d));
 43:   PetscCall(VecSetSizes(ctx->d, PETSC_DECIDE, ctx->m));
 44:   PetscCall(VecSetFromOptions(ctx->d));
 45:   PetscCall(VecSetRandom(ctx->d, ctx->rctx));
 46:   PetscFunctionReturn(PETSC_SUCCESS);
 47: }

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

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

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

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

105:   PetscFunctionBegin;
106:   PetscCall(MatCreateVecs(ctx->F, &ctx->workLeft[0], &ctx->workRight[0]));
107:   for (i = 1; i < NWORKLEFT; i++) PetscCall(VecDuplicate(ctx->workLeft[0], &ctx->workLeft[i]));
108:   for (i = 1; i < NWORKRIGHT; i++) PetscCall(VecDuplicate(ctx->workRight[0], &ctx->workRight[i]));
109:   PetscFunctionReturn(PETSC_SUCCESS);
110: }

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

158: static PetscErrorCode DestroyContext(UserCtx *ctx)
159: {
160:   PetscInt i;

162:   PetscFunctionBegin;
163:   PetscCall(MatDestroy(&((*ctx)->F)));
164:   PetscCall(MatDestroy(&((*ctx)->W)));
165:   PetscCall(MatDestroy(&((*ctx)->Hm)));
166:   PetscCall(MatDestroy(&((*ctx)->Hr)));
167:   PetscCall(VecDestroy(&((*ctx)->d)));
168:   for (i = 0; i < NWORKLEFT; i++) PetscCall(VecDestroy(&((*ctx)->workLeft[i])));
169:   for (i = 0; i < NWORKRIGHT; i++) PetscCall(VecDestroy(&((*ctx)->workRight[i])));
170:   PetscCall(PetscRandomDestroy(&((*ctx)->rctx)));
171:   PetscCall(PetscFree(*ctx));
172:   PetscFunctionReturn(PETSC_SUCCESS);
173: }

175: /* compute (1/2) * ||F x - d||^2 */
176: static PetscErrorCode ObjectiveMisfit(Tao tao, Vec x, PetscReal *J, void *_ctx)
177: {
178:   UserCtx ctx = (UserCtx)_ctx;
179:   Vec     y;

181:   PetscFunctionBegin;
182:   y = ctx->workLeft[0];
183:   PetscCall(MatMult(ctx->F, x, y));
184:   PetscCall(VecAXPY(y, -1., ctx->d));
185:   PetscCall(VecDot(y, y, J));
186:   *J *= 0.5;
187:   PetscFunctionReturn(PETSC_SUCCESS);
188: }

190: /* compute V = FTFx - FTd */
191: static PetscErrorCode GradientMisfit(Tao tao, Vec x, Vec V, void *_ctx)
192: {
193:   UserCtx ctx = (UserCtx)_ctx;
194:   Vec     FTFx, FTd;

196:   PetscFunctionBegin;
197:   /* work1 is A^T Ax, work2 is Ab, W is A^T A*/
198:   FTFx = ctx->workRight[0];
199:   FTd  = ctx->workRight[1];
200:   PetscCall(MatMult(ctx->W, x, FTFx));
201:   PetscCall(MatMultTranspose(ctx->F, ctx->d, FTd));
202:   PetscCall(VecWAXPY(V, -1., FTd, FTFx));
203:   PetscFunctionReturn(PETSC_SUCCESS);
204: }

206: /* returns FTF */
207: static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
208: {
209:   UserCtx ctx = (UserCtx)_ctx;

211:   PetscFunctionBegin;
212:   if (H != ctx->W) PetscCall(MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN));
213:   if (Hpre != ctx->W) PetscCall(MatCopy(ctx->W, Hpre, DIFFERENT_NONZERO_PATTERN));
214:   PetscFunctionReturn(PETSC_SUCCESS);
215: }

217: /* computes augment Lagrangian objective (with scaled dual):
218:  * 0.5 * ||F x - d||^2  + 0.5 * mu ||x - z + u||^2 */
219: static PetscErrorCode ObjectiveMisfitADMM(Tao tao, Vec x, PetscReal *J, void *_ctx)
220: {
221:   UserCtx   ctx = (UserCtx)_ctx;
222:   PetscReal mu, workNorm, misfit;
223:   Vec       z, u, temp;

225:   PetscFunctionBegin;
226:   mu   = ctx->mu;
227:   z    = ctx->workRight[5];
228:   u    = ctx->workRight[6];
229:   temp = ctx->workRight[10];
230:   /* misfit = f(x) */
231:   PetscCall(ObjectiveMisfit(tao, x, &misfit, _ctx));
232:   PetscCall(VecCopy(x, temp));
233:   /* temp = x - z + u */
234:   PetscCall(VecAXPBYPCZ(temp, -1., 1., 1., z, u));
235:   /* workNorm = ||x - z + u||^2 */
236:   PetscCall(VecDot(temp, temp, &workNorm));
237:   /* augment Lagrangian objective (with scaled dual): f(x) + 0.5 * mu ||x -z + u||^2 */
238:   *J = misfit + 0.5 * mu * workNorm;
239:   PetscFunctionReturn(PETSC_SUCCESS);
240: }

242: /* computes FTFx - FTd  mu*(x - z + u) */
243: static PetscErrorCode GradientMisfitADMM(Tao tao, Vec x, Vec V, void *_ctx)
244: {
245:   UserCtx   ctx = (UserCtx)_ctx;
246:   PetscReal mu;
247:   Vec       z, u, temp;

249:   PetscFunctionBegin;
250:   mu   = ctx->mu;
251:   z    = ctx->workRight[5];
252:   u    = ctx->workRight[6];
253:   temp = ctx->workRight[10];
254:   PetscCall(GradientMisfit(tao, x, V, _ctx));
255:   PetscCall(VecCopy(x, temp));
256:   /* temp = x - z + u */
257:   PetscCall(VecAXPBYPCZ(temp, -1., 1., 1., z, u));
258:   /* V =  FTFx - FTd  mu*(x - z + u) */
259:   PetscCall(VecAXPY(V, mu, temp));
260:   PetscFunctionReturn(PETSC_SUCCESS);
261: }

263: /* returns FTF + diag(mu) */
264: static PetscErrorCode HessianMisfitADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
265: {
266:   UserCtx ctx = (UserCtx)_ctx;

268:   PetscFunctionBegin;
269:   PetscCall(MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN));
270:   PetscCall(MatShift(H, ctx->mu));
271:   if (Hpre != H) PetscCall(MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN));
272:   PetscFunctionReturn(PETSC_SUCCESS);
273: }

275: /* computes || x ||_p (mult by 0.5 in case of NORM_2) */
276: static PetscErrorCode ObjectiveRegularization(Tao tao, Vec x, PetscReal *J, void *_ctx)
277: {
278:   UserCtx   ctx = (UserCtx)_ctx;
279:   PetscReal norm;

281:   PetscFunctionBegin;
282:   *J = 0;
283:   PetscCall(VecNorm(x, ctx->p, &norm));
284:   if (ctx->p == NORM_2) norm = 0.5 * norm * norm;
285:   *J = ctx->alpha * norm;
286:   PetscFunctionReturn(PETSC_SUCCESS);
287: }

289: /* NORM_2 Case: return x
290:  * NORM_1 Case: x/(|x| + eps)
291:  * Else: TODO */
292: static PetscErrorCode GradientRegularization(Tao tao, Vec x, Vec V, void *_ctx)
293: {
294:   UserCtx   ctx = (UserCtx)_ctx;
295:   PetscReal eps = ctx->eps;

297:   PetscFunctionBegin;
298:   if (ctx->p == NORM_2) {
299:     PetscCall(VecCopy(x, V));
300:   } else if (ctx->p == NORM_1) {
301:     PetscCall(VecCopy(x, ctx->workRight[1]));
302:     PetscCall(VecAbs(ctx->workRight[1]));
303:     PetscCall(VecShift(ctx->workRight[1], eps));
304:     PetscCall(VecPointwiseDivide(V, x, ctx->workRight[1]));
305:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
306:   PetscFunctionReturn(PETSC_SUCCESS);
307: }

309: /* NORM_2 Case: returns diag(mu)
310:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)))  */
311: static PetscErrorCode HessianRegularization(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
312: {
313:   UserCtx   ctx = (UserCtx)_ctx;
314:   PetscReal eps = ctx->eps;
315:   Vec       copy1, copy2, copy3;

317:   PetscFunctionBegin;
318:   if (ctx->p == NORM_2) {
319:     /* Identity matrix scaled by mu */
320:     PetscCall(MatZeroEntries(H));
321:     PetscCall(MatShift(H, ctx->mu));
322:     if (Hpre != H) {
323:       PetscCall(MatZeroEntries(Hpre));
324:       PetscCall(MatShift(Hpre, ctx->mu));
325:     }
326:   } else if (ctx->p == NORM_1) {
327:     /* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps)) */
328:     copy1 = ctx->workRight[1];
329:     copy2 = ctx->workRight[2];
330:     copy3 = ctx->workRight[3];
331:     /* copy1 : 1/sqrt(x_i^2 + eps) */
332:     PetscCall(VecCopy(x, copy1));
333:     PetscCall(VecPow(copy1, 2));
334:     PetscCall(VecShift(copy1, eps));
335:     PetscCall(VecSqrtAbs(copy1));
336:     PetscCall(VecReciprocal(copy1));
337:     /* copy2:  x_i^2.*/
338:     PetscCall(VecCopy(x, copy2));
339:     PetscCall(VecPow(copy2, 2));
340:     /* copy3: abs(x_i^2 + eps) */
341:     PetscCall(VecCopy(x, copy3));
342:     PetscCall(VecPow(copy3, 2));
343:     PetscCall(VecShift(copy3, eps));
344:     PetscCall(VecAbs(copy3));
345:     /* copy2: 1 - x_i^2/abs(x_i^2 + eps) */
346:     PetscCall(VecPointwiseDivide(copy2, copy2, copy3));
347:     PetscCall(VecScale(copy2, -1.));
348:     PetscCall(VecShift(copy2, 1.));
349:     PetscCall(VecAXPY(copy1, 1., copy2));
350:     PetscCall(VecScale(copy1, ctx->mu));
351:     PetscCall(MatZeroEntries(H));
352:     PetscCall(MatDiagonalSet(H, copy1, INSERT_VALUES));
353:     if (Hpre != H) {
354:       PetscCall(MatZeroEntries(Hpre));
355:       PetscCall(MatDiagonalSet(Hpre, copy1, INSERT_VALUES));
356:     }
357:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
358:   PetscFunctionReturn(PETSC_SUCCESS);
359: }

361: /* NORM_2 Case: 0.5 || x ||_2 + 0.5 * mu * ||x + u - z||^2
362:  * Else : || x ||_2 + 0.5 * mu * ||x + u - z||^2 */
363: static PetscErrorCode ObjectiveRegularizationADMM(Tao tao, Vec z, PetscReal *J, void *_ctx)
364: {
365:   UserCtx   ctx = (UserCtx)_ctx;
366:   PetscReal mu, workNorm, reg;
367:   Vec       x, u, temp;

369:   PetscFunctionBegin;
370:   mu   = ctx->mu;
371:   x    = ctx->workRight[4];
372:   u    = ctx->workRight[6];
373:   temp = ctx->workRight[10];
374:   PetscCall(ObjectiveRegularization(tao, z, &reg, _ctx));
375:   PetscCall(VecCopy(z, temp));
376:   /* temp = x + u -z */
377:   PetscCall(VecAXPBYPCZ(temp, 1., 1., -1., x, u));
378:   /* workNorm = ||x + u - z ||^2 */
379:   PetscCall(VecDot(temp, temp, &workNorm));
380:   *J = reg + 0.5 * mu * workNorm;
381:   PetscFunctionReturn(PETSC_SUCCESS);
382: }

384: /* NORM_2 Case: x - mu*(x + u - z)
385:  * NORM_1 Case: x/(|x| + eps) - mu*(x + u - z)
386:  * Else: TODO */
387: static PetscErrorCode GradientRegularizationADMM(Tao tao, Vec z, Vec V, void *_ctx)
388: {
389:   UserCtx   ctx = (UserCtx)_ctx;
390:   PetscReal mu;
391:   Vec       x, u, temp;

393:   PetscFunctionBegin;
394:   mu   = ctx->mu;
395:   x    = ctx->workRight[4];
396:   u    = ctx->workRight[6];
397:   temp = ctx->workRight[10];
398:   PetscCall(GradientRegularization(tao, z, V, _ctx));
399:   PetscCall(VecCopy(z, temp));
400:   /* temp = x + u -z */
401:   PetscCall(VecAXPBYPCZ(temp, 1., 1., -1., x, u));
402:   PetscCall(VecAXPY(V, -mu, temp));
403:   PetscFunctionReturn(PETSC_SUCCESS);
404: }

406: /* NORM_2 Case: returns diag(mu)
407:  * NORM_1 Case: FTF + diag(mu) */
408: static PetscErrorCode HessianRegularizationADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx)
409: {
410:   UserCtx ctx = (UserCtx)_ctx;

412:   PetscFunctionBegin;
413:   if (ctx->p == NORM_2) {
414:     /* Identity matrix scaled by mu */
415:     PetscCall(MatZeroEntries(H));
416:     PetscCall(MatShift(H, ctx->mu));
417:     if (Hpre != H) {
418:       PetscCall(MatZeroEntries(Hpre));
419:       PetscCall(MatShift(Hpre, ctx->mu));
420:     }
421:   } else if (ctx->p == NORM_1) {
422:     PetscCall(HessianMisfit(tao, x, H, Hpre, (void *)ctx));
423:     PetscCall(MatShift(H, ctx->mu));
424:     if (Hpre != H) PetscCall(MatShift(Hpre, ctx->mu));
425:   } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2");
426:   PetscFunctionReturn(PETSC_SUCCESS);
427: }

429: /* NORM_2 Case : (1/2) * ||F x - d||^2 + 0.5 * || x ||_p
430: *  NORM_1 Case : (1/2) * ||F x - d||^2 + || x ||_p */
431: static PetscErrorCode ObjectiveComplete(Tao tao, Vec x, PetscReal *J, void *ctx)
432: {
433:   PetscReal Jm, Jr;

435:   PetscFunctionBegin;
436:   PetscCall(ObjectiveMisfit(tao, x, &Jm, ctx));
437:   PetscCall(ObjectiveRegularization(tao, x, &Jr, ctx));
438:   *J = Jm + Jr;
439:   PetscFunctionReturn(PETSC_SUCCESS);
440: }

442: /* NORM_2 Case: FTFx - FTd + x
443:  * NORM_1 Case: FTFx - FTd + x/(|x| + eps) */
444: static PetscErrorCode GradientComplete(Tao tao, Vec x, Vec V, void *ctx)
445: {
446:   UserCtx cntx = (UserCtx)ctx;

448:   PetscFunctionBegin;
449:   PetscCall(GradientMisfit(tao, x, cntx->workRight[2], ctx));
450:   PetscCall(GradientRegularization(tao, x, cntx->workRight[3], ctx));
451:   PetscCall(VecWAXPY(V, 1, cntx->workRight[2], cntx->workRight[3]));
452:   PetscFunctionReturn(PETSC_SUCCESS);
453: }

455: /* NORM_2 Case: diag(mu) + FTF
456:  * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) + FTF  */
457: static PetscErrorCode HessianComplete(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx)
458: {
459:   Mat tempH;

461:   PetscFunctionBegin;
462:   PetscCall(MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH));
463:   PetscCall(HessianMisfit(tao, x, H, H, ctx));
464:   PetscCall(HessianRegularization(tao, x, tempH, tempH, ctx));
465:   PetscCall(MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN));
466:   if (Hpre != H) PetscCall(MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN));
467:   PetscCall(MatDestroy(&tempH));
468:   PetscFunctionReturn(PETSC_SUCCESS);
469: }

471: static PetscErrorCode TaoSolveADMM(UserCtx ctx, Vec x)
472: {
473:   PetscInt  i;
474:   PetscReal u_norm, r_norm, s_norm, primal, dual, x_norm, z_norm;
475:   Tao       tao1, tao2;
476:   Vec       xk, z, u, diff, zold, zdiff, temp;
477:   PetscReal mu;

479:   PetscFunctionBegin;
480:   xk    = ctx->workRight[4];
481:   z     = ctx->workRight[5];
482:   u     = ctx->workRight[6];
483:   diff  = ctx->workRight[7];
484:   zold  = ctx->workRight[8];
485:   zdiff = ctx->workRight[9];
486:   temp  = ctx->workRight[11];
487:   mu    = ctx->mu;
488:   PetscCall(VecSet(u, 0.));
489:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao1));
490:   PetscCall(TaoSetType(tao1, TAONLS));
491:   PetscCall(TaoSetObjective(tao1, ObjectiveMisfitADMM, (void *)ctx));
492:   PetscCall(TaoSetGradient(tao1, NULL, GradientMisfitADMM, (void *)ctx));
493:   PetscCall(TaoSetHessian(tao1, ctx->Hm, ctx->Hm, HessianMisfitADMM, (void *)ctx));
494:   PetscCall(VecSet(xk, 0.));
495:   PetscCall(TaoSetSolution(tao1, xk));
496:   PetscCall(TaoSetOptionsPrefix(tao1, "misfit_"));
497:   PetscCall(TaoSetFromOptions(tao1));
498:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao2));
499:   if (ctx->p == NORM_2) {
500:     PetscCall(TaoSetType(tao2, TAONLS));
501:     PetscCall(TaoSetObjective(tao2, ObjectiveRegularizationADMM, (void *)ctx));
502:     PetscCall(TaoSetGradient(tao2, NULL, GradientRegularizationADMM, (void *)ctx));
503:     PetscCall(TaoSetHessian(tao2, ctx->Hr, ctx->Hr, HessianRegularizationADMM, (void *)ctx));
504:   }
505:   PetscCall(VecSet(z, 0.));
506:   PetscCall(TaoSetSolution(tao2, z));
507:   PetscCall(TaoSetOptionsPrefix(tao2, "reg_"));
508:   PetscCall(TaoSetFromOptions(tao2));

510:   for (i = 0; i < ctx->iter; i++) {
511:     PetscCall(VecCopy(z, zold));
512:     PetscCall(TaoSolve(tao1)); /* Updates xk */
513:     if (ctx->p == NORM_1) {
514:       PetscCall(VecWAXPY(temp, 1., xk, u));
515:       PetscCall(TaoSoftThreshold(temp, -ctx->alpha / mu, ctx->alpha / mu, z));
516:     } else {
517:       PetscCall(TaoSolve(tao2)); /* Update zk */
518:     }
519:     /* u = u + xk -z */
520:     PetscCall(VecAXPBYPCZ(u, 1., -1., 1., xk, z));
521:     /* r_norm : norm(x-z) */
522:     PetscCall(VecWAXPY(diff, -1., z, xk));
523:     PetscCall(VecNorm(diff, NORM_2, &r_norm));
524:     /* s_norm : norm(-mu(z-zold)) */
525:     PetscCall(VecWAXPY(zdiff, -1., zold, z));
526:     PetscCall(VecNorm(zdiff, NORM_2, &s_norm));
527:     s_norm = s_norm * mu;
528:     /* primal : sqrt(n)*ABSTOL + RELTOL*max(norm(x), norm(-z))*/
529:     PetscCall(VecNorm(xk, NORM_2, &x_norm));
530:     PetscCall(VecNorm(z, NORM_2, &z_norm));
531:     primal = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * PetscMax(x_norm, z_norm);
532:     /* Duality : sqrt(n)*ABSTOL + RELTOL*norm(mu*u)*/
533:     PetscCall(VecNorm(u, NORM_2, &u_norm));
534:     dual = PetscSqrtReal(ctx->n) * ctx->abstol + ctx->reltol * u_norm * mu;
535:     PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao1), "Iter %" PetscInt_FMT " : ||x-z||: %g, mu*||z-zold||: %g\n", i, (double)r_norm, (double)s_norm));
536:     if (r_norm < primal && s_norm < dual) break;
537:   }
538:   PetscCall(VecCopy(xk, x));
539:   PetscCall(TaoDestroy(&tao1));
540:   PetscCall(TaoDestroy(&tao2));
541:   PetscFunctionReturn(PETSC_SUCCESS);
542: }

544: /* Second order Taylor remainder convergence test */
545: static PetscErrorCode TaylorTest(UserCtx ctx, Tao tao, Vec x, PetscReal *C)
546: {
547:   PetscReal  h, J, temp;
548:   PetscInt   i, j;
549:   PetscInt   numValues;
550:   PetscReal  Jx, Jxhat_comp, Jxhat_pred;
551:   PetscReal *Js, *hs;
552:   PetscReal  gdotdx;
553:   PetscReal  minrate = PETSC_MAX_REAL;
554:   MPI_Comm   comm    = PetscObjectComm((PetscObject)x);
555:   Vec        g, dx, xhat;

557:   PetscFunctionBegin;
558:   PetscCall(VecDuplicate(x, &g));
559:   PetscCall(VecDuplicate(x, &xhat));
560:   /* choose a perturbation direction */
561:   PetscCall(VecDuplicate(x, &dx));
562:   PetscCall(VecSetRandom(dx, ctx->rctx));
563:   /* evaluate objective at x: J(x) */
564:   PetscCall(TaoComputeObjective(tao, x, &Jx));
565:   /* evaluate gradient at x, save in vector g */
566:   PetscCall(TaoComputeGradient(tao, x, g));
567:   PetscCall(VecDot(g, dx, &gdotdx));

569:   for (numValues = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor) numValues++;
570:   PetscCall(PetscCalloc2(numValues, &Js, numValues, &hs));
571:   for (i = 0, h = ctx->hStart; h >= ctx->hMin; h *= ctx->hFactor, i++) {
572:     PetscCall(VecWAXPY(xhat, h, dx, x));
573:     PetscCall(TaoComputeObjective(tao, xhat, &Jxhat_comp));
574:     /* J(\hat(x)) \approx J(x) + g^T (xhat - x) = J(x) + h * g^T dx */
575:     Jxhat_pred = Jx + h * gdotdx;
576:     /* Vector to dJdm scalar? Dot?*/
577:     J = PetscAbsReal(Jxhat_comp - Jxhat_pred);
578:     PetscCall(PetscPrintf(comm, "J(xhat): %g, predicted: %g, diff %g\n", (double)Jxhat_comp, (double)Jxhat_pred, (double)J));
579:     Js[i] = J;
580:     hs[i] = h;
581:   }
582:   for (j = 1; j < numValues; j++) {
583:     temp = PetscLogReal(Js[j] / Js[j - 1]) / PetscLogReal(hs[j] / hs[j - 1]);
584:     PetscCall(PetscPrintf(comm, "Convergence rate step %" PetscInt_FMT ": %g\n", j - 1, (double)temp));
585:     minrate = PetscMin(minrate, temp);
586:   }
587:   /* If O is not ~2, then the test is wrong */
588:   PetscCall(PetscFree2(Js, hs));
589:   *C = minrate;
590:   PetscCall(VecDestroy(&dx));
591:   PetscCall(VecDestroy(&xhat));
592:   PetscCall(VecDestroy(&g));
593:   PetscFunctionReturn(PETSC_SUCCESS);
594: }

596: int main(int argc, char **argv)
597: {
598:   UserCtx ctx;
599:   Tao     tao;
600:   Vec     x;
601:   Mat     H;

603:   PetscFunctionBeginUser;
604:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
605:   PetscCall(PetscNew(&ctx));
606:   PetscCall(ConfigureContext(ctx));
607:   /* Define two functions that could pass as objectives to TaoSetObjective(): one
608:    * for the misfit component, and one for the regularization component */
609:   /* ObjectiveMisfit() and ObjectiveRegularization() */

611:   /* Define a single function that calls both components adds them together: the complete objective,
612:    * in the absence of a Tao implementation that handles separability */
613:   /* ObjectiveComplete() */
614:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
615:   PetscCall(TaoSetType(tao, TAONM));
616:   PetscCall(TaoSetObjective(tao, ObjectiveComplete, (void *)ctx));
617:   PetscCall(TaoSetGradient(tao, NULL, GradientComplete, (void *)ctx));
618:   PetscCall(MatDuplicate(ctx->W, MAT_SHARE_NONZERO_PATTERN, &H));
619:   PetscCall(TaoSetHessian(tao, H, H, HessianComplete, (void *)ctx));
620:   PetscCall(MatCreateVecs(ctx->F, NULL, &x));
621:   PetscCall(VecSet(x, 0.));
622:   PetscCall(TaoSetSolution(tao, x));
623:   PetscCall(TaoSetFromOptions(tao));
624:   if (ctx->use_admm) PetscCall(TaoSolveADMM(ctx, x));
625:   else PetscCall(TaoSolve(tao));
626:   /* examine solution */
627:   PetscCall(VecViewFromOptions(x, NULL, "-view_sol"));
628:   if (ctx->taylor) {
629:     PetscReal rate;
630:     PetscCall(TaylorTest(ctx, tao, x, &rate));
631:   }
632:   if (ctx->soft) { PetscCall(TaoSoftThreshold(x, 0., 0., x)); }
633:   PetscCall(MatDestroy(&H));
634:   PetscCall(TaoDestroy(&tao));
635:   PetscCall(VecDestroy(&x));
636:   PetscCall(DestroyContext(&ctx));
637:   PetscCall(PetscFinalize());
638:   return 0;
639: }

641: /*TEST

643:   build:
644:     requires: !complex

646:   test:
647:     suffix: 0
648:     args:

650:   test:
651:     suffix: l1_1
652:     args: -p 1 -tao_type lmvm -alpha 1. -epsilon 1.e-7 -m 64 -n 64 -view_sol -matrix_format 1

654:   test:
655:     suffix: hessian_1
656:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nls

658:   test:
659:     suffix: hessian_2
660:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nls

662:   test:
663:     suffix: nm_1
664:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nm -tao_max_it 50

666:   test:
667:     suffix: nm_2
668:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nm -tao_max_it 50

670:   test:
671:     suffix: lmvm_1
672:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type lmvm -tao_max_it 40

674:   test:
675:     suffix: lmvm_2
676:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type lmvm -tao_max_it 15

678:   test:
679:     suffix: soft_threshold_admm_1
680:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm

682:   test:
683:     suffix: hessian_admm_1
684:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nls -misfit_tao_type nls

686:   test:
687:     suffix: hessian_admm_2
688:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nls -misfit_tao_type nls

690:   test:
691:     suffix: nm_admm_1
692:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nm -misfit_tao_type nm

694:   test:
695:     suffix: nm_admm_2
696:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nm -misfit_tao_type nm -iter 7

698:   test:
699:     suffix: lmvm_admm_1
700:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm

702:   test:
703:     suffix: lmvm_admm_2
704:     args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm

706:   test:
707:     suffix: soft
708:     args: -taylor 0 -soft 1

710: TEST*/