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; 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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 of `func`: 230 $ PetscErrorCode func(DMDALocalInfo *info, PetscReal t, void *x, void *f, void *ctx) 231 + info - defines the subdomain to evaluate the residual on 232 . t - time at which to evaluate residual 233 . x - array of local state information 234 . f - output array of local residual information 235 - ctx - optional user context 236 237 Level: beginner 238 239 .seealso: [](ch_ts), `DMDA`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()` 240 @*/ 241 PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocal func, void *ctx) 242 { 243 DMTS sdm; 244 DMTS_DA *dmdats; 245 246 PetscFunctionBegin; 247 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 248 PetscCall(DMGetDMTSWrite(dm, &sdm)); 249 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 250 dmdats->rhsfunctionlocalimode = imode; 251 dmdats->rhsfunctionlocal = func; 252 dmdats->rhsfunctionlocalctx = ctx; 253 PetscCall(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats)); 254 PetscFunctionReturn(PETSC_SUCCESS); 255 } 256 257 /*@C 258 DMDATSSetRHSJacobianLocal - set a local residual evaluation function for use with `DMDA` 259 260 Logically Collective 261 262 Input Parameters: 263 + dm - `DM` to associate callback with 264 . func - local RHS Jacobian evaluation routine 265 - ctx - optional context for local jacobian evaluation 266 267 Calling sequence of `func`: 268 $ PetscErrorCode func(DMDALocalInfo *info, PetscReal t, void* x, Mat J, Mat B, void *ctx) 269 + info - defines the subdomain to evaluate the residual on 270 . t - time at which to evaluate residual 271 . x - array of local state information 272 . J - Jacobian matrix 273 . B - preconditioner matrix; often same as `J` 274 - ctx - optional context passed above 275 276 Level: beginner 277 278 .seealso: [](ch_ts), `DMDA`, `DMTSSetRHSJacobian()`, `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()` 279 @*/ 280 PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocal func, void *ctx) 281 { 282 DMTS sdm; 283 DMTS_DA *dmdats; 284 285 PetscFunctionBegin; 286 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 287 PetscCall(DMGetDMTSWrite(dm, &sdm)); 288 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 289 dmdats->rhsjacobianlocal = func; 290 dmdats->rhsjacobianlocalctx = ctx; 291 PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats)); 292 PetscFunctionReturn(PETSC_SUCCESS); 293 } 294 295 /*@C 296 DMDATSSetIFunctionLocal - set a local residual evaluation function for use with `DMDA` 297 298 Logically Collective 299 300 Input Parameters: 301 + dm - `DM` to associate callback with 302 . func - local residual evaluation 303 - ctx - optional context for local residual evaluation 304 305 Calling sequence of `func`: 306 $ PetscErrorCode func(DMDALocalInfo *info, PetscReal t, Vec x, Vec xdot, Vec f, void *ctx) 307 + info - defines the subdomain to evaluate the residual on 308 . t - time at which to evaluate residual 309 . x - array of local state information 310 . xdot - array of local time derivative information 311 . f - output array of local function evaluation information 312 - ctx - optional context passed above 313 314 Level: beginner 315 316 .seealso: [](ch_ts), `DMDA`, `DMTSSetIFunction()`, `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()` 317 @*/ 318 PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocal func, void *ctx) 319 { 320 DMTS sdm; 321 DMTS_DA *dmdats; 322 323 PetscFunctionBegin; 324 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 325 PetscCall(DMGetDMTSWrite(dm, &sdm)); 326 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 327 dmdats->ifunctionlocalimode = imode; 328 dmdats->ifunctionlocal = func; 329 dmdats->ifunctionlocalctx = ctx; 330 PetscCall(DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats)); 331 PetscFunctionReturn(PETSC_SUCCESS); 332 } 333 334 /*@C 335 DMDATSSetIJacobianLocal - set a local residual evaluation function for use with `DMDA` 336 337 Logically Collective 338 339 Input Parameters: 340 + dm - `DM` to associate callback with 341 . func - local residual evaluation 342 - ctx - optional context for local residual evaluation 343 344 Calling sequence of `func`: 345 $ PetscErrorCode func(DMDALocalInfo *info, PetscReal t, void* x, void *xdot, PetscScalar shift, Mat J, Mat B, void *ctx) 346 + info - defines the subdomain to evaluate the residual on 347 . t - time at which to evaluate the jacobian 348 . x - array of local state information 349 . xdot - time derivative at this state 350 . shift - see `TSSetIJacobian()` for the meaning of this parameter 351 . J - Jacobian matrix 352 . B - preconditioner matrix; often same as `J` 353 - ctx - optional context passed above 354 355 Level: beginner 356 357 .seealso: [](ch_ts), `DMDA`, `DMTSSetJacobian()`, `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()` 358 @*/ 359 PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocal func, void *ctx) 360 { 361 DMTS sdm; 362 DMTS_DA *dmdats; 363 364 PetscFunctionBegin; 365 PetscValidHeaderSpecific(dm, DM_CLASSID, 1); 366 PetscCall(DMGetDMTSWrite(dm, &sdm)); 367 PetscCall(DMDATSGetContext(dm, sdm, &dmdats)); 368 dmdats->ijacobianlocal = func; 369 dmdats->ijacobianlocalctx = ctx; 370 PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats)); 371 PetscFunctionReturn(PETSC_SUCCESS); 372 } 373 374 PetscErrorCode TSMonitorDMDARayDestroy(void **mctx) 375 { 376 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx; 377 378 PetscFunctionBegin; 379 if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx)); 380 PetscCall(VecDestroy(&rayctx->ray)); 381 PetscCall(VecScatterDestroy(&rayctx->scatter)); 382 PetscCall(PetscViewerDestroy(&rayctx->viewer)); 383 PetscCall(PetscFree(rayctx)); 384 PetscFunctionReturn(PETSC_SUCCESS); 385 } 386 387 PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx) 388 { 389 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx; 390 Vec solution; 391 392 PetscFunctionBegin; 393 PetscCall(TSGetSolution(ts, &solution)); 394 PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD)); 395 PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD)); 396 if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer)); 397 PetscFunctionReturn(PETSC_SUCCESS); 398 } 399 400 PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx) 401 { 402 TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx; 403 TSMonitorLGCtx lgctx = (TSMonitorLGCtx)rayctx->lgctx; 404 Vec v = rayctx->ray; 405 const PetscScalar *a; 406 PetscInt dim; 407 408 PetscFunctionBegin; 409 PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD)); 410 PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD)); 411 if (!step) { 412 PetscDrawAxis axis; 413 414 PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis)); 415 PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution")); 416 PetscCall(VecGetLocalSize(rayctx->ray, &dim)); 417 PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim)); 418 PetscCall(PetscDrawLGReset(lgctx->lg)); 419 } 420 PetscCall(VecGetArrayRead(v, &a)); 421 #if defined(PETSC_USE_COMPLEX) 422 { 423 PetscReal *areal; 424 PetscInt i, n; 425 PetscCall(VecGetLocalSize(v, &n)); 426 PetscCall(PetscMalloc1(n, &areal)); 427 for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]); 428 PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal)); 429 PetscCall(PetscFree(areal)); 430 } 431 #else 432 PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a)); 433 #endif 434 PetscCall(VecRestoreArrayRead(v, &a)); 435 if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) { 436 PetscCall(PetscDrawLGDraw(lgctx->lg)); 437 PetscCall(PetscDrawLGSave(lgctx->lg)); 438 } 439 PetscFunctionReturn(PETSC_SUCCESS); 440 } 441