1 #include <petscdmda.h> /*I "petscdmda.h" I*/ 2 #include <petsc/private/dmimpl.h> 3 #include <petsc/private/tsimpl.h> /*I "petscts.h" I*/ 4 #include <petscdraw.h> 5 6 /* This structure holds the user-provided DMDA callbacks */ 7 typedef struct { 8 PetscErrorCode (*ifunctionlocal)(DMDALocalInfo *, PetscReal, void *, void *, void *, void *); 9 PetscErrorCode (*rhsfunctionlocal)(DMDALocalInfo *, PetscReal, void *, void *, void *); 10 PetscErrorCode (*ijacobianlocal)(DMDALocalInfo *, PetscReal, void *, void *, PetscReal, Mat, Mat, void *); 11 PetscErrorCode (*rhsjacobianlocal)(DMDALocalInfo *, PetscReal, void *, Mat, Mat, void *); 12 void *ifunctionlocalctx; 13 void *ijacobianlocalctx; 14 void *rhsfunctionlocalctx; 15 void *rhsjacobianlocalctx; 16 InsertMode ifunctionlocalimode; 17 InsertMode rhsfunctionlocalimode; 18 } DMTS_DA; 19 20 static PetscErrorCode DMTSDestroy_DMDA(DMTS sdm) 21 { 22 PetscFunctionBegin; 23 PetscCall(PetscFree(sdm->data)); 24 PetscFunctionReturn(PETSC_SUCCESS); 25 } 26 27 static PetscErrorCode DMTSDuplicate_DMDA(DMTS oldsdm, DMTS sdm) 28 { 29 PetscFunctionBegin; 30 PetscCall(PetscNew((DMTS_DA **)&sdm->data)); 31 if (oldsdm->data) PetscCall(PetscMemcpy(sdm->data, oldsdm->data, sizeof(DMTS_DA))); 32 PetscFunctionReturn(PETSC_SUCCESS); 33 } 34 35 static PetscErrorCode DMDATSGetContext(DM dm, DMTS sdm, DMTS_DA **dmdats) 36 { 37 PetscFunctionBegin; 38 *dmdats = NULL; 39 if (!sdm->data) { 40 PetscCall(PetscNew((DMTS_DA **)&sdm->data)); 41 sdm->ops->destroy = DMTSDestroy_DMDA; 42 sdm->ops->duplicate = DMTSDuplicate_DMDA; 43 } 44 *dmdats = (DMTS_DA *)sdm->data; 45 PetscFunctionReturn(PETSC_SUCCESS); 46 } 47 48 static PetscErrorCode TSComputeIFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec Xdot, Vec F, void *ctx) 49 { 50 DM dm; 51 DMTS_DA *dmdats = (DMTS_DA *)ctx; 52 DMDALocalInfo info; 53 Vec Xloc, Xdotloc; 54 void *x, *f, *xdot; 55 56 PetscFunctionBegin; 57 PetscValidHeaderSpecific(ts, TS_CLASSID, 1); 58 PetscValidHeaderSpecific(X, VEC_CLASSID, 3); 59 PetscValidHeaderSpecific(F, VEC_CLASSID, 5); 60 PetscCheck(dmdats->ifunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context"); 61 PetscCall(TSGetDM(ts, &dm)); 62 PetscCall(DMGetLocalVector(dm, &Xdotloc)); 63 PetscCall(DMGlobalToLocalBegin(dm, Xdot, INSERT_VALUES, Xdotloc)); 64 PetscCall(DMGlobalToLocalEnd(dm, Xdot, INSERT_VALUES, Xdotloc)); 65 PetscCall(DMGetLocalVector(dm, &Xloc)); 66 PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc)); 67 PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc)); 68 PetscCall(DMDAGetLocalInfo(dm, &info)); 69 PetscCall(DMDAVecGetArray(dm, Xloc, &x)); 70 PetscCall(DMDAVecGetArray(dm, Xdotloc, &xdot)); 71 switch (dmdats->ifunctionlocalimode) { 72 case INSERT_VALUES: { 73 PetscCall(DMDAVecGetArray(dm, F, &f)); 74 CHKMEMQ; 75 PetscCall((*dmdats->ifunctionlocal)(&info, ptime, x, xdot, f, dmdats->ifunctionlocalctx)); 76 CHKMEMQ; 77 PetscCall(DMDAVecRestoreArray(dm, F, &f)); 78 } break; 79 case ADD_VALUES: { 80 Vec Floc; 81 PetscCall(DMGetLocalVector(dm, &Floc)); 82 PetscCall(VecZeroEntries(Floc)); 83 PetscCall(DMDAVecGetArray(dm, Floc, &f)); 84 CHKMEMQ; 85 PetscCall((*dmdats->ifunctionlocal)(&info, ptime, x, xdot, f, dmdats->ifunctionlocalctx)); 86 CHKMEMQ; 87 PetscCall(DMDAVecRestoreArray(dm, Floc, &f)); 88 PetscCall(VecZeroEntries(F)); 89 PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F)); 90 PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F)); 91 PetscCall(DMRestoreLocalVector(dm, &Floc)); 92 } break; 93 default: 94 SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->ifunctionlocalimode); 95 } 96 PetscCall(DMDAVecRestoreArray(dm, Xloc, &x)); 97 PetscCall(DMRestoreLocalVector(dm, &Xloc)); 98 PetscCall(DMDAVecRestoreArray(dm, Xdotloc, &xdot)); 99 PetscCall(DMRestoreLocalVector(dm, &Xdotloc)); 100 PetscFunctionReturn(PETSC_SUCCESS); 101 } 102 103 static PetscErrorCode TSComputeIJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Vec Xdot, PetscReal shift, Mat A, Mat B, void *ctx) 104 { 105 DM dm; 106 DMTS_DA *dmdats = (DMTS_DA *)ctx; 107 DMDALocalInfo info; 108 Vec Xloc, Xdotloc; 109 void *x, *xdot; 110 111 PetscFunctionBegin; 112 PetscCheck(dmdats->ifunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context"); 113 PetscCall(TSGetDM(ts, &dm)); 114 115 PetscCheck(dmdats->ijacobianlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()"); 116 PetscCall(DMGetLocalVector(dm, &Xloc)); 117 PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc)); 118 PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc)); 119 PetscCall(DMGetLocalVector(dm, &Xdotloc)); 120 PetscCall(DMGlobalToLocalBegin(dm, Xdot, INSERT_VALUES, Xdotloc)); 121 PetscCall(DMGlobalToLocalEnd(dm, Xdot, INSERT_VALUES, Xdotloc)); 122 PetscCall(DMDAGetLocalInfo(dm, &info)); 123 PetscCall(DMDAVecGetArray(dm, Xloc, &x)); 124 PetscCall(DMDAVecGetArray(dm, Xdotloc, &xdot)); 125 CHKMEMQ; 126 PetscCall((*dmdats->ijacobianlocal)(&info, ptime, x, xdot, shift, A, B, dmdats->ijacobianlocalctx)); 127 CHKMEMQ; 128 PetscCall(DMDAVecRestoreArray(dm, Xloc, &x)); 129 PetscCall(DMDAVecRestoreArray(dm, Xdotloc, &xdot)); 130 PetscCall(DMRestoreLocalVector(dm, &Xloc)); 131 PetscCall(DMRestoreLocalVector(dm, &Xdotloc)); 132 /* This will be redundant if the user called both, but it's too common to forget. */ 133 if (A != B) { 134 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 135 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 136 } 137 PetscFunctionReturn(PETSC_SUCCESS); 138 } 139 140 static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec F, void *ctx) 141 { 142 DM dm; 143 DMTS_DA *dmdats = (DMTS_DA *)ctx; 144 DMDALocalInfo info; 145 Vec Xloc; 146 void *x, *f; 147 148 PetscFunctionBegin; 149 PetscValidHeaderSpecific(ts, TS_CLASSID, 1); 150 PetscValidHeaderSpecific(X, VEC_CLASSID, 3); 151 PetscValidHeaderSpecific(F, VEC_CLASSID, 4); 152 PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context"); 153 PetscCall(TSGetDM(ts, &dm)); 154 PetscCall(DMGetLocalVector(dm, &Xloc)); 155 PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc)); 156 PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc)); 157 PetscCall(DMDAGetLocalInfo(dm, &info)); 158 PetscCall(DMDAVecGetArray(dm, Xloc, &x)); 159 switch (dmdats->rhsfunctionlocalimode) { 160 case INSERT_VALUES: { 161 PetscCall(DMDAVecGetArray(dm, F, &f)); 162 CHKMEMQ; 163 PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx)); 164 CHKMEMQ; 165 PetscCall(DMDAVecRestoreArray(dm, F, &f)); 166 } break; 167 case ADD_VALUES: { 168 Vec Floc; 169 PetscCall(DMGetLocalVector(dm, &Floc)); 170 PetscCall(VecZeroEntries(Floc)); 171 PetscCall(DMDAVecGetArray(dm, Floc, &f)); 172 CHKMEMQ; 173 PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx)); 174 CHKMEMQ; 175 PetscCall(DMDAVecRestoreArray(dm, Floc, &f)); 176 PetscCall(VecZeroEntries(F)); 177 PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F)); 178 PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F)); 179 PetscCall(DMRestoreLocalVector(dm, &Floc)); 180 } break; 181 default: 182 SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->rhsfunctionlocalimode); 183 } 184 PetscCall(DMDAVecRestoreArray(dm, Xloc, &x)); 185 PetscCall(DMRestoreLocalVector(dm, &Xloc)); 186 PetscFunctionReturn(PETSC_SUCCESS); 187 } 188 189 static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Mat A, Mat B, void *ctx) 190 { 191 DM dm; 192 DMTS_DA *dmdats = (DMTS_DA *)ctx; 193 DMDALocalInfo info; 194 Vec Xloc; 195 void *x; 196 197 PetscFunctionBegin; 198 PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context"); 199 PetscCall(TSGetDM(ts, &dm)); 200 201 PetscCheck(dmdats->rhsjacobianlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeRHSJacobian_DMDA() called without calling DMDATSSetRHSJacobian()"); 202 PetscCall(DMGetLocalVector(dm, &Xloc)); 203 PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc)); 204 PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc)); 205 PetscCall(DMDAGetLocalInfo(dm, &info)); 206 PetscCall(DMDAVecGetArray(dm, Xloc, &x)); 207 CHKMEMQ; 208 PetscCall((*dmdats->rhsjacobianlocal)(&info, ptime, x, A, B, dmdats->rhsjacobianlocalctx)); 209 CHKMEMQ; 210 PetscCall(DMDAVecRestoreArray(dm, Xloc, &x)); 211 PetscCall(DMRestoreLocalVector(dm, &Xloc)); 212 /* This will be redundant if the user called both, but it's too common to forget. */ 213 if (A != B) { 214 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 215 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 216 } 217 PetscFunctionReturn(PETSC_SUCCESS); 218 } 219 220 /*@C 221 DMDATSSetRHSFunctionLocal - set a local residual evaluation function for use with `DMDA` 222 223 Logically Collective 224 225 Input Parameters: 226 + dm - `DM` to associate callback with 227 . imode - insert mode for the residual 228 . func - local residual evaluation, see `DMDATSRHSFunctionLocalFn` for the calling sequence 229 - ctx - optional context for local residual evaluation 230 231 Level: beginner 232 233 .seealso: [](ch_ts), `DMDA`, `DMDATSRHSFunctionLocalFn`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()` 234 @*/ 235 PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocalFn *func, void *ctx) 236 { 237 DMTS sdm; 238 DMTS_DA *dmdats; 239 240 PetscFunctionBegin; 241 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 242 PetscCall(DMGetDMTSWrite(dm, &sdm)); 243 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 244 dmdats->rhsfunctionlocalimode = imode; 245 dmdats->rhsfunctionlocal = func; 246 dmdats->rhsfunctionlocalctx = ctx; 247 PetscCall(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats)); 248 PetscFunctionReturn(PETSC_SUCCESS); 249 } 250 251 /*@C 252 DMDATSSetRHSJacobianLocal - set a local residual evaluation function for use with `DMDA` 253 254 Logically Collective 255 256 Input Parameters: 257 + dm - `DM` to associate callback with 258 . func - local RHS Jacobian evaluation routine, see `DMDATSRHSJacobianLocalFn` for the calling sequence 259 - ctx - optional context for local jacobian evaluation 260 261 Level: beginner 262 263 .seealso: [](ch_ts), `DMDA`, `DMDATSRHSJacobianLocalFn`, `DMTSSetRHSJacobian()`, 264 `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()` 265 @*/ 266 PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocalFn *func, void *ctx) 267 { 268 DMTS sdm; 269 DMTS_DA *dmdats; 270 271 PetscFunctionBegin; 272 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 273 PetscCall(DMGetDMTSWrite(dm, &sdm)); 274 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 275 dmdats->rhsjacobianlocal = func; 276 dmdats->rhsjacobianlocalctx = ctx; 277 PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats)); 278 PetscFunctionReturn(PETSC_SUCCESS); 279 } 280 281 /*@C 282 DMDATSSetIFunctionLocal - set a local residual evaluation function for use with `DMDA` 283 284 Logically Collective 285 286 Input Parameters: 287 + dm - `DM` to associate callback with 288 . imode - the insert mode of the function 289 . func - local residual evaluation, see `DMDATSIFunctionLocalFn` for the calling sequence 290 - ctx - optional context for local residual evaluation 291 292 Level: beginner 293 294 .seealso: [](ch_ts), `DMDA`, `DMDATSIFunctionLocalFn`, `DMTSSetIFunction()`, 295 `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()` 296 @*/ 297 PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocalFn *func, void *ctx) 298 { 299 DMTS sdm; 300 DMTS_DA *dmdats; 301 302 PetscFunctionBegin; 303 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 304 PetscCall(DMGetDMTSWrite(dm, &sdm)); 305 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 306 dmdats->ifunctionlocalimode = imode; 307 dmdats->ifunctionlocal = func; 308 dmdats->ifunctionlocalctx = ctx; 309 PetscCall(DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats)); 310 PetscFunctionReturn(PETSC_SUCCESS); 311 } 312 313 /*@C 314 DMDATSSetIJacobianLocal - set a local residual evaluation function for use with `DMDA` 315 316 Logically Collective 317 318 Input Parameters: 319 + dm - `DM` to associate callback with 320 . func - local residual evaluation, see `DMDATSIJacobianLocalFn` for the calling sequence 321 - ctx - optional context for local residual evaluation 322 323 Level: beginner 324 325 .seealso: [](ch_ts), `DMDA`, `DMDATSIJacobianLocalFn`, `DMTSSetIJacobian()`, 326 `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()` 327 @*/ 328 PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocalFn *func, void *ctx) 329 { 330 DMTS sdm; 331 DMTS_DA *dmdats; 332 333 PetscFunctionBegin; 334 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 335 PetscCall(DMGetDMTSWrite(dm, &sdm)); 336 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 337 dmdats->ijacobianlocal = func; 338 dmdats->ijacobianlocalctx = ctx; 339 PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats)); 340 PetscFunctionReturn(PETSC_SUCCESS); 341 } 342 343 PetscErrorCode TSMonitorDMDARayDestroy(void **mctx) 344 { 345 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx; 346 347 PetscFunctionBegin; 348 if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx)); 349 PetscCall(VecDestroy(&rayctx->ray)); 350 PetscCall(VecScatterDestroy(&rayctx->scatter)); 351 PetscCall(PetscViewerDestroy(&rayctx->viewer)); 352 PetscCall(PetscFree(rayctx)); 353 PetscFunctionReturn(PETSC_SUCCESS); 354 } 355 356 PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx) 357 { 358 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx; 359 Vec solution; 360 361 PetscFunctionBegin; 362 PetscCall(TSGetSolution(ts, &solution)); 363 PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD)); 364 PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD)); 365 if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer)); 366 PetscFunctionReturn(PETSC_SUCCESS); 367 } 368 369 PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx) 370 { 371 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx; 372 TSMonitorLGCtx lgctx = rayctx->lgctx; 373 Vec v = rayctx->ray; 374 const PetscScalar *a; 375 PetscInt dim; 376 377 PetscFunctionBegin; 378 PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD)); 379 PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD)); 380 if (!step) { 381 PetscDrawAxis axis; 382 383 PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis)); 384 PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution")); 385 PetscCall(VecGetLocalSize(rayctx->ray, &dim)); 386 PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim)); 387 PetscCall(PetscDrawLGReset(lgctx->lg)); 388 } 389 PetscCall(VecGetArrayRead(v, &a)); 390 #if defined(PETSC_USE_COMPLEX) 391 { 392 PetscReal *areal; 393 PetscInt i, n; 394 PetscCall(VecGetLocalSize(v, &n)); 395 PetscCall(PetscMalloc1(n, &areal)); 396 for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]); 397 PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal)); 398 PetscCall(PetscFree(areal)); 399 } 400 #else 401 PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a)); 402 #endif 403 PetscCall(VecRestoreArrayRead(v, &a)); 404 if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) { 405 PetscCall(PetscDrawLGDraw(lgctx->lg)); 406 PetscCall(PetscDrawLGSave(lgctx->lg)); 407 } 408 PetscFunctionReturn(PETSC_SUCCESS); 409 } 410