1 /* 2 Defines a SNES that can consist of a collection of SNESes on patches of the domain 3 */ 4 #include <petsc/private/snesimpl.h> /*I "petscsnes.h" I*/ 5 #include <petsc/private/pcpatchimpl.h> /* We need internal access to PCPatch right now, until that part is moved to Plex */ 6 #include <petscsf.h> 7 8 typedef struct { 9 PC pc; /* The linear patch preconditioner */ 10 SNESCompositeType type; 11 } SNES_Patch; 12 13 static PetscErrorCode SNESPatchComputeResidual_Private(SNES snes, Vec x, Vec F, void *ctx) 14 { 15 PC pc = (PC) ctx; 16 PC_PATCH *pcpatch = (PC_PATCH *) pc->data; 17 PetscErrorCode ierr; 18 19 PetscFunctionBegin; 20 ierr = PCPatchComputeFunction_Internal(pc, x, F, pcpatch->currentPatch);CHKERRQ(ierr); 21 PetscFunctionReturn(0); 22 } 23 24 static PetscErrorCode SNESPatchComputeJacobian_Private(SNES snes, Vec x, Mat J, Mat M, void *ctx) 25 { 26 PC pc = (PC) ctx; 27 PC_PATCH *pcpatch = (PC_PATCH *) pc->data; 28 PetscErrorCode ierr; 29 30 PetscFunctionBegin; 31 ierr = PCPatchComputeOperator_Internal(pc, x, M, pcpatch->currentPatch, PETSC_FALSE);CHKERRQ(ierr); 32 PetscFunctionReturn(0); 33 } 34 35 static PetscErrorCode PCSetUp_PATCH_Nonlinear(PC pc) 36 { 37 PC_PATCH *patch = (PC_PATCH *) pc->data; 38 const char *prefix; 39 PetscInt i; 40 PetscErrorCode ierr; 41 42 PetscFunctionBegin; 43 if (!pc->setupcalled) { 44 ierr = PetscMalloc1(patch->npatch, &patch->solver);CHKERRQ(ierr); 45 ierr = PCGetOptionsPrefix(pc, &prefix);CHKERRQ(ierr); 46 for (i = 0; i < patch->npatch; ++i) { 47 SNES snes; 48 KSP subksp; 49 50 ierr = SNESCreate(PETSC_COMM_SELF, &snes);CHKERRQ(ierr); 51 ierr = SNESSetOptionsPrefix(snes, prefix);CHKERRQ(ierr); 52 ierr = SNESAppendOptionsPrefix(snes, "sub_");CHKERRQ(ierr); 53 ierr = PetscObjectIncrementTabLevel((PetscObject) snes, (PetscObject) pc, 2);CHKERRQ(ierr); 54 ierr = SNESGetKSP(snes, &subksp);CHKERRQ(ierr); 55 ierr = PetscObjectIncrementTabLevel((PetscObject) subksp, (PetscObject) pc, 2);CHKERRQ(ierr); 56 ierr = PetscLogObjectParent((PetscObject) pc, (PetscObject) snes);CHKERRQ(ierr); 57 patch->solver[i] = (PetscObject) snes; 58 } 59 60 ierr = PetscMalloc1(patch->npatch, &patch->patchResidual);CHKERRQ(ierr); 61 ierr = PetscMalloc1(patch->npatch, &patch->patchState);CHKERRQ(ierr); 62 for (i = 0; i < patch->npatch; ++i) { 63 ierr = VecDuplicate(patch->patchRHS[i], &patch->patchResidual[i]);CHKERRQ(ierr); 64 ierr = VecDuplicate(patch->patchUpdate[i], &patch->patchState[i]);CHKERRQ(ierr); 65 } 66 ierr = VecDuplicate(patch->localUpdate, &patch->localState);CHKERRQ(ierr); 67 } 68 for (i = 0; i < patch->npatch; ++i) { 69 SNES snes = (SNES) patch->solver[i]; 70 71 ierr = SNESSetFunction(snes, patch->patchResidual[i], SNESPatchComputeResidual_Private, pc);CHKERRQ(ierr); 72 ierr = SNESSetJacobian(snes, patch->mat[i], patch->mat[i], SNESPatchComputeJacobian_Private, pc);CHKERRQ(ierr); 73 } 74 if (!pc->setupcalled && patch->optionsSet) for (i = 0; i < patch->npatch; ++i) {ierr = SNESSetFromOptions((SNES) patch->solver[i]);CHKERRQ(ierr);} 75 PetscFunctionReturn(0); 76 } 77 78 static PetscErrorCode PCApply_PATCH_Nonlinear(PC pc, PetscInt i, Vec patchRHS, Vec patchUpdate) 79 { 80 PC_PATCH *patch = (PC_PATCH *) pc->data; 81 PetscInt pStart; 82 PetscErrorCode ierr; 83 84 PetscFunctionBegin; 85 patch->currentPatch = i; 86 ierr = PetscLogEventBegin(PC_Patch_Solve, pc, 0, 0, 0);CHKERRQ(ierr); 87 88 /* Scatter the overlapped global state to our patch state vector */ 89 ierr = PetscSectionGetChart(patch->gtolCounts, &pStart, NULL);CHKERRQ(ierr); 90 ierr = PCPatch_ScatterLocal_Private(pc, i+pStart, patch->localState, patch->patchState[i], INSERT_VALUES, SCATTER_FORWARD, PETSC_FALSE);CHKERRQ(ierr); 91 92 /* Set initial guess to be current state*/ 93 ierr = VecCopy(patch->patchState[i], patchUpdate);CHKERRQ(ierr); 94 /* Solve for new state */ 95 ierr = SNESSolve((SNES) patch->solver[i], patchRHS, patchUpdate);CHKERRQ(ierr); 96 /* To compute update, subtract off previous state */ 97 ierr = VecAXPY(patchUpdate, -1.0, patch->patchState[i]);CHKERRQ(ierr); 98 99 ierr = PetscLogEventEnd(PC_Patch_Solve, pc, 0, 0, 0);CHKERRQ(ierr); 100 PetscFunctionReturn(0); 101 } 102 103 static PetscErrorCode PCReset_PATCH_Nonlinear(PC pc) 104 { 105 PC_PATCH *patch = (PC_PATCH *) pc->data; 106 PetscInt i; 107 PetscErrorCode ierr; 108 109 PetscFunctionBegin; 110 111 if (patch->solver) { 112 for (i = 0; i < patch->npatch; ++i) {ierr = SNESReset((SNES) patch->solver[i]);CHKERRQ(ierr);} 113 } 114 115 if (patch->patchResidual) { 116 for (i = 0; i < patch->npatch; ++i) {ierr = VecDestroy(&patch->patchResidual[i]);CHKERRQ(ierr);} 117 ierr = PetscFree(patch->patchResidual);CHKERRQ(ierr); 118 } 119 120 if (patch->patchState) { 121 for (i = 0; i < patch->npatch; ++i) {ierr = VecDestroy(&patch->patchState[i]);CHKERRQ(ierr);} 122 ierr = PetscFree(patch->patchState);CHKERRQ(ierr); 123 } 124 125 ierr = VecDestroy(&patch->localState);CHKERRQ(ierr); 126 127 PetscFunctionReturn(0); 128 } 129 130 static PetscErrorCode PCDestroy_PATCH_Nonlinear(PC pc) 131 { 132 PC_PATCH *patch = (PC_PATCH *) pc->data; 133 PetscInt i; 134 PetscErrorCode ierr; 135 136 PetscFunctionBegin; 137 if (patch->solver) { 138 for (i = 0; i < patch->npatch; ++i) {ierr = SNESDestroy((SNES *) &patch->solver[i]);CHKERRQ(ierr);} 139 ierr = PetscFree(patch->solver);CHKERRQ(ierr); 140 } 141 PetscFunctionReturn(0); 142 } 143 144 static PetscErrorCode SNESSetUp_Patch(SNES snes) 145 { 146 SNES_Patch *patch = (SNES_Patch *) snes->data; 147 DM dm; 148 Mat dummy; 149 Vec F; 150 PetscInt n, N; 151 PetscErrorCode ierr; 152 153 PetscFunctionBegin; 154 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 155 ierr = PCSetDM(patch->pc, dm);CHKERRQ(ierr); 156 ierr = SNESGetFunction(snes, &F, NULL, NULL);CHKERRQ(ierr); 157 ierr = VecGetLocalSize(F, &n);CHKERRQ(ierr); 158 ierr = VecGetSize(F, &N);CHKERRQ(ierr); 159 ierr = MatCreateShell(PetscObjectComm((PetscObject) snes), n, n, N, N, (void *) snes, &dummy);CHKERRQ(ierr); 160 ierr = PCSetOperators(patch->pc, dummy, dummy);CHKERRQ(ierr); 161 ierr = MatDestroy(&dummy);CHKERRQ(ierr); 162 ierr = PCSetUp(patch->pc);CHKERRQ(ierr); 163 /* allocate workspace */ 164 PetscFunctionReturn(0); 165 } 166 167 static PetscErrorCode SNESReset_Patch(SNES snes) 168 { 169 SNES_Patch *patch = (SNES_Patch *) snes->data; 170 PetscErrorCode ierr; 171 172 PetscFunctionBegin; 173 ierr = PCReset(patch->pc);CHKERRQ(ierr); 174 PetscFunctionReturn(0); 175 } 176 177 static PetscErrorCode SNESDestroy_Patch(SNES snes) 178 { 179 SNES_Patch *patch = (SNES_Patch *) snes->data; 180 PetscErrorCode ierr; 181 182 PetscFunctionBegin; 183 ierr = SNESReset_Patch(snes);CHKERRQ(ierr); 184 ierr = PCDestroy(&patch->pc);CHKERRQ(ierr); 185 ierr = PetscFree(snes->data);CHKERRQ(ierr); 186 PetscFunctionReturn(0); 187 } 188 189 static PetscErrorCode SNESSetFromOptions_Patch(PetscOptionItems *PetscOptionsObject, SNES snes) 190 { 191 SNES_Patch *patch = (SNES_Patch *) snes->data; 192 PetscBool flg; 193 const char *prefix; 194 PetscErrorCode ierr; 195 196 PetscFunctionBegin; 197 ierr = PetscOptionsHead(PetscOptionsObject, "Patch nonlinear preconditioner options");CHKERRQ(ierr); 198 ierr = PetscOptionsEnum("-snes_patch_type", "Type of composition", "SNESPatchSetType", SNESCompositeTypes, (PetscEnum) patch->type, (PetscEnum *) &patch->type, &flg);CHKERRQ(ierr); 199 if (flg) {ierr = SNESPatchSetType(snes, patch->type);CHKERRQ(ierr);} 200 ierr = PetscOptionsTail();CHKERRQ(ierr); 201 202 ierr = PetscObjectGetOptionsPrefix((PetscObject)snes, &prefix);CHKERRQ(ierr); 203 ierr = PetscObjectSetOptionsPrefix((PetscObject)patch->pc, prefix);CHKERRQ(ierr); 204 ierr = PCSetFromOptions(patch->pc);CHKERRQ(ierr); 205 PetscFunctionReturn(0); 206 } 207 208 static PetscErrorCode SNESView_Patch(SNES snes,PetscViewer viewer) 209 { 210 SNES_Patch *patch = (SNES_Patch *) snes->data; 211 PetscBool iascii; 212 PetscErrorCode ierr; 213 214 PetscFunctionBegin; 215 ierr = PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);CHKERRQ(ierr); 216 if (iascii) { 217 ierr = PetscViewerASCIIPrintf(viewer," type - %s\n",SNESCompositeTypes[patch->type]);CHKERRQ(ierr); 218 } 219 ierr = PetscViewerASCIIPushTab(viewer);CHKERRQ(ierr); 220 ierr = PCView(patch->pc, viewer);CHKERRQ(ierr); 221 ierr = PetscViewerASCIIPopTab(viewer);CHKERRQ(ierr); 222 PetscFunctionReturn(0); 223 } 224 225 static PetscErrorCode SNESSolve_Patch(SNES snes) 226 { 227 SNES_Patch *patch = (SNES_Patch *) snes->data; 228 PC_PATCH *pcpatch = (PC_PATCH *) patch->pc->data; 229 Vec rhs, update, state; 230 const PetscScalar *globalState = NULL; 231 PetscScalar *localState = NULL; 232 233 PetscErrorCode ierr; 234 235 PetscFunctionBegin; 236 237 ierr = SNESGetSolution(snes, &state);CHKERRQ(ierr); 238 ierr = SNESGetSolutionUpdate(snes, &update);CHKERRQ(ierr); 239 ierr = SNESGetRhs(snes, &rhs);CHKERRQ(ierr); 240 241 /* These happen in a loop */ 242 { 243 /* Scatter state vector to overlapped vector on all patches. 244 The vector pcpatch->localState is scattered to each patch 245 in PCApply_PATCH_Nonlinear. */ 246 ierr = VecGetArrayRead(state, &globalState);CHKERRQ(ierr); 247 ierr = VecGetArray(pcpatch->localState, &localState);CHKERRQ(ierr); 248 ierr = PetscSFBcastBegin(pcpatch->defaultSF, MPIU_SCALAR, globalState, localState);CHKERRQ(ierr); 249 ierr = PetscSFBcastEnd(pcpatch->defaultSF, MPIU_SCALAR, globalState, localState);CHKERRQ(ierr); 250 ierr = VecRestoreArray(pcpatch->localState, &localState);CHKERRQ(ierr); 251 ierr = VecRestoreArrayRead(state, &globalState);CHKERRQ(ierr); 252 253 ierr = PCApply(patch->pc, rhs, update); 254 ierr = VecAXPY(state, 1.0, update);CHKERRQ(ierr); /* FIXME: replace with line search */ 255 } 256 257 ierr = SNESSetConvergedReason(snes, SNES_CONVERGED_ITS);CHKERRQ(ierr); 258 PetscFunctionReturn(0); 259 } 260 261 static PetscErrorCode SNESPatchSetType_Patch(SNES snes, SNESCompositeType type) 262 { 263 SNES_Patch *patch = (SNES_Patch *) snes->data; 264 265 PetscFunctionBegin; 266 patch->type = type; 267 PetscFunctionReturn(0); 268 } 269 270 /*MC 271 SNESPATCH - Solve a nonlinear problem by composing together many nonlinear solvers on patches 272 273 Options Database Keys: 274 . -snes_patch_type <type: one of additive, multiplicative, symmetric_multiplicative> - Sets composition type 275 276 Level: intermediate 277 278 Concepts: composing solvers 279 280 .seealso: SNESCreate(), SNESSetType(), SNESType (for list of available types), SNES, 281 PCPATCH 282 283 References: 284 . 1. - Peter R. Brune, Matthew G. Knepley, Barry F. Smith, and Xuemin Tu, "Composing Scalable Nonlinear Algebraic Solvers", SIAM Review, 57(4), 2015 285 286 M*/ 287 PETSC_EXTERN PetscErrorCode SNESCreate_Patch(SNES snes) 288 { 289 PetscErrorCode ierr; 290 SNES_Patch *patch; 291 PC_PATCH *patchpc; 292 293 PetscFunctionBegin; 294 ierr = PetscNewLog(snes, &patch);CHKERRQ(ierr); 295 296 snes->ops->solve = SNESSolve_Patch; 297 snes->ops->setup = SNESSetUp_Patch; 298 snes->ops->reset = SNESReset_Patch; 299 snes->ops->destroy = SNESDestroy_Patch; 300 snes->ops->setfromoptions = SNESSetFromOptions_Patch; 301 snes->ops->view = SNESView_Patch; 302 303 snes->alwayscomputesfinalresidual = PETSC_FALSE; 304 305 snes->data = (void *) patch; 306 patch->type = SNES_COMPOSITE_ADDITIVE; 307 ierr = PCCreate(PetscObjectComm((PetscObject) snes), &patch->pc);CHKERRQ(ierr); 308 ierr = PCSetType(patch->pc, PCPATCH);CHKERRQ(ierr); 309 310 patchpc = (PC_PATCH*) patch->pc->data; 311 patchpc->classname = "snes"; 312 313 patchpc->setupsolver = PCSetUp_PATCH_Nonlinear; 314 patchpc->applysolver = PCApply_PATCH_Nonlinear; 315 patchpc->resetsolver = PCReset_PATCH_Nonlinear; 316 patchpc->destroysolver = PCDestroy_PATCH_Nonlinear; 317 318 ierr = PetscObjectComposeFunction((PetscObject) snes, "SNESPatchSetType_C", SNESPatchSetType_Patch);CHKERRQ(ierr); 319 PetscFunctionReturn(0); 320 } 321 322 /*@ 323 SNESPatchSetType - Sets the type of composition. 324 325 Logically Collective on SNES 326 327 Input Parameter: 328 + snes - the preconditioner context 329 - type - SNES_COMPOSITE_ADDITIVE (default), SNES_COMPOSITE_MULTIPLICATIVE, etc. 330 331 Options Database Key: 332 . -snes_composite_type <type: one of additive, multiplicative, etc> - Sets composite preconditioner type 333 334 Level: Developer 335 336 .keywords: SNES, set, type, composite preconditioner, additive, multiplicative 337 @*/ 338 PetscErrorCode SNESPatchSetType(SNES snes, SNESCompositeType type) 339 { 340 PetscErrorCode ierr; 341 342 PetscFunctionBegin; 343 PetscValidHeaderSpecific(snes, SNES_CLASSID, 1); 344 PetscValidLogicalCollectiveEnum(snes,type, 2); 345 ierr = PetscTryMethod(snes, "SNESPatchSetType_C", (SNES,SNESCompositeType), (snes,type));CHKERRQ(ierr); 346 PetscFunctionReturn(0); 347 } 348 349 PetscErrorCode SNESPatchSetDiscretisationInfo(SNES snes, PetscInt nsubspaces, DM *dms, PetscInt *bs, PetscInt *nodesPerCell, const PetscInt **cellNodeMap, 350 const PetscInt *subspaceOffsets, PetscInt numGhostBcs, const PetscInt *ghostBcNodes, PetscInt numGlobalBcs, const PetscInt *globalBcNodes) 351 { 352 SNES_Patch *patch = (SNES_Patch *) snes->data; 353 PetscErrorCode ierr; 354 DM dm; 355 356 PetscFunctionBegin; 357 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 358 if (!dm) SETERRQ(PetscObjectComm((PetscObject)snes), PETSC_ERR_ARG_WRONGSTATE, "DM not yet set on patch SNES\n"); 359 ierr = PCSetDM(patch->pc, dm);CHKERRQ(ierr); 360 ierr = PCPatchSetDiscretisationInfo(patch->pc, nsubspaces, dms, bs, nodesPerCell, cellNodeMap, subspaceOffsets, numGhostBcs, ghostBcNodes, numGlobalBcs, globalBcNodes);CHKERRQ(ierr); 361 PetscFunctionReturn(0); 362 } 363 364 PetscErrorCode SNESPatchSetComputeOperator(SNES snes, PetscErrorCode (*func)(PC, PetscInt, Vec, Mat, IS, PetscInt, const PetscInt *, void *), void *ctx) 365 { 366 SNES_Patch *patch = (SNES_Patch *) snes->data; 367 PetscErrorCode ierr; 368 369 PetscFunctionBegin; 370 ierr = PCPatchSetComputeOperator(patch->pc, func, ctx);CHKERRQ(ierr); 371 PetscFunctionReturn(0); 372 } 373 374 PetscErrorCode SNESPatchSetComputeFunction(SNES snes, PetscErrorCode (*func)(PC, PetscInt, Vec, Vec, IS, PetscInt, const PetscInt *, void *), void *ctx) 375 { 376 SNES_Patch *patch = (SNES_Patch *) snes->data; 377 PetscErrorCode ierr; 378 379 PetscFunctionBegin; 380 ierr = PCPatchSetComputeFunction(patch->pc, func, ctx);CHKERRQ(ierr); 381 PetscFunctionReturn(0); 382 } 383 384 PetscErrorCode SNESPatchSetConstructType(SNES snes, PCPatchConstructType ctype, PetscErrorCode (*func)(PC, PetscInt *, IS **, IS *, void *), void *ctx) 385 { 386 SNES_Patch *patch = (SNES_Patch *) snes->data; 387 PetscErrorCode ierr; 388 389 PetscFunctionBegin; 390 ierr = PCPatchSetConstructType(patch->pc, ctype, func, ctx);CHKERRQ(ierr); 391 PetscFunctionReturn(0); 392 } 393 394 PetscErrorCode SNESPatchSetCellNumbering(SNES snes, PetscSection cellNumbering) 395 { 396 SNES_Patch *patch = (SNES_Patch *) snes->data; 397 PetscErrorCode ierr; 398 399 PetscFunctionBegin; 400 ierr = PCPatchSetCellNumbering(patch->pc, cellNumbering);CHKERRQ(ierr); 401 PetscFunctionReturn(0); 402 } 403