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 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; 421 PetscErrorCode ierr; 422 423 PetscFunctionBegin; 424 mu = ctx->mu; 425 x = ctx->workRight[4]; 426 u = ctx->workRight[6]; 427 temp = ctx->workRight[10]; 428 ierr = GradientRegularization(tao, z, V, _ctx);CHKERRQ(ierr); 429 ierr = VecCopy(z, temp);CHKERRQ(ierr); 430 /* temp = x + u -z */ 431 ierr = VecAXPBYPCZ(temp,1.,1.,-1.,x,u);CHKERRQ(ierr); 432 ierr = VecAXPY(V, -mu, temp);CHKERRQ(ierr); 433 PetscFunctionReturn(0); 434 } 435 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; 441 PetscErrorCode ierr; 442 443 PetscFunctionBegin; 444 if (ctx->p == NORM_2) { 445 /* Identity matrix scaled by mu */ 446 ierr = MatZeroEntries(H);CHKERRQ(ierr); 447 ierr = MatShift(H,ctx->mu);CHKERRQ(ierr); 448 if (Hpre != H) { 449 ierr = MatZeroEntries(Hpre);CHKERRQ(ierr); 450 ierr = MatShift(Hpre,ctx->mu);CHKERRQ(ierr); 451 } 452 } else if (ctx->p == NORM_1) { 453 ierr = HessianMisfit(tao, x, H, Hpre, (void*) ctx);CHKERRQ(ierr); 454 ierr = MatShift(H, ctx->mu);CHKERRQ(ierr); 455 if (Hpre != H) {ierr = MatShift(Hpre, ctx->mu);CHKERRQ(ierr);} 456 } else SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_OUTOFRANGE, "Example only works for NORM_1 and NORM_2"); 457 PetscFunctionReturn(0); 458 } 459 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; 465 PetscErrorCode ierr; 466 467 PetscFunctionBegin; 468 ierr = ObjectiveMisfit(tao, x, &Jm, ctx);CHKERRQ(ierr); 469 ierr = ObjectiveRegularization(tao, x, &Jr, ctx);CHKERRQ(ierr); 470 *J = Jm + Jr; 471 PetscFunctionReturn(0); 472 } 473 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; 479 PetscErrorCode ierr; 480 481 PetscFunctionBegin; 482 ierr = GradientMisfit(tao, x, cntx->workRight[2], ctx);CHKERRQ(ierr); 483 ierr = GradientRegularization(tao, x, cntx->workRight[3], ctx);CHKERRQ(ierr); 484 ierr = VecWAXPY(V,1,cntx->workRight[2],cntx->workRight[3]);CHKERRQ(ierr); 485 PetscFunctionReturn(0); 486 } 487 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; 493 PetscErrorCode ierr; 494 495 PetscFunctionBegin; 496 ierr = MatDuplicate(H, MAT_SHARE_NONZERO_PATTERN, &tempH);CHKERRQ(ierr); 497 ierr = HessianMisfit(tao, x, H, H, ctx);CHKERRQ(ierr); 498 ierr = HessianRegularization(tao, x, tempH, tempH, ctx);CHKERRQ(ierr); 499 ierr = MatAXPY(H, 1., tempH, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr); 500 if (Hpre != H) { 501 ierr = MatCopy(H, Hpre, DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr); 502 } 503 ierr = MatDestroy(&tempH);CHKERRQ(ierr); 504 PetscFunctionReturn(0); 505 } 506 507 static PetscErrorCode TaoSolveADMM(UserCtx ctx, Vec x) 508 { 509 PetscErrorCode ierr; 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; 515 516 PetscFunctionBegin; 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 ierr = VecSet(u, 0.);CHKERRQ(ierr); 526 ierr = TaoCreate(PETSC_COMM_WORLD, &tao1);CHKERRQ(ierr); 527 ierr = TaoSetType(tao1,TAONLS);CHKERRQ(ierr); 528 ierr = TaoSetObjectiveRoutine(tao1, ObjectiveMisfitADMM, (void*) ctx);CHKERRQ(ierr); 529 ierr = TaoSetGradientRoutine(tao1, GradientMisfitADMM, (void*) ctx);CHKERRQ(ierr); 530 ierr = TaoSetHessianRoutine(tao1, ctx->Hm, ctx->Hm, HessianMisfitADMM, (void*) ctx);CHKERRQ(ierr); 531 ierr = VecSet(xk, 0.);CHKERRQ(ierr); 532 ierr = TaoSetInitialVector(tao1, xk);CHKERRQ(ierr); 533 ierr = TaoSetOptionsPrefix(tao1, "misfit_");CHKERRQ(ierr); 534 ierr = TaoSetFromOptions(tao1);CHKERRQ(ierr); 535 ierr = TaoCreate(PETSC_COMM_WORLD, &tao2);CHKERRQ(ierr); 536 if (ctx->p == NORM_2) { 537 ierr = TaoSetType(tao2,TAONLS);CHKERRQ(ierr); 538 ierr = TaoSetObjectiveRoutine(tao2, ObjectiveRegularizationADMM, (void*) ctx);CHKERRQ(ierr); 539 ierr = TaoSetGradientRoutine(tao2, GradientRegularizationADMM, (void*) ctx);CHKERRQ(ierr); 540 ierr = TaoSetHessianRoutine(tao2, ctx->Hr, ctx->Hr, HessianRegularizationADMM, (void*) ctx);CHKERRQ(ierr); 541 } 542 ierr = VecSet(z, 0.);CHKERRQ(ierr); 543 ierr = TaoSetInitialVector(tao2, z);CHKERRQ(ierr); 544 ierr = TaoSetOptionsPrefix(tao2, "reg_");CHKERRQ(ierr); 545 ierr = TaoSetFromOptions(tao2);CHKERRQ(ierr); 546 547 for (i=0; i<ctx->iter; i++) { 548 ierr = VecCopy(z,zold);CHKERRQ(ierr); 549 ierr = TaoSolve(tao1);CHKERRQ(ierr); /* Updates xk */ 550 if (ctx->p == NORM_1){ 551 ierr = VecWAXPY(temp,1.,xk,u);CHKERRQ(ierr); 552 ierr = TaoSoftThreshold(temp,-ctx->alpha/mu,ctx->alpha/mu,z);CHKERRQ(ierr); 553 } else { 554 ierr = TaoSolve(tao2);CHKERRQ(ierr); /* Update zk */ 555 } 556 /* u = u + xk -z */ 557 ierr = VecAXPBYPCZ(u,1.,-1.,1.,xk,z);CHKERRQ(ierr); 558 /* r_norm : norm(x-z) */ 559 ierr = VecWAXPY(diff,-1.,z,xk);CHKERRQ(ierr); 560 ierr = VecNorm(diff,NORM_2,&r_norm);CHKERRQ(ierr); 561 /* s_norm : norm(-mu(z-zold)) */ 562 ierr = VecWAXPY(zdiff, -1.,zold,z);CHKERRQ(ierr); 563 ierr = VecNorm(zdiff,NORM_2,&s_norm);CHKERRQ(ierr); 564 s_norm = s_norm * mu; 565 /* primal : sqrt(n)*ABSTOL + RELTOL*max(norm(x), norm(-z))*/ 566 ierr = VecNorm(xk,NORM_2,&x_norm);CHKERRQ(ierr); 567 ierr = VecNorm(z,NORM_2,&z_norm);CHKERRQ(ierr); 568 primal = PetscSqrtReal(ctx->n)*ctx->abstol + ctx->reltol*PetscMax(x_norm,z_norm); 569 /* Duality : sqrt(n)*ABSTOL + RELTOL*norm(mu*u)*/ 570 ierr = VecNorm(u,NORM_2,&u_norm);CHKERRQ(ierr); 571 dual = PetscSqrtReal(ctx->n)*ctx->abstol + ctx->reltol*u_norm*mu; 572 ierr = PetscPrintf(PetscObjectComm((PetscObject)tao1),"Iter %D : ||x-z||: %g, mu*||z-zold||: %g\n", i, (double) r_norm, (double) s_norm);CHKERRQ(ierr); 573 if (r_norm < primal && s_norm < dual) break; 574 } 575 ierr = VecCopy(xk, x);CHKERRQ(ierr); 576 ierr = TaoDestroy(&tao1);CHKERRQ(ierr); 577 ierr = TaoDestroy(&tao2);CHKERRQ(ierr); 578 PetscFunctionReturn(0); 579 } 580 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; 593 PetscErrorCode ierr; 594 595 PetscFunctionBegin; 596 ierr = VecDuplicate(x, &g);CHKERRQ(ierr); 597 ierr = VecDuplicate(x, &xhat);CHKERRQ(ierr); 598 /* choose a perturbation direction */ 599 ierr = VecDuplicate(x, &dx);CHKERRQ(ierr); 600 ierr = VecSetRandom(dx,ctx->rctx);CHKERRQ(ierr); 601 /* evaluate objective at x: J(x) */ 602 ierr = TaoComputeObjective(tao, x, &Jx);CHKERRQ(ierr); 603 /* evaluate gradient at x, save in vector g */ 604 ierr = TaoComputeGradient(tao, x, g);CHKERRQ(ierr); 605 ierr = VecDot(g, dx, &gdotdx);CHKERRQ(ierr); 606 607 for (numValues=0, h=ctx->hStart; h>=ctx->hMin; h*=ctx->hFactor) numValues++; 608 ierr = PetscCalloc2(numValues, &Js, numValues, &hs);CHKERRQ(ierr); 609 for (i=0, h=ctx->hStart; h>=ctx->hMin; h*=ctx->hFactor, i++) { 610 ierr = VecWAXPY(xhat, h, dx, x);CHKERRQ(ierr); 611 ierr = TaoComputeObjective(tao, xhat, &Jxhat_comp);CHKERRQ(ierr); 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 ierr = PetscPrintf (comm, "J(xhat): %g, predicted: %g, diff %g\n", (double) Jxhat_comp,(double) Jxhat_pred, (double) J);CHKERRQ(ierr); 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 ierr = PetscPrintf (comm, "Convergence rate step %D: %g\n", j-1, (double) temp);CHKERRQ(ierr); 623 minrate = PetscMin(minrate, temp); 624 } 625 /* If O is not ~2, then the test is wrong */ 626 ierr = PetscFree2(Js, hs);CHKERRQ(ierr); 627 *C = minrate; 628 ierr = VecDestroy(&dx);CHKERRQ(ierr); 629 ierr = VecDestroy(&xhat);CHKERRQ(ierr); 630 ierr = VecDestroy(&g);CHKERRQ(ierr); 631 PetscFunctionReturn(0); 632 } 633 634 int main(int argc, char ** argv) 635 { 636 UserCtx ctx; 637 Tao tao; 638 Vec x; 639 Mat H; 640 PetscErrorCode ierr; 641 642 ierr = PetscInitialize(&argc, &argv, NULL,help);if (ierr) return ierr; 643 ierr = PetscNew(&ctx);CHKERRQ(ierr); 644 ierr = ConfigureContext(ctx);CHKERRQ(ierr); 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() */ 648 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 ierr = TaoCreate(PETSC_COMM_WORLD, &tao);CHKERRQ(ierr); 653 ierr = TaoSetType(tao,TAONM);CHKERRQ(ierr); 654 ierr = TaoSetObjectiveRoutine(tao, ObjectiveComplete, (void*) ctx);CHKERRQ(ierr); 655 ierr = TaoSetGradientRoutine(tao, GradientComplete, (void*) ctx);CHKERRQ(ierr); 656 ierr = MatDuplicate(ctx->W, MAT_SHARE_NONZERO_PATTERN, &H);CHKERRQ(ierr); 657 ierr = TaoSetHessianRoutine(tao, H, H, HessianComplete, (void*) ctx);CHKERRQ(ierr); 658 ierr = MatCreateVecs(ctx->F, NULL, &x);CHKERRQ(ierr); 659 ierr = VecSet(x, 0.);CHKERRQ(ierr); 660 ierr = TaoSetInitialVector(tao, x);CHKERRQ(ierr); 661 ierr = TaoSetFromOptions(tao);CHKERRQ(ierr); 662 if (ctx->use_admm) { 663 ierr = TaoSolveADMM(ctx,x); CHKERRQ(ierr); 664 } else {ierr = TaoSolve(tao);CHKERRQ(ierr);} 665 /* examine solution */ 666 ierr = VecViewFromOptions(x, NULL, "-view_sol");CHKERRQ(ierr); 667 if (ctx->taylor) { 668 PetscReal rate; 669 ierr = TaylorTest(ctx, tao, x, &rate);CHKERRQ(ierr); 670 } 671 ierr = MatDestroy(&H);CHKERRQ(ierr); 672 ierr = TaoDestroy(&tao);CHKERRQ(ierr); 673 ierr = VecDestroy(&x);CHKERRQ(ierr); 674 ierr = DestroyContext(&ctx);CHKERRQ(ierr); 675 ierr = PetscFinalize();CHKERRQ(ierr); 676 return ierr; 677 } 678 679 /*TEST 680 681 build: 682 requires: !complex 683 684 test: 685 suffix: 0 686 args: 687 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 691 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 695 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 699 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 703 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 707 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 711 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 715 716 test: 717 suffix: soft_threshold_admm_1 718 args: -matrix_format 1 -m 100 -n 100 -tao_monitor -p 1 -use_admm 719 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 723 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 727 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 731 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 735 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 739 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 743 744 TEST*/ 745