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