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(0); 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(0); 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(0); 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(0); 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; 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 if (dmdats->ijacobianlocal) { 116 PetscCall(DMGetLocalVector(dm, &Xloc)); 117 PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc)); 118 PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc)); 119 PetscCall(DMDAGetLocalInfo(dm, &info)); 120 PetscCall(DMDAVecGetArray(dm, Xloc, &x)); 121 PetscCall(DMDAVecGetArray(dm, Xdot, &xdot)); 122 CHKMEMQ; 123 PetscCall((*dmdats->ijacobianlocal)(&info, ptime, x, xdot, shift, A, B, dmdats->ijacobianlocalctx)); 124 CHKMEMQ; 125 PetscCall(DMDAVecRestoreArray(dm, Xloc, &x)); 126 PetscCall(DMDAVecRestoreArray(dm, Xdot, &xdot)); 127 PetscCall(DMRestoreLocalVector(dm, &Xloc)); 128 } else SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()"); 129 /* This will be redundant if the user called both, but it's too common to forget. */ 130 if (A != B) { 131 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 132 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 133 } 134 PetscFunctionReturn(0); 135 } 136 137 static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec F, void *ctx) 138 { 139 DM dm; 140 DMTS_DA *dmdats = (DMTS_DA *)ctx; 141 DMDALocalInfo info; 142 Vec Xloc; 143 void *x, *f; 144 145 PetscFunctionBegin; 146 PetscValidHeaderSpecific(ts, TS_CLASSID, 1); 147 PetscValidHeaderSpecific(X, VEC_CLASSID, 3); 148 PetscValidHeaderSpecific(F, VEC_CLASSID, 4); 149 PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context"); 150 PetscCall(TSGetDM(ts, &dm)); 151 PetscCall(DMGetLocalVector(dm, &Xloc)); 152 PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc)); 153 PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc)); 154 PetscCall(DMDAGetLocalInfo(dm, &info)); 155 PetscCall(DMDAVecGetArray(dm, Xloc, &x)); 156 switch (dmdats->rhsfunctionlocalimode) { 157 case INSERT_VALUES: { 158 PetscCall(DMDAVecGetArray(dm, F, &f)); 159 CHKMEMQ; 160 PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx)); 161 CHKMEMQ; 162 PetscCall(DMDAVecRestoreArray(dm, F, &f)); 163 } break; 164 case ADD_VALUES: { 165 Vec Floc; 166 PetscCall(DMGetLocalVector(dm, &Floc)); 167 PetscCall(VecZeroEntries(Floc)); 168 PetscCall(DMDAVecGetArray(dm, Floc, &f)); 169 CHKMEMQ; 170 PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx)); 171 CHKMEMQ; 172 PetscCall(DMDAVecRestoreArray(dm, Floc, &f)); 173 PetscCall(VecZeroEntries(F)); 174 PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F)); 175 PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F)); 176 PetscCall(DMRestoreLocalVector(dm, &Floc)); 177 } break; 178 default: 179 SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->rhsfunctionlocalimode); 180 } 181 PetscCall(DMDAVecRestoreArray(dm, Xloc, &x)); 182 PetscCall(DMRestoreLocalVector(dm, &Xloc)); 183 PetscFunctionReturn(0); 184 } 185 186 static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Mat A, Mat B, void *ctx) 187 { 188 DM dm; 189 DMTS_DA *dmdats = (DMTS_DA *)ctx; 190 DMDALocalInfo info; 191 Vec Xloc; 192 void *x; 193 194 PetscFunctionBegin; 195 PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context"); 196 PetscCall(TSGetDM(ts, &dm)); 197 198 if (dmdats->rhsjacobianlocal) { 199 PetscCall(DMGetLocalVector(dm, &Xloc)); 200 PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc)); 201 PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc)); 202 PetscCall(DMDAGetLocalInfo(dm, &info)); 203 PetscCall(DMDAVecGetArray(dm, Xloc, &x)); 204 CHKMEMQ; 205 PetscCall((*dmdats->rhsjacobianlocal)(&info, ptime, x, A, B, dmdats->rhsjacobianlocalctx)); 206 CHKMEMQ; 207 PetscCall(DMDAVecRestoreArray(dm, Xloc, &x)); 208 PetscCall(DMRestoreLocalVector(dm, &Xloc)); 209 } else SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeRHSJacobian_DMDA() called without calling DMDATSSetRHSJacobian()"); 210 /* This will be redundant if the user called both, but it's too common to forget. */ 211 if (A != B) { 212 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 213 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 214 } 215 PetscFunctionReturn(0); 216 } 217 218 /*@C 219 DMDATSSetRHSFunctionLocal - set a local residual evaluation function for use with `DMDA` 220 221 Logically Collective 222 223 Input Parameters: 224 + dm - `DM` to associate callback with 225 . imode - insert mode for the residual 226 . func - local residual evaluation 227 - ctx - optional context for local residual evaluation 228 229 Calling sequence for func: 230 231 $ func(DMDALocalInfo info,PetscReal t,void *x,void *f,void *ctx) 232 233 + info - `DMDALocalInfo` defining the subdomain to evaluate the residual on 234 . t - time at which to evaluate residual 235 . x - array of local state information 236 . f - output array of local residual information 237 - ctx - optional user context 238 239 Level: beginner 240 241 .seealso: [](chapter_ts), `DMDA`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()` 242 @*/ 243 PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocal func, void *ctx) 244 { 245 DMTS sdm; 246 DMTS_DA *dmdats; 247 248 PetscFunctionBegin; 249 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 250 PetscCall(DMGetDMTSWrite(dm, &sdm)); 251 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 252 dmdats->rhsfunctionlocalimode = imode; 253 dmdats->rhsfunctionlocal = func; 254 dmdats->rhsfunctionlocalctx = ctx; 255 PetscCall(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats)); 256 PetscFunctionReturn(0); 257 } 258 259 /*@C 260 DMDATSSetRHSJacobianLocal - set a local residual evaluation function for use with `DMDA` 261 262 Logically Collective 263 264 Input Parameters: 265 + dm - `DM` to associate callback with 266 . func - local RHS Jacobian evaluation routine 267 - ctx - optional context for local jacobian evaluation 268 269 Calling sequence for func: 270 271 $ func(DMDALocalInfo* info,PetscReal t,void* x,Mat J,Mat B,void *ctx); 272 273 + info - `DMDALocalInfo` defining the subdomain to evaluate the residual on 274 . t - time at which to evaluate residual 275 . x - array of local state information 276 . J - Jacobian matrix 277 . B - preconditioner matrix; often same as J 278 - ctx - optional context passed above 279 280 Level: beginner 281 282 .seealso: [](chapter_ts), `DMDA`, `DMTSSetRHSJacobian()`, `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()` 283 @*/ 284 PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocal func, void *ctx) 285 { 286 DMTS sdm; 287 DMTS_DA *dmdats; 288 289 PetscFunctionBegin; 290 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 291 PetscCall(DMGetDMTSWrite(dm, &sdm)); 292 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 293 dmdats->rhsjacobianlocal = func; 294 dmdats->rhsjacobianlocalctx = ctx; 295 PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats)); 296 PetscFunctionReturn(0); 297 } 298 299 /*@C 300 DMDATSSetIFunctionLocal - set a local residual evaluation function for use with `DMDA` 301 302 Logically Collective 303 304 Input Parameters: 305 + dm - `DM` to associate callback with 306 . func - local residual evaluation 307 - ctx - optional context for local residual evaluation 308 309 Calling sequence for func: 310 + info - `DMDALocalInfo` defining the subdomain to evaluate the residual on 311 . t - time at which to evaluate residual 312 . x - array of local state information 313 . xdot - array of local time derivative information 314 . f - output array of local function evaluation information 315 - ctx - optional context passed above 316 317 Level: beginner 318 319 .seealso: [](chapter_ts), `DMDA`, `DMTSSetIFunction()`, `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()` 320 @*/ 321 PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocal func, void *ctx) 322 { 323 DMTS sdm; 324 DMTS_DA *dmdats; 325 326 PetscFunctionBegin; 327 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 328 PetscCall(DMGetDMTSWrite(dm, &sdm)); 329 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 330 dmdats->ifunctionlocalimode = imode; 331 dmdats->ifunctionlocal = func; 332 dmdats->ifunctionlocalctx = ctx; 333 PetscCall(DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats)); 334 PetscFunctionReturn(0); 335 } 336 337 /*@C 338 DMDATSSetIJacobianLocal - set a local residual evaluation function for use with `DMDA` 339 340 Logically Collective 341 342 Input Parameters: 343 + dm - `DM` to associate callback with 344 . func - local residual evaluation 345 - ctx - optional context for local residual evaluation 346 347 Calling sequence for func: 348 349 $ func(DMDALocalInfo* info,PetscReal t,void* x,void *xdot,PetscScalar shift,Mat J,Mat B,void *ctx); 350 351 + info - `DMDALocalInfo` defining the subdomain to evaluate the residual on 352 . t - time at which to evaluate the jacobian 353 . x - array of local state information 354 . xdot - time derivative at this state 355 . shift - see TSSetIJacobian() for the meaning of this parameter 356 . J - Jacobian matrix 357 . B - preconditioner matrix; often same as J 358 - ctx - optional context passed above 359 360 Level: beginner 361 362 .seealso: [](chapter_ts), `DMDA`, `DMTSSetJacobian()`, `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()` 363 @*/ 364 PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocal func, void *ctx) 365 { 366 DMTS sdm; 367 DMTS_DA *dmdats; 368 369 PetscFunctionBegin; 370 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 371 PetscCall(DMGetDMTSWrite(dm, &sdm)); 372 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 373 dmdats->ijacobianlocal = func; 374 dmdats->ijacobianlocalctx = ctx; 375 PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats)); 376 PetscFunctionReturn(0); 377 } 378 379 PetscErrorCode TSMonitorDMDARayDestroy(void **mctx) 380 { 381 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx; 382 383 PetscFunctionBegin; 384 if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx)); 385 PetscCall(VecDestroy(&rayctx->ray)); 386 PetscCall(VecScatterDestroy(&rayctx->scatter)); 387 PetscCall(PetscViewerDestroy(&rayctx->viewer)); 388 PetscCall(PetscFree(rayctx)); 389 PetscFunctionReturn(0); 390 } 391 392 PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx) 393 { 394 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx; 395 Vec solution; 396 397 PetscFunctionBegin; 398 PetscCall(TSGetSolution(ts, &solution)); 399 PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD)); 400 PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD)); 401 if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer)); 402 PetscFunctionReturn(0); 403 } 404 405 PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx) 406 { 407 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx; 408 TSMonitorLGCtx lgctx = (TSMonitorLGCtx)rayctx->lgctx; 409 Vec v = rayctx->ray; 410 const PetscScalar *a; 411 PetscInt dim; 412 413 PetscFunctionBegin; 414 PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD)); 415 PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD)); 416 if (!step) { 417 PetscDrawAxis axis; 418 419 PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis)); 420 PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution")); 421 PetscCall(VecGetLocalSize(rayctx->ray, &dim)); 422 PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim)); 423 PetscCall(PetscDrawLGReset(lgctx->lg)); 424 } 425 PetscCall(VecGetArrayRead(v, &a)); 426 #if defined(PETSC_USE_COMPLEX) 427 { 428 PetscReal *areal; 429 PetscInt i, n; 430 PetscCall(VecGetLocalSize(v, &n)); 431 PetscCall(PetscMalloc1(n, &areal)); 432 for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]); 433 PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal)); 434 PetscCall(PetscFree(areal)); 435 } 436 #else 437 PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a)); 438 #endif 439 PetscCall(VecRestoreArrayRead(v, &a)); 440 if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) { 441 PetscCall(PetscDrawLGDraw(lgctx->lg)); 442 PetscCall(PetscDrawLGSave(lgctx->lg)); 443 } 444 PetscFunctionReturn(0); 445 } 446