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