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