Date: 2022-08-02T03:15:09+00:00

Actual source code: dmdats.c

      1: #include <petscdmda.h>
      2: #include <petsc/private/dmimpl.h>
      3: #include <petsc/private/tsimpl.h>
      4: #include <petscdraw.h>

      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;

     20: static PetscErrorCode DMTSDestroy_DMDA(DMTS sdm)
     21: {
     22:   PetscFree(sdm->data);
     23:   return 0;
     24: }

     26: static PetscErrorCode DMTSDuplicate_DMDA(DMTS oldsdm,DMTS sdm)
     27: {
     28:   PetscNewLog(sdm,(DMTS_DA**)&sdm->data);
     29:   if (oldsdm->data) PetscMemcpy(sdm->data,oldsdm->data,sizeof(DMTS_DA));
     30:   return 0;
     31: }

     33: static PetscErrorCode DMDATSGetContext(DM dm,DMTS sdm,DMTS_DA **dmdats)
     34: {
     35:   *dmdats = NULL;
     36:   if (!sdm->data) {
     37:     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:   return 0;
     43: }

     45: static PetscErrorCode TSComputeIFunction_DMDA(TS ts,PetscReal ptime,Vec X,Vec Xdot,Vec F,void *ctx)
     46: {
     47:   DM             dm;
     48:   DMTS_DA        *dmdats = (DMTS_DA*)ctx;
     49:   DMDALocalInfo  info;
     50:   Vec            Xloc,Xdotloc;
     51:   void           *x,*f,*xdot;

     57:   TSGetDM(ts,&dm);
     58:   DMGetLocalVector(dm,&Xdotloc);
     59:   DMGlobalToLocalBegin(dm,Xdot,INSERT_VALUES,Xdotloc);
     60:   DMGlobalToLocalEnd(dm,Xdot,INSERT_VALUES,Xdotloc);
     61:   DMGetLocalVector(dm,&Xloc);
     62:   DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
     63:   DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
     64:   DMDAGetLocalInfo(dm,&info);
     65:   DMDAVecGetArray(dm,Xloc,&x);
     66:   DMDAVecGetArray(dm,Xdotloc,&xdot);
     67:   switch (dmdats->ifunctionlocalimode) {
     68:   case INSERT_VALUES: {
     69:     DMDAVecGetArray(dm,F,&f);
     70:     CHKMEMQ;
     71:     (*dmdats->ifunctionlocal)(&info,ptime,x,xdot,f,dmdats->ifunctionlocalctx);
     72:     CHKMEMQ;
     73:     DMDAVecRestoreArray(dm,F,&f);
     74:   } break;
     75:   case ADD_VALUES: {
     76:     Vec Floc;
     77:     DMGetLocalVector(dm,&Floc);
     78:     VecZeroEntries(Floc);
     79:     DMDAVecGetArray(dm,Floc,&f);
     80:     CHKMEMQ;
     81:     (*dmdats->ifunctionlocal)(&info,ptime,x,xdot,f,dmdats->ifunctionlocalctx);
     82:     CHKMEMQ;
     83:     DMDAVecRestoreArray(dm,Floc,&f);
     84:     VecZeroEntries(F);
     85:     DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);
     86:     DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);
     87:     DMRestoreLocalVector(dm,&Floc);
     88:   } break;
     89:   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdats->ifunctionlocalimode);
     90:   }
     91:   DMDAVecRestoreArray(dm,Xloc,&x);
     92:   DMRestoreLocalVector(dm,&Xloc);
     93:   DMDAVecRestoreArray(dm,Xdotloc,&xdot);
     94:   DMRestoreLocalVector(dm,&Xdotloc);
     95:   return 0;
     96: }

     98: static PetscErrorCode TSComputeIJacobian_DMDA(TS ts,PetscReal ptime,Vec X,Vec Xdot,PetscReal shift,Mat A,Mat B,void *ctx)
     99: {
    100:   DM             dm;
    101:   DMTS_DA        *dmdats = (DMTS_DA*)ctx;
    102:   DMDALocalInfo  info;
    103:   Vec            Xloc;
    104:   void           *x,*xdot;

    107:   TSGetDM(ts,&dm);

    109:   if (dmdats->ijacobianlocal) {
    110:     DMGetLocalVector(dm,&Xloc);
    111:     DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
    112:     DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
    113:     DMDAGetLocalInfo(dm,&info);
    114:     DMDAVecGetArray(dm,Xloc,&x);
    115:     DMDAVecGetArray(dm,Xdot,&xdot);
    116:     CHKMEMQ;
    117:     (*dmdats->ijacobianlocal)(&info,ptime,x,xdot,shift,A,B,dmdats->ijacobianlocalctx);
    118:     CHKMEMQ;
    119:     DMDAVecRestoreArray(dm,Xloc,&x);
    120:     DMDAVecRestoreArray(dm,Xdot,&xdot);
    121:     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:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
    126:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
    127:   }
    128:   return 0;
    129: }

    131: static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts,PetscReal ptime,Vec X,Vec F,void *ctx)
    132: {
    133:   DM             dm;
    134:   DMTS_DA        *dmdats = (DMTS_DA*)ctx;
    135:   DMDALocalInfo  info;
    136:   Vec            Xloc;
    137:   void           *x,*f;

    143:   TSGetDM(ts,&dm);
    144:   DMGetLocalVector(dm,&Xloc);
    145:   DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
    146:   DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
    147:   DMDAGetLocalInfo(dm,&info);
    148:   DMDAVecGetArray(dm,Xloc,&x);
    149:   switch (dmdats->rhsfunctionlocalimode) {
    150:   case INSERT_VALUES: {
    151:     DMDAVecGetArray(dm,F,&f);
    152:     CHKMEMQ;
    153:     (*dmdats->rhsfunctionlocal)(&info,ptime,x,f,dmdats->rhsfunctionlocalctx);
    154:     CHKMEMQ;
    155:     DMDAVecRestoreArray(dm,F,&f);
    156:   } break;
    157:   case ADD_VALUES: {
    158:     Vec Floc;
    159:     DMGetLocalVector(dm,&Floc);
    160:     VecZeroEntries(Floc);
    161:     DMDAVecGetArray(dm,Floc,&f);
    162:     CHKMEMQ;
    163:     (*dmdats->rhsfunctionlocal)(&info,ptime,x,f,dmdats->rhsfunctionlocalctx);
    164:     CHKMEMQ;
    165:     DMDAVecRestoreArray(dm,Floc,&f);
    166:     VecZeroEntries(F);
    167:     DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);
    168:     DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);
    169:     DMRestoreLocalVector(dm,&Floc);
    170:   } break;
    171:   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdats->rhsfunctionlocalimode);
    172:   }
    173:   DMDAVecRestoreArray(dm,Xloc,&x);
    174:   DMRestoreLocalVector(dm,&Xloc);
    175:   return 0;
    176: }

    178: static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts,PetscReal ptime,Vec X,Mat A,Mat B,void *ctx)
    179: {
    180:   DM             dm;
    181:   DMTS_DA        *dmdats = (DMTS_DA*)ctx;
    182:   DMDALocalInfo  info;
    183:   Vec            Xloc;
    184:   void           *x;

    187:   TSGetDM(ts,&dm);

    189:   if (dmdats->rhsjacobianlocal) {
    190:     DMGetLocalVector(dm,&Xloc);
    191:     DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
    192:     DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
    193:     DMDAGetLocalInfo(dm,&info);
    194:     DMDAVecGetArray(dm,Xloc,&x);
    195:     CHKMEMQ;
    196:     (*dmdats->rhsjacobianlocal)(&info,ptime,x,A,B,dmdats->rhsjacobianlocalctx);
    197:     CHKMEMQ;
    198:     DMDAVecRestoreArray(dm,Xloc,&x);
    199:     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:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
    204:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
    205:   }
    206:   return 0;
    207: }

    209: /*@C
    210:    DMDATSSetRHSFunctionLocal - set a local residual evaluation function

    212:    Logically Collective

    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

    220:    Calling sequence for func:

    222: $ func(DMDALocalInfo info,PetscReal t,void *x,void *f,void *ctx)

    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

    230:    Level: beginner

    232: .seealso: DMTSSetRHSFunction(), DMDATSSetRHSJacobianLocal(), DMDASNESSetFunctionLocal()
    233: @*/
    234: PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm,InsertMode imode,DMDATSRHSFunctionLocal func,void *ctx)
    235: {
    236:   DMTS           sdm;
    237:   DMTS_DA        *dmdats;

    240:   DMGetDMTSWrite(dm,&sdm);
    241:   DMDATSGetContext(dm,sdm,&dmdats);
    242:   dmdats->rhsfunctionlocalimode = imode;
    243:   dmdats->rhsfunctionlocal      = func;
    244:   dmdats->rhsfunctionlocalctx   = ctx;
    245:   DMTSSetRHSFunction(dm,TSComputeRHSFunction_DMDA,dmdats);
    246:   return 0;
    247: }

    249: /*@C
    250:    DMDATSSetRHSJacobianLocal - set a local residual evaluation function

    252:    Logically Collective

    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

    259:    Calling sequence for func:

    261: $ func(DMDALocalInfo* info,PetscReal t,void* x,Mat J,Mat B,void *ctx);

    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

    270:    Level: beginner

    272: .seealso: DMTSSetRHSJacobian(), DMDATSSetRHSFunctionLocal(), DMDASNESSetJacobianLocal()
    273: @*/
    274: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm,DMDATSRHSJacobianLocal func,void *ctx)
    275: {
    276:   DMTS           sdm;
    277:   DMTS_DA        *dmdats;

    280:   DMGetDMTSWrite(dm,&sdm);
    281:   DMDATSGetContext(dm,sdm,&dmdats);
    282:   dmdats->rhsjacobianlocal    = func;
    283:   dmdats->rhsjacobianlocalctx = ctx;
    284:   DMTSSetRHSJacobian(dm,TSComputeRHSJacobian_DMDA,dmdats);
    285:   return 0;
    286: }

    288: /*@C
    289:    DMDATSSetIFunctionLocal - set a local residual evaluation function

    291:    Logically Collective

    293:    Input Parameters:
    294: +  dm   - DM to associate callback with
    295: .  func - local residual evaluation
    296: -  ctx  - optional context for local residual evaluation

    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

    306:    Level: beginner

    308: .seealso: DMTSSetIFunction(), DMDATSSetIJacobianLocal(), DMDASNESSetFunctionLocal()
    309: @*/
    310: PetscErrorCode DMDATSSetIFunctionLocal(DM dm,InsertMode imode,DMDATSIFunctionLocal func,void *ctx)
    311: {
    312:   DMTS           sdm;
    313:   DMTS_DA        *dmdats;

    316:   DMGetDMTSWrite(dm,&sdm);
    317:   DMDATSGetContext(dm,sdm,&dmdats);
    318:   dmdats->ifunctionlocalimode = imode;
    319:   dmdats->ifunctionlocal      = func;
    320:   dmdats->ifunctionlocalctx   = ctx;
    321:   DMTSSetIFunction(dm,TSComputeIFunction_DMDA,dmdats);
    322:   return 0;
    323: }

    325: /*@C
    326:    DMDATSSetIJacobianLocal - set a local residual evaluation function

    328:    Logically Collective

    330:    Input Parameters:
    331: +  dm   - DM to associate callback with
    332: .  func - local residual evaluation
    333: -  ctx   - optional context for local residual evaluation

    335:    Calling sequence for func:

    337: $ func(DMDALocalInfo* info,PetscReal t,void* x,void *xdot,PetscScalar shift,Mat J,Mat B,void *ctx);

    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

    348:    Level: beginner

    350: .seealso: DMTSSetJacobian(), DMDATSSetIFunctionLocal(), DMDASNESSetJacobianLocal()
    351: @*/
    352: PetscErrorCode DMDATSSetIJacobianLocal(DM dm,DMDATSIJacobianLocal func,void *ctx)
    353: {
    354:   DMTS           sdm;
    355:   DMTS_DA        *dmdats;

    358:   DMGetDMTSWrite(dm,&sdm);
    359:   DMDATSGetContext(dm,sdm,&dmdats);
    360:   dmdats->ijacobianlocal    = func;
    361:   dmdats->ijacobianlocalctx = ctx;
    362:   DMTSSetIJacobian(dm,TSComputeIJacobian_DMDA,dmdats);
    363:   return 0;
    364: }

    366: PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
    367: {
    368:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *) *mctx;

    370:   if (rayctx->lgctx) TSMonitorLGCtxDestroy(&rayctx->lgctx);
    371:   VecDestroy(&rayctx->ray);
    372:   VecScatterDestroy(&rayctx->scatter);
    373:   PetscViewerDestroy(&rayctx->viewer);
    374:   PetscFree(rayctx);
    375:   return 0;
    376: }

    378: PetscErrorCode TSMonitorDMDARay(TS ts,PetscInt steps,PetscReal time,Vec u,void *mctx)
    379: {
    380:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx*)mctx;
    381:   Vec                 solution;

    383:   TSGetSolution(ts,&solution);
    384:   VecScatterBegin(rayctx->scatter,solution,rayctx->ray,INSERT_VALUES,SCATTER_FORWARD);
    385:   VecScatterEnd(rayctx->scatter,solution,rayctx->ray,INSERT_VALUES,SCATTER_FORWARD);
    386:   if (rayctx->viewer) {
    387:     VecView(rayctx->ray,rayctx->viewer);
    388:   }
    389:   return 0;
    390: }

    392: PetscErrorCode  TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx)
    393: {
    394:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *) ctx;
    395:   TSMonitorLGCtx       lgctx  = (TSMonitorLGCtx) rayctx->lgctx;
    396:   Vec                  v      = rayctx->ray;
    397:   const PetscScalar   *a;
    398:   PetscInt             dim;

    400:   VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD);
    401:   VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD);
    402:   if (!step) {
    403:     PetscDrawAxis axis;

    405:     PetscDrawLGGetAxis(lgctx->lg, &axis);
    406:     PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution");
    407:     VecGetLocalSize(rayctx->ray, &dim);
    408:     PetscDrawLGSetDimension(lgctx->lg, dim);
    409:     PetscDrawLGReset(lgctx->lg);
    410:   }
    411:   VecGetArrayRead(v, &a);
    412: #if defined(PETSC_USE_COMPLEX)
    413:   {
    414:     PetscReal *areal;
    415:     PetscInt   i,n;
    416:     VecGetLocalSize(v, &n);
    417:     PetscMalloc1(n, &areal);
    418:     for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]);
    419:     PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal);
    420:     PetscFree(areal);
    421:   }
    422: #else
    423:   PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a);
    424: #endif
    425:   VecRestoreArrayRead(v, &a);
    426:   if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) {
    427:     PetscDrawLGDraw(lgctx->lg);
    428:     PetscDrawLGSave(lgctx->lg);
    429:   }
    430:   return 0;
    431: }