1 static char help[] = "Simple example to test separable objective optimizers.\n"; 2 3 #include <petsc.h> 4 #include <petsctao.h> 5 #include <petscvec.h> 6 #include <petscmath.h> 7 8 #define NWORKLEFT 4 9 #define NWORKRIGHT 12 10 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; 37 38 static PetscErrorCode CreateRHS(UserCtx ctx) 39 { 40 PetscErrorCode ierr; 41 42 PetscFunctionBegin; 43 /* build the rhs d in ctx */ 44 ierr = VecCreate(PETSC_COMM_WORLD,&(ctx->d));CHKERRQ(ierr); 45 ierr = VecSetSizes(ctx->d,PETSC_DECIDE,ctx->m);CHKERRQ(ierr); 46 ierr = VecSetFromOptions(ctx->d);CHKERRQ(ierr); 47 ierr = VecSetRandom(ctx->d,ctx->rctx);CHKERRQ(ierr); 48 PetscFunctionReturn(0); 49 } 50 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 57 PetscErrorCode ierr; 58 59 PetscFunctionBegin; 60 /* build the matrix F in ctx */ 61 ierr = MatCreate(PETSC_COMM_WORLD, &(ctx->F));CHKERRQ(ierr); 62 ierr = MatSetSizes(ctx->F,PETSC_DECIDE, PETSC_DECIDE, ctx->m, ctx->n);CHKERRQ(ierr); 63 ierr = MatSetType(ctx->F,MATAIJ);CHKERRQ(ierr); /* TODO: Decide specific SetType other than dummy*/ 64 ierr = MatMPIAIJSetPreallocation(ctx->F, 5, NULL, 5, NULL);CHKERRQ(ierr); /*TODO: some number other than 5?*/ 65 ierr = MatSeqAIJSetPreallocation(ctx->F, 5, NULL);CHKERRQ(ierr); 66 ierr = MatSetUp(ctx->F);CHKERRQ(ierr); 67 ierr = MatGetOwnershipRange(ctx->F,&Istart,&Iend);CHKERRQ(ierr); 68 ierr = PetscLogStageRegister("Assembly", &stage);CHKERRQ(ierr); 69 ierr = PetscLogStagePush(stage);CHKERRQ(ierr); 70 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 ierr = MatSetValue(ctx->F, Ii, Ii, 4., INSERT_VALUES);CHKERRQ(ierr); 87 ierr = MatSetValue(ctx->F, Ii, I_n, -1., INSERT_VALUES);CHKERRQ(ierr); 88 ierr = MatSetValue(ctx->F, Ii, I_s, -1., INSERT_VALUES);CHKERRQ(ierr); 89 ierr = MatSetValue(ctx->F, Ii, I_e, -1., INSERT_VALUES);CHKERRQ(ierr); 90 ierr = MatSetValue(ctx->F, Ii, I_w, -1., INSERT_VALUES);CHKERRQ(ierr); 91 } 92 } else {ierr = MatSetRandom(ctx->F, ctx->rctx);CHKERRQ(ierr);} 93 ierr = MatAssemblyBegin(ctx->F, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 94 ierr = MatAssemblyEnd(ctx->F, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 95 ierr = PetscLogStagePop();CHKERRQ(ierr); 96 /* Stencil matrix is symmetric. Setting symmetric flag for ICC/Cholesky preconditioner */ 97 if (!(ctx->matops)) { 98 ierr = MatSetOption(ctx->F,MAT_SYMMETRIC,PETSC_TRUE);CHKERRQ(ierr); 99 } 100 ierr = MatTransposeMatMult(ctx->F,ctx->F, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &(ctx->W));CHKERRQ(ierr); 101 /* Setup Hessian Workspace in same shape as W */ 102 ierr = MatDuplicate(ctx->W,MAT_DO_NOT_COPY_VALUES,&(ctx->Hm));CHKERRQ(ierr); 103 ierr = MatDuplicate(ctx->W,MAT_DO_NOT_COPY_VALUES,&(ctx->Hr));CHKERRQ(ierr); 104 PetscFunctionReturn(0); 105 } 106 107 static PetscErrorCode SetupWorkspace(UserCtx ctx) 108 { 109 PetscInt i; 110 PetscErrorCode ierr; 111 112 PetscFunctionBegin; 113 ierr = MatCreateVecs(ctx->F, &ctx->workLeft[0], &ctx->workRight[0]);CHKERRQ(ierr); 114 for (i=1; i<NWORKLEFT; i++) { 115 ierr = VecDuplicate(ctx->workLeft[0], &(ctx->workLeft[i]));CHKERRQ(ierr); 116 } 117 for (i=1; i<NWORKRIGHT; i++) { 118 ierr = VecDuplicate(ctx->workRight[0], &(ctx->workRight[i]));CHKERRQ(ierr); 119 } 120 PetscFunctionReturn(0); 121 } 122 123 static PetscErrorCode ConfigureContext(UserCtx ctx) 124 { 125 PetscErrorCode ierr; 126 127 PetscFunctionBegin; 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 ierr = PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Configure separable objection example", "ex4.c");CHKERRQ(ierr); 144 ierr = PetscOptionsInt("-m", "The row dimension of matrix F", "ex4.c", ctx->m, &(ctx->m), NULL);CHKERRQ(ierr); 145 ierr = PetscOptionsInt("-n", "The column dimension of matrix F", "ex4.c", ctx->n, &(ctx->n), NULL);CHKERRQ(ierr); 146 ierr = PetscOptionsInt("-matrix_format","Decide format of F matrix. 0 for stencil, 1 for random", "ex4.c", ctx->matops, &(ctx->matops), NULL);CHKERRQ(ierr); 147 ierr = PetscOptionsInt("-iter","Iteration number ADMM", "ex4.c", ctx->iter, &(ctx->iter), NULL);CHKERRQ(ierr); 148 ierr = PetscOptionsReal("-alpha", "The regularization multiplier. 1 default", "ex4.c", ctx->alpha, &(ctx->alpha), NULL);CHKERRQ(ierr); 149 ierr = 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);CHKERRQ(ierr); 150 ierr = PetscOptionsReal("-mu", "The augmented lagrangian multiplier in ADMM", "ex4.c", ctx->mu, &(ctx->mu), NULL);CHKERRQ(ierr); 151 ierr = PetscOptionsReal("-hStart", "Taylor test starting point. 1 default.", "ex4.c", ctx->hStart, &(ctx->hStart), NULL);CHKERRQ(ierr); 152 ierr = PetscOptionsReal("-hFactor", "Taylor test multiplier factor. 0.5 default", "ex4.c", ctx->hFactor, &(ctx->hFactor), NULL);CHKERRQ(ierr); 153 ierr = PetscOptionsReal("-hMin", "Taylor test ending condition. 1.e-3 default", "ex4.c", ctx->hMin, &(ctx->hMin), NULL);CHKERRQ(ierr); 154 ierr = PetscOptionsReal("-abstol", "Absolute stopping criterion for ADMM", "ex4.c", ctx->abstol, &(ctx->abstol), NULL);CHKERRQ(ierr); 155 ierr = PetscOptionsReal("-reltol", "Relative stopping criterion for ADMM", "ex4.c", ctx->reltol, &(ctx->reltol), NULL);CHKERRQ(ierr); 156 ierr = PetscOptionsBool("-taylor","Flag for Taylor test. Default is true.", "ex4.c", ctx->taylor, &(ctx->taylor), NULL);CHKERRQ(ierr); 157 ierr = PetscOptionsBool("-use_admm","Use the ADMM solver in this example.", "ex4.c", ctx->use_admm, &(ctx->use_admm), NULL);CHKERRQ(ierr); 158 ierr = PetscOptionsEnum("-p","Norm type.", "ex4.c", NormTypes, (PetscEnum)ctx->p, (PetscEnum *) &(ctx->p), NULL);CHKERRQ(ierr); 159 ierr = PetscOptionsEnd();CHKERRQ(ierr); 160 /* Creating random ctx */ 161 ierr = PetscRandomCreate(PETSC_COMM_WORLD,&(ctx->rctx));CHKERRQ(ierr); 162 ierr = PetscRandomSetFromOptions(ctx->rctx);CHKERRQ(ierr); 163 ierr = CreateMatrix(ctx);CHKERRQ(ierr); 164 ierr = CreateRHS(ctx);CHKERRQ(ierr); 165 ierr = SetupWorkspace(ctx);CHKERRQ(ierr); 166 PetscFunctionReturn(0); 167 } 168 169 static PetscErrorCode DestroyContext(UserCtx *ctx) 170 { 171 PetscInt i; 172 PetscErrorCode ierr; 173 174 PetscFunctionBegin; 175 ierr = MatDestroy(&((*ctx)->F));CHKERRQ(ierr); 176 ierr = MatDestroy(&((*ctx)->W));CHKERRQ(ierr); 177 ierr = MatDestroy(&((*ctx)->Hm));CHKERRQ(ierr); 178 ierr = MatDestroy(&((*ctx)->Hr));CHKERRQ(ierr); 179 ierr = VecDestroy(&((*ctx)->d));CHKERRQ(ierr); 180 for (i=0; i<NWORKLEFT; i++) { 181 ierr = VecDestroy(&((*ctx)->workLeft[i]));CHKERRQ(ierr); 182 } 183 for (i=0; i<NWORKRIGHT; i++) { 184 ierr = VecDestroy(&((*ctx)->workRight[i]));CHKERRQ(ierr); 185 } 186 ierr = PetscRandomDestroy(&((*ctx)->rctx));CHKERRQ(ierr); 187 ierr = PetscFree(*ctx);CHKERRQ(ierr); 188 PetscFunctionReturn(0); 189 } 190 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; 195 PetscErrorCode ierr; 196 Vec y; 197 198 PetscFunctionBegin; 199 y = ctx->workLeft[0]; 200 ierr = MatMult(ctx->F, x, y);CHKERRQ(ierr); 201 ierr = VecAXPY(y, -1., ctx->d);CHKERRQ(ierr); 202 ierr = VecDot(y, y, J);CHKERRQ(ierr); 203 *J *= 0.5; 204 PetscFunctionReturn(0); 205 } 206 207 /* compute V = FTFx - FTd */ 208 static PetscErrorCode GradientMisfit(Tao tao, Vec x, Vec V, void *_ctx) 209 { 210 UserCtx ctx = (UserCtx) _ctx; 211 PetscErrorCode ierr; 212 Vec FTFx, FTd; 213 214 PetscFunctionBegin; 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 ierr = MatMult(ctx->W,x,FTFx);CHKERRQ(ierr); 219 ierr = MatMultTranspose(ctx->F, ctx->d, FTd);CHKERRQ(ierr); 220 ierr = VecWAXPY(V, -1., FTd, FTFx);CHKERRQ(ierr); 221 PetscFunctionReturn(0); 222 } 223 224 /* returns FTF */ 225 static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx) 226 { 227 UserCtx ctx = (UserCtx) _ctx; 228 PetscErrorCode ierr; 229 230 PetscFunctionBegin; 231 if (H != ctx->W) {ierr = MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr);} 232 if (Hpre != ctx->W) {ierr = MatCopy(ctx->W, Hpre, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr);} 233 PetscFunctionReturn(0); 234 } 235 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; 243 PetscErrorCode ierr; 244 245 PetscFunctionBegin; 246 mu = ctx->mu; 247 z = ctx->workRight[5]; 248 u = ctx->workRight[6]; 249 temp = ctx->workRight[10]; 250 /* misfit = f(x) */ 251 ierr = ObjectiveMisfit(tao, x, &misfit, _ctx);CHKERRQ(ierr); 252 ierr = VecCopy(x,temp);CHKERRQ(ierr); 253 /* temp = x - z + u */ 254 ierr = VecAXPBYPCZ(temp,-1.,1.,1.,z,u);CHKERRQ(ierr); 255 /* workNorm = ||x - z + u||^2 */ 256 ierr = VecDot(temp, temp, &workNorm);CHKERRQ(ierr); 257 /* augment Lagrangian objective (with scaled dual): f(x) + 0.5 * mu ||x -z + u||^2 */ 258 *J = misfit + 0.5 * mu * workNorm; 259 PetscFunctionReturn(0); 260 } 261 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; 268 PetscErrorCode ierr; 269 270 PetscFunctionBegin; 271 mu = ctx->mu; 272 z = ctx->workRight[5]; 273 u = ctx->workRight[6]; 274 temp = ctx->workRight[10]; 275 ierr = GradientMisfit(tao, x, V, _ctx);CHKERRQ(ierr); 276 ierr = VecCopy(x, temp);CHKERRQ(ierr); 277 /* temp = x - z + u */ 278 ierr = VecAXPBYPCZ(temp,-1.,1.,1.,z,u);CHKERRQ(ierr); 279 /* V = FTFx - FTd mu*(x - z + u) */ 280 ierr = VecAXPY(V, mu, temp);CHKERRQ(ierr); 281 PetscFunctionReturn(0); 282 } 283 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; 288 PetscErrorCode ierr; 289 290 PetscFunctionBegin; 291 ierr = MatCopy(ctx->W, H, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr); 292 ierr = MatShift(H, ctx->mu);CHKERRQ(ierr); 293 if (Hpre != H) { 294 ierr = MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr); 295 } 296 PetscFunctionReturn(0); 297 } 298 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; 304 PetscErrorCode ierr; 305 306 PetscFunctionBegin; 307 *J = 0; 308 ierr = VecNorm (x, ctx->p, &norm);CHKERRQ(ierr); 309 if (ctx->p == NORM_2) norm = 0.5 * norm * norm; 310 *J = ctx->alpha * norm; 311 PetscFunctionReturn(0); 312 } 313 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; 320 PetscErrorCode ierr; 321 PetscReal eps = ctx->eps; 322 323 PetscFunctionBegin; 324 if (ctx->p == NORM_2) { 325 ierr = VecCopy(x, V);CHKERRQ(ierr); 326 } else if (ctx->p == NORM_1) { 327 ierr = VecCopy(x, ctx->workRight[1]);CHKERRQ(ierr); 328 ierr = VecAbs(ctx->workRight[1]);CHKERRQ(ierr); 329 ierr = VecShift(ctx->workRight[1], eps);CHKERRQ(ierr); 330 ierr = VecPointwiseDivide(V, x, ctx->workRight[1]);CHKERRQ(ierr); 331 } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2"); 332 PetscFunctionReturn(0); 333 } 334 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; 342 PetscErrorCode ierr; 343 344 PetscFunctionBegin; 345 if (ctx->p == NORM_2) { 346 /* Identity matrix scaled by mu */ 347 ierr = MatZeroEntries(H);CHKERRQ(ierr); 348 ierr = MatShift(H,ctx->mu);CHKERRQ(ierr); 349 if (Hpre != H) { 350 ierr = MatZeroEntries(Hpre);CHKERRQ(ierr); 351 ierr = MatShift(Hpre,ctx->mu);CHKERRQ(ierr); 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 ierr = VecCopy(x, copy1);CHKERRQ(ierr); 360 ierr = VecPow(copy1,2);CHKERRQ(ierr); 361 ierr = VecShift(copy1, eps);CHKERRQ(ierr); 362 ierr = VecSqrtAbs(copy1);CHKERRQ(ierr); 363 ierr = VecReciprocal(copy1);CHKERRQ(ierr); 364 /* copy2: x_i^2.*/ 365 ierr = VecCopy(x,copy2);CHKERRQ(ierr); 366 ierr = VecPow(copy2,2);CHKERRQ(ierr); 367 /* copy3: abs(x_i^2 + eps) */ 368 ierr = VecCopy(x,copy3);CHKERRQ(ierr); 369 ierr = VecPow(copy3,2);CHKERRQ(ierr); 370 ierr = VecShift(copy3, eps);CHKERRQ(ierr); 371 ierr = VecAbs(copy3);CHKERRQ(ierr); 372 /* copy2: 1 - x_i^2/abs(x_i^2 + eps) */ 373 ierr = VecPointwiseDivide(copy2, copy2,copy3);CHKERRQ(ierr); 374 ierr = VecScale(copy2, -1.);CHKERRQ(ierr); 375 ierr = VecShift(copy2, 1.);CHKERRQ(ierr); 376 ierr = VecAXPY(copy1,1.,copy2);CHKERRQ(ierr); 377 ierr = VecScale(copy1, ctx->mu);CHKERRQ(ierr); 378 ierr = MatZeroEntries(H);CHKERRQ(ierr); 379 ierr = MatDiagonalSet(H, copy1,INSERT_VALUES);CHKERRQ(ierr); 380 if (Hpre != H) { 381 ierr = MatZeroEntries(Hpre);CHKERRQ(ierr); 382 ierr = MatDiagonalSet(Hpre, copy1,INSERT_VALUES);CHKERRQ(ierr); 383 } 384 } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2"); 385 PetscFunctionReturn(0); 386 } 387 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; 395 PetscErrorCode ierr; 396 397 PetscFunctionBegin; 398 mu = ctx->mu; 399 x = ctx->workRight[4]; 400 u = ctx->workRight[6]; 401 temp = ctx->workRight[10]; 402 ierr = ObjectiveRegularization(tao, z, ®, _ctx);CHKERRQ(ierr); 403 ierr = VecCopy(z,temp);CHKERRQ(ierr); 404 /* temp = x + u -z */ 405 ierr = VecAXPBYPCZ(temp,1.,1.,-1.,x,u);CHKERRQ(ierr); 406 /* workNorm = ||x + u - z ||^2 */ 407 ierr = VecDot(temp, temp, &workNorm);CHKERRQ(ierr); 408 *J = reg + 0.5 * mu * workNorm; 409 PetscFunctionReturn(0); 410 } 411 412 /* NORM_2 Case: x - mu*(x + u - z) 413 * NORM_1 Case: x/(|x| + eps) - mu*(x + u - z) 414 * Else: TODO */ 415 static PetscErrorCode GradientRegularizationADMM(Tao tao, Vec z, Vec V, void *_ctx) 416 { 417 UserCtx ctx = (UserCtx) _ctx; 418 PetscReal mu; 419 Vec x, u, temp; 420 PetscErrorCode ierr; 421 422 PetscFunctionBegin; 423 mu = ctx->mu; 424 x = ctx->workRight[4]; 425 u = ctx->workRight[6]; 426 temp = ctx->workRight[10]; 427 ierr = GradientRegularization(tao, z, V, _ctx);CHKERRQ(ierr); 428 ierr = VecCopy(z, temp);CHKERRQ(ierr); 429 /* temp = x + u -z */ 430 ierr = VecAXPBYPCZ(temp,1.,1.,-1.,x,u);CHKERRQ(ierr); 431 ierr = VecAXPY(V, -mu, temp);CHKERRQ(ierr); 432 PetscFunctionReturn(0); 433 } 434 435 /* NORM_2 Case: returns diag(mu) 436 * NORM_1 Case: FTF + diag(mu) */ 437 static PetscErrorCode HessianRegularizationADMM(Tao tao, Vec x, Mat H, Mat Hpre, void *_ctx) 438 { 439 UserCtx ctx = (UserCtx) _ctx; 440 PetscErrorCode ierr; 441 442 PetscFunctionBegin; 443 if (ctx->p == NORM_2) { 444 /* Identity matrix scaled by mu */ 445 ierr = MatZeroEntries(H);CHKERRQ(ierr); 446 ierr = MatShift(H,ctx->mu);CHKERRQ(ierr); 447 if (Hpre != H) { 448 ierr = MatZeroEntries(Hpre);CHKERRQ(ierr); 449 ierr = MatShift(Hpre,ctx->mu);CHKERRQ(ierr); 450 } 451 } else if (ctx->p == NORM_1) { 452 ierr = HessianMisfit(tao, x, H, Hpre, (void*) ctx);CHKERRQ(ierr); 453 ierr = MatShift(H, ctx->mu);CHKERRQ(ierr); 454 if (Hpre != H) {ierr = MatShift(Hpre, ctx->mu);CHKERRQ(ierr);} 455 } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2"); 456 PetscFunctionReturn(0); 457 } 458 459 /* NORM_2 Case : (1/2) * ||F x - d||^2 + 0.5 * || x ||_p 460 * NORM_1 Case : (1/2) * ||F x - d||^2 + || x ||_p */ 461 static PetscErrorCode ObjectiveComplete(Tao tao, Vec x, PetscReal *J, void *ctx) 462 { 463 PetscReal Jm, Jr; 464 PetscErrorCode ierr; 465 466 PetscFunctionBegin; 467 ierr = ObjectiveMisfit(tao, x, &Jm, ctx);CHKERRQ(ierr); 468 ierr = ObjectiveRegularization(tao, x, &Jr, ctx);CHKERRQ(ierr); 469 *J = Jm + Jr; 470 PetscFunctionReturn(0); 471 } 472 473 /* NORM_2 Case: FTFx - FTd + x 474 * NORM_1 Case: FTFx - FTd + x/(|x| + eps) */ 475 static PetscErrorCode GradientComplete(Tao tao, Vec x, Vec V, void *ctx) 476 { 477 UserCtx cntx = (UserCtx) ctx; 478 PetscErrorCode ierr; 479 480 PetscFunctionBegin; 481 ierr = GradientMisfit(tao, x, cntx->workRight[2], ctx);CHKERRQ(ierr); 482 ierr = GradientRegularization(tao, x, cntx->workRight[3], ctx);CHKERRQ(ierr); 483 ierr = VecWAXPY(V,1,cntx->workRight[2],cntx->workRight[3]);CHKERRQ(ierr); 484 PetscFunctionReturn(0); 485 } 486 487 /* NORM_2 Case: diag(mu) + FTF 488 * NORM_1 Case: diag(mu* 1/sqrt(x_i^2 + eps) * (1 - x_i^2/ABS(x_i^2+eps))) + FTF */ 489 static PetscErrorCode HessianComplete(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx) 490 { 491 Mat tempH; 492 PetscErrorCode ierr; 493 494 PetscFunctionBegin; 495 ierr = MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH);CHKERRQ(ierr); 496 ierr = HessianMisfit(tao, x, H, H, ctx);CHKERRQ(ierr); 497 ierr = HessianRegularization(tao, x, tempH, tempH, ctx);CHKERRQ(ierr); 498 ierr = MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr); 499 if (Hpre != H) { 500 ierr = MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr); 501 } 502 ierr = MatDestroy(&tempH);CHKERRQ(ierr); 503 PetscFunctionReturn(0); 504 } 505 506 static PetscErrorCode TaoSolveADMM(UserCtx ctx, Vec x) 507 { 508 PetscErrorCode ierr; 509 PetscInt i; 510 PetscReal u_norm, r_norm, s_norm, primal, dual, x_norm, z_norm; 511 Tao tao1,tao2; 512 Vec xk,z,u,diff,zold,zdiff,temp; 513 PetscReal mu; 514 515 PetscFunctionBegin; 516 xk = ctx->workRight[4]; 517 z = ctx->workRight[5]; 518 u = ctx->workRight[6]; 519 diff = ctx->workRight[7]; 520 zold = ctx->workRight[8]; 521 zdiff = ctx->workRight[9]; 522 temp = ctx->workRight[11]; 523 mu = ctx->mu; 524 ierr = VecSet(u, 0.);CHKERRQ(ierr); 525 ierr = TaoCreate(PETSC_COMM_WORLD, &tao1);CHKERRQ(ierr); 526 ierr = TaoSetType(tao1,TAONLS);CHKERRQ(ierr); 527 ierr = TaoSetObjectiveRoutine(tao1, ObjectiveMisfitADMM, (void*) ctx);CHKERRQ(ierr); 528 ierr = TaoSetGradientRoutine(tao1, GradientMisfitADMM, (void*) ctx);CHKERRQ(ierr); 529 ierr = TaoSetHessianRoutine(tao1, ctx->Hm, ctx->Hm, HessianMisfitADMM, (void*) ctx);CHKERRQ(ierr); 530 ierr = VecSet(xk, 0.);CHKERRQ(ierr); 531 ierr = TaoSetInitialVector(tao1, xk);CHKERRQ(ierr); 532 ierr = TaoSetOptionsPrefix(tao1, "misfit_");CHKERRQ(ierr); 533 ierr = TaoSetFromOptions(tao1);CHKERRQ(ierr); 534 ierr = TaoCreate(PETSC_COMM_WORLD, &tao2);CHKERRQ(ierr); 535 if (ctx->p == NORM_2) { 536 ierr = TaoSetType(tao2,TAONLS);CHKERRQ(ierr); 537 ierr = TaoSetObjectiveRoutine(tao2, ObjectiveRegularizationADMM, (void*) ctx);CHKERRQ(ierr); 538 ierr = TaoSetGradientRoutine(tao2, GradientRegularizationADMM, (void*) ctx);CHKERRQ(ierr); 539 ierr = TaoSetHessianRoutine(tao2, ctx->Hr, ctx->Hr, HessianRegularizationADMM, (void*) ctx);CHKERRQ(ierr); 540 } 541 ierr = VecSet(z, 0.);CHKERRQ(ierr); 542 ierr = TaoSetInitialVector(tao2, z);CHKERRQ(ierr); 543 ierr = TaoSetOptionsPrefix(tao2, "reg_");CHKERRQ(ierr); 544 ierr = TaoSetFromOptions(tao2);CHKERRQ(ierr); 545 546 for (i=0; i<ctx->iter; i++) { 547 ierr = VecCopy(z,zold);CHKERRQ(ierr); 548 ierr = TaoSolve(tao1);CHKERRQ(ierr); /* Updates xk */ 549 if (ctx->p == NORM_1) { 550 ierr = VecWAXPY(temp,1.,xk,u);CHKERRQ(ierr); 551 ierr = TaoSoftThreshold(temp,-ctx->alpha/mu,ctx->alpha/mu,z);CHKERRQ(ierr); 552 } else { 553 ierr = TaoSolve(tao2);CHKERRQ(ierr); /* Update zk */ 554 } 555 /* u = u + xk -z */ 556 ierr = VecAXPBYPCZ(u,1.,-1.,1.,xk,z);CHKERRQ(ierr); 557 /* r_norm : norm(x-z) */ 558 ierr = VecWAXPY(diff,-1.,z,xk);CHKERRQ(ierr); 559 ierr = VecNorm(diff,NORM_2,&r_norm);CHKERRQ(ierr); 560 /* s_norm : norm(-mu(z-zold)) */ 561 ierr = VecWAXPY(zdiff, -1.,zold,z);CHKERRQ(ierr); 562 ierr = VecNorm(zdiff,NORM_2,&s_norm);CHKERRQ(ierr); 563 s_norm = s_norm * mu; 564 /* primal : sqrt(n)*ABSTOL + RELTOL*max(norm(x), norm(-z))*/ 565 ierr = VecNorm(xk,NORM_2,&x_norm);CHKERRQ(ierr); 566 ierr = VecNorm(z,NORM_2,&z_norm);CHKERRQ(ierr); 567 primal = PetscSqrtReal(ctx->n)*ctx->abstol + ctx->reltol*PetscMax(x_norm,z_norm); 568 /* Duality : sqrt(n)*ABSTOL + RELTOL*norm(mu*u)*/ 569 ierr = VecNorm(u,NORM_2,&u_norm);CHKERRQ(ierr); 570 dual = PetscSqrtReal(ctx->n)*ctx->abstol + ctx->reltol*u_norm*mu; 571 ierr = PetscPrintf(PetscObjectComm((PetscObject)tao1),"Iter %D : ||x-z||: %g, mu*||z-zold||: %g\n", i, (double) r_norm, (double) s_norm);CHKERRQ(ierr); 572 if (r_norm < primal && s_norm < dual) break; 573 } 574 ierr = VecCopy(xk, x);CHKERRQ(ierr); 575 ierr = TaoDestroy(&tao1);CHKERRQ(ierr); 576 ierr = TaoDestroy(&tao2);CHKERRQ(ierr); 577 PetscFunctionReturn(0); 578 } 579 580 /* Second order Taylor remainder convergence test */ 581 static PetscErrorCode TaylorTest(UserCtx ctx, Tao tao, Vec x, PetscReal *C) 582 { 583 PetscReal h,J,temp; 584 PetscInt i,j; 585 PetscInt numValues; 586 PetscReal Jx,Jxhat_comp,Jxhat_pred; 587 PetscReal *Js, *hs; 588 PetscReal gdotdx; 589 PetscReal minrate = PETSC_MAX_REAL; 590 MPI_Comm comm = PetscObjectComm((PetscObject)x); 591 Vec g, dx, xhat; 592 PetscErrorCode ierr; 593 594 PetscFunctionBegin; 595 ierr = VecDuplicate(x, &g);CHKERRQ(ierr); 596 ierr = VecDuplicate(x, &xhat);CHKERRQ(ierr); 597 /* choose a perturbation direction */ 598 ierr = VecDuplicate(x, &dx);CHKERRQ(ierr); 599 ierr = VecSetRandom(dx,ctx->rctx);CHKERRQ(ierr); 600 /* evaluate objective at x: J(x) */ 601 ierr = TaoComputeObjective(tao, x, &Jx);CHKERRQ(ierr); 602 /* evaluate gradient at x, save in vector g */ 603 ierr = TaoComputeGradient(tao, x, g);CHKERRQ(ierr); 604 ierr = VecDot(g, dx, &gdotdx);CHKERRQ(ierr); 605 606 for (numValues=0, h=ctx->hStart; h>=ctx->hMin; h*=ctx->hFactor) numValues++; 607 ierr = PetscCalloc2(numValues, &Js, numValues, &hs);CHKERRQ(ierr); 608 for (i=0, h=ctx->hStart; h>=ctx->hMin; h*=ctx->hFactor, i++) { 609 ierr = VecWAXPY(xhat, h, dx, x);CHKERRQ(ierr); 610 ierr = TaoComputeObjective(tao, xhat, &Jxhat_comp);CHKERRQ(ierr); 611 /* J(\hat(x)) \approx J(x) + g^T (xhat - x) = J(x) + h * g^T dx */ 612 Jxhat_pred = Jx + h * gdotdx; 613 /* Vector to dJdm scalar? Dot?*/ 614 J = PetscAbsReal(Jxhat_comp - Jxhat_pred); 615 ierr = PetscPrintf (comm, "J(xhat): %g, predicted: %g, diff %g\n", (double) Jxhat_comp,(double) Jxhat_pred, (double) J);CHKERRQ(ierr); 616 Js[i] = J; 617 hs[i] = h; 618 } 619 for (j=1; j<numValues; j++) { 620 temp = PetscLogReal(Js[j]/Js[j-1]) / PetscLogReal (hs[j]/hs[j-1]); 621 ierr = PetscPrintf (comm, "Convergence rate step %D: %g\n", j-1, (double) temp);CHKERRQ(ierr); 622 minrate = PetscMin(minrate, temp); 623 } 624 /* If O is not ~2, then the test is wrong */ 625 ierr = PetscFree2(Js, hs);CHKERRQ(ierr); 626 *C = minrate; 627 ierr = VecDestroy(&dx);CHKERRQ(ierr); 628 ierr = VecDestroy(&xhat);CHKERRQ(ierr); 629 ierr = VecDestroy(&g);CHKERRQ(ierr); 630 PetscFunctionReturn(0); 631 } 632 633 int main(int argc, char ** argv) 634 { 635 UserCtx ctx; 636 Tao tao; 637 Vec x; 638 Mat H; 639 PetscErrorCode ierr; 640 641 ierr = PetscInitialize(&argc, &argv, NULL,help);if (ierr) return ierr; 642 ierr = PetscNew(&ctx);CHKERRQ(ierr); 643 ierr = ConfigureContext(ctx);CHKERRQ(ierr); 644 /* Define two functions that could pass as objectives to TaoSetObjectiveRoutine(): one 645 * for the misfit component, and one for the regularization component */ 646 /* ObjectiveMisfit() and ObjectiveRegularization() */ 647 648 /* Define a single function that calls both components adds them together: the complete objective, 649 * in the absence of a Tao implementation that handles separability */ 650 /* ObjectiveComplete() */ 651 ierr = TaoCreate(PETSC_COMM_WORLD, &tao);CHKERRQ(ierr); 652 ierr = TaoSetType(tao,TAONM);CHKERRQ(ierr); 653 ierr = TaoSetObjectiveRoutine(tao, ObjectiveComplete, (void*) ctx);CHKERRQ(ierr); 654 ierr = TaoSetGradientRoutine(tao, GradientComplete, (void*) ctx);CHKERRQ(ierr); 655 ierr = MatDuplicate(ctx->W, MAT_SHARE_NONZERO_PATTERN, &H);CHKERRQ(ierr); 656 ierr = TaoSetHessianRoutine(tao, H, H, HessianComplete, (void*) ctx);CHKERRQ(ierr); 657 ierr = MatCreateVecs(ctx->F, NULL, &x);CHKERRQ(ierr); 658 ierr = VecSet(x, 0.);CHKERRQ(ierr); 659 ierr = TaoSetInitialVector(tao, x);CHKERRQ(ierr); 660 ierr = TaoSetFromOptions(tao);CHKERRQ(ierr); 661 if (ctx->use_admm) { 662 ierr = TaoSolveADMM(ctx,x);CHKERRQ(ierr); 663 } else {ierr = TaoSolve(tao);CHKERRQ(ierr);} 664 /* examine solution */ 665 ierr = VecViewFromOptions(x, NULL, "-view_sol");CHKERRQ(ierr); 666 if (ctx->taylor) { 667 PetscReal rate; 668 ierr = TaylorTest(ctx, tao, x, &rate);CHKERRQ(ierr); 669 } 670 ierr = MatDestroy(&H);CHKERRQ(ierr); 671 ierr = TaoDestroy(&tao);CHKERRQ(ierr); 672 ierr = VecDestroy(&x);CHKERRQ(ierr); 673 ierr = DestroyContext(&ctx);CHKERRQ(ierr); 674 ierr = PetscFinalize();CHKERRQ(ierr); 675 return ierr; 676 } 677 678 /*TEST 679 680 build: 681 requires: !complex 682 683 test: 684 suffix: 0 685 args: 686 687 test: 688 suffix: l1_1 689 args: -p 1 -tao_type lmvm -alpha 1. -epsilon 1.e-7 -m 64 -n 64 -view_sol -matrix_format 1 690 691 test: 692 suffix: hessian_1 693 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nls -tao_nls_ksp_monitor 694 695 test: 696 suffix: hessian_2 697 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nls -tao_nls_ksp_monitor 698 699 test: 700 suffix: nm_1 701 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type nm -tao_max_it 50 702 703 test: 704 suffix: nm_2 705 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type nm -tao_max_it 50 706 707 test: 708 suffix: lmvm_1 709 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -tao_type lmvm -tao_max_it 40 710 711 test: 712 suffix: lmvm_2 713 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -tao_type lmvm -tao_max_it 15 714 715 test: 716 suffix: soft_threshold_admm_1 717 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm 718 719 test: 720 suffix: hessian_admm_1 721 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nls -misfit_tao_type nls 722 723 test: 724 suffix: hessian_admm_2 725 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nls -misfit_tao_type nls 726 727 test: 728 suffix: nm_admm_1 729 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type nm -misfit_tao_type nm 730 731 test: 732 suffix: nm_admm_2 733 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type nm -misfit_tao_type nm -iter 7 734 735 test: 736 suffix: lmvm_admm_1 737 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm 738 739 test: 740 suffix: lmvm_admm_2 741 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 2 -use_admm -reg_tao_type lmvm -misfit_tao_type lmvm 742 743 TEST*/ 744