Date: 2022-08-02T03:14:57+00:00

Actual source code: dmdasnes.c

      1: #include <petscdmda.h>
      2: #include <petsc/private/dmimpl.h>
      3: #include <petsc/private/snesimpl.h>

      5: /* This structure holds the user-provided DMDA callbacks */
      6: typedef struct {
      7:   PetscErrorCode (*residuallocal)(DMDALocalInfo*,void*,void*,void*);
      8:   PetscErrorCode (*jacobianlocal)(DMDALocalInfo*,void*,Mat,Mat,void*);
      9:   PetscErrorCode (*objectivelocal)(DMDALocalInfo*,void*,PetscReal*,void*);
     10:   void       *residuallocalctx;
     11:   void       *jacobianlocalctx;
     12:   void       *objectivelocalctx;
     13:   InsertMode residuallocalimode;

     15:   /*   For Picard iteration defined locally */
     16:   PetscErrorCode (*rhsplocal)(DMDALocalInfo*,void*,void*,void*);
     17:   PetscErrorCode (*jacobianplocal)(DMDALocalInfo*,void*,Mat,Mat,void*);
     18:   void *picardlocalctx;
     19: } DMSNES_DA;

     21: static PetscErrorCode DMSNESDestroy_DMDA(DMSNES sdm)
     22: {
     23:   PetscFree(sdm->data);
     24:   return 0;
     25: }

     27: static PetscErrorCode DMSNESDuplicate_DMDA(DMSNES oldsdm,DMSNES sdm)
     28: {
     29:   PetscNewLog(sdm,(DMSNES_DA**)&sdm->data);
     30:   if (oldsdm->data) {
     31:     PetscMemcpy(sdm->data,oldsdm->data,sizeof(DMSNES_DA));
     32:   }
     33:   return 0;
     34: }

     36: static PetscErrorCode DMDASNESGetContext(DM dm,DMSNES sdm,DMSNES_DA  **dmdasnes)
     37: {
     38:   *dmdasnes = NULL;
     39:   if (!sdm->data) {
     40:     PetscNewLog(dm,(DMSNES_DA**)&sdm->data);
     41:     sdm->ops->destroy   = DMSNESDestroy_DMDA;
     42:     sdm->ops->duplicate = DMSNESDuplicate_DMDA;
     43:   }
     44:   *dmdasnes = (DMSNES_DA*)sdm->data;
     45:   return 0;
     46: }

     48: static PetscErrorCode SNESComputeFunction_DMDA(SNES snes,Vec X,Vec F,void *ctx)
     49: {
     50:   DM             dm;
     51:   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
     52:   DMDALocalInfo  info;
     53:   Vec            Xloc;
     54:   void           *x,*f;

     60:   SNESGetDM(snes,&dm);
     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:   switch (dmdasnes->residuallocalimode) {
     67:   case INSERT_VALUES: {
     68:     DMDAVecGetArray(dm,F,&f);
     69:     PetscLogEventBegin(SNES_FunctionEval,snes,X,F,0);
     70:     CHKMEMQ;
     71:     (*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx);
     72:     CHKMEMQ;
     73:     PetscLogEventEnd(SNES_FunctionEval,snes,X,F,0);
     74:     DMDAVecRestoreArray(dm,F,&f);
     75:   } break;
     76:   case ADD_VALUES: {
     77:     Vec Floc;
     78:     DMGetLocalVector(dm,&Floc);
     79:     VecZeroEntries(Floc);
     80:     DMDAVecGetArray(dm,Floc,&f);
     81:     PetscLogEventBegin(SNES_FunctionEval,snes,X,F,0);
     82:     CHKMEMQ;
     83:     (*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx);
     84:     CHKMEMQ;
     85:     PetscLogEventEnd(SNES_FunctionEval,snes,X,F,0);
     86:     DMDAVecRestoreArray(dm,Floc,&f);
     87:     VecZeroEntries(F);
     88:     DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);
     89:     DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);
     90:     DMRestoreLocalVector(dm,&Floc);
     91:   } break;
     92:   default: SETERRQ(PetscObjectComm((PetscObject)snes),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdasnes->residuallocalimode);
     93:   }
     94:   DMDAVecRestoreArray(dm,Xloc,&x);
     95:   DMRestoreLocalVector(dm,&Xloc);
     96:   if (snes->domainerror) {
     97:     VecSetInf(F);
     98:   }
     99:   return 0;
    100: }

    102: static PetscErrorCode SNESComputeObjective_DMDA(SNES snes,Vec X,PetscReal *ob,void *ctx)
    103: {
    104:   DM             dm;
    105:   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
    106:   DMDALocalInfo  info;
    107:   Vec            Xloc;
    108:   void           *x;

    114:   SNESGetDM(snes,&dm);
    115:   DMGetLocalVector(dm,&Xloc);
    116:   DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
    117:   DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
    118:   DMDAGetLocalInfo(dm,&info);
    119:   DMDAVecGetArray(dm,Xloc,&x);
    120:   CHKMEMQ;
    121:   (*dmdasnes->objectivelocal)(&info,x,ob,dmdasnes->objectivelocalctx);
    122:   CHKMEMQ;
    123:   DMDAVecRestoreArray(dm,Xloc,&x);
    124:   DMRestoreLocalVector(dm,&Xloc);
    125:   return 0;
    126: }

    128: /* Routine is called by example, hence must be labeled PETSC_EXTERN */
    129: PETSC_EXTERN PetscErrorCode SNESComputeJacobian_DMDA(SNES snes,Vec X,Mat A,Mat B,void *ctx)
    130: {
    131:   DM             dm;
    132:   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
    133:   DMDALocalInfo  info;
    134:   Vec            Xloc;
    135:   void           *x;

    138:   SNESGetDM(snes,&dm);

    140:   if (dmdasnes->jacobianlocal) {
    141:     DMGetLocalVector(dm,&Xloc);
    142:     DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
    143:     DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
    144:     DMDAGetLocalInfo(dm,&info);
    145:     DMDAVecGetArray(dm,Xloc,&x);
    146:     CHKMEMQ;
    147:     (*dmdasnes->jacobianlocal)(&info,x,A,B,dmdasnes->jacobianlocalctx);
    148:     CHKMEMQ;
    149:     DMDAVecRestoreArray(dm,Xloc,&x);
    150:     DMRestoreLocalVector(dm,&Xloc);
    151:   } else {
    152:     MatFDColoring fdcoloring;
    153:     PetscObjectQuery((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject*)&fdcoloring);
    154:     if (!fdcoloring) {
    155:       ISColoring coloring;

    157:       DMCreateColoring(dm,dm->coloringtype,&coloring);
    158:       MatFDColoringCreate(B,coloring,&fdcoloring);
    159:       switch (dm->coloringtype) {
    160:       case IS_COLORING_GLOBAL:
    161:         MatFDColoringSetFunction(fdcoloring,(PetscErrorCode (*)(void))SNESComputeFunction_DMDA,dmdasnes);
    162:         break;
    163:       default: SETERRQ(PetscObjectComm((PetscObject)snes),PETSC_ERR_SUP,"No support for coloring type '%s'",ISColoringTypes[dm->coloringtype]);
    164:       }
    165:       PetscObjectSetOptionsPrefix((PetscObject)fdcoloring,((PetscObject)dm)->prefix);
    166:       MatFDColoringSetFromOptions(fdcoloring);
    167:       MatFDColoringSetUp(B,coloring,fdcoloring);
    168:       ISColoringDestroy(&coloring);
    169:       PetscObjectCompose((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject)fdcoloring);
    170:       PetscObjectDereference((PetscObject)fdcoloring);

    172:       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
    173:        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
    174:        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
    175:        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
    176:        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
    177:        */
    178:       PetscObjectDereference((PetscObject)dm);
    179:     }
    180:     MatFDColoringApply(B,fdcoloring,X,snes);
    181:   }
    182:   /* This will be redundant if the user called both, but it's too common to forget. */
    183:   if (A != B) {
    184:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
    185:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
    186:   }
    187:   return 0;
    188: }

    190: /*@C
    191:    DMDASNESSetFunctionLocal - set a local residual evaluation function

    193:    Logically Collective

    195:    Input Parameters:
    196: +  dm - DM to associate callback with
    197: .  imode - INSERT_VALUES if local function computes owned part, ADD_VALUES if it contributes to ghosted part
    198: .  func - local residual evaluation
    199: -  ctx - optional context for local residual evaluation

    201:    Calling sequence:
    202:    For PetscErrorCode (*func)(DMDALocalInfo *info,void *x, void *f, void *ctx),
    203: +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
    204: .  x - dimensional pointer to state at which to evaluate residual (e.g. PetscScalar *x or **x or ***x)
    205: .  f - dimensional pointer to residual, write the residual here (e.g. PetscScalar *f or **f or ***f)
    206: -  ctx - optional context passed above

    208:    Level: beginner

    210: .seealso: DMDASNESSetJacobianLocal(), DMSNESSetFunction(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
    211: @*/
    212: PetscErrorCode DMDASNESSetFunctionLocal(DM dm,InsertMode imode,PetscErrorCode (*func)(DMDALocalInfo*,void*,void*,void*),void *ctx)
    213: {
    214:   DMSNES         sdm;
    215:   DMSNES_DA      *dmdasnes;

    218:   DMGetDMSNESWrite(dm,&sdm);
    219:   DMDASNESGetContext(dm,sdm,&dmdasnes);

    221:   dmdasnes->residuallocalimode = imode;
    222:   dmdasnes->residuallocal      = func;
    223:   dmdasnes->residuallocalctx   = ctx;

    225:   DMSNESSetFunction(dm,SNESComputeFunction_DMDA,dmdasnes);
    226:   if (!sdm->ops->computejacobian) {  /* Call us for the Jacobian too, can be overridden by the user. */
    227:     DMSNESSetJacobian(dm,SNESComputeJacobian_DMDA,dmdasnes);
    228:   }
    229:   return 0;
    230: }

    232: /*@C
    233:    DMDASNESSetJacobianLocal - set a local Jacobian evaluation function

    235:    Logically Collective

    237:    Input Parameters:
    238: +  dm - DM to associate callback with
    239: .  func - local Jacobian evaluation
    240: -  ctx - optional context for local Jacobian evaluation

    242:    Calling sequence:
    243:    For PetscErrorCode (*func)(DMDALocalInfo *info,void *x,Mat J,Mat M,void *ctx),
    244: +  info - DMDALocalInfo defining the subdomain to evaluate the Jacobian at
    245: .  x - dimensional pointer to state at which to evaluate Jacobian (e.g. PetscScalar *x or **x or ***x)
    246: .  J - Mat object for the Jacobian
    247: .  M - Mat object for the Jacobian preconditioner matrix
    248: -  ctx - optional context passed above

    250:    Level: beginner

    252: .seealso: DMDASNESSetFunctionLocal(), DMSNESSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
    253: @*/
    254: PetscErrorCode DMDASNESSetJacobianLocal(DM dm,PetscErrorCode (*func)(DMDALocalInfo*,void*,Mat,Mat,void*),void *ctx)
    255: {
    256:   DMSNES         sdm;
    257:   DMSNES_DA      *dmdasnes;

    260:   DMGetDMSNESWrite(dm,&sdm);
    261:   DMDASNESGetContext(dm,sdm,&dmdasnes);

    263:   dmdasnes->jacobianlocal    = func;
    264:   dmdasnes->jacobianlocalctx = ctx;

    266:   DMSNESSetJacobian(dm,SNESComputeJacobian_DMDA,dmdasnes);
    267:   return 0;
    268: }

    270: /*@C
    271:    DMDASNESSetObjectiveLocal - set a local residual evaluation function

    273:    Logically Collective

    275:    Input Parameters:
    276: +  dm - DM to associate callback with
    277: .  func - local objective evaluation
    278: -  ctx - optional context for local residual evaluation

    280:    Calling sequence for func:
    281: +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
    282: .  x - dimensional pointer to state at which to evaluate residual
    283: .  ob - eventual objective value
    284: -  ctx - optional context passed above

    286:    Level: beginner

    288: .seealso: DMSNESSetFunction(), DMDASNESSetJacobianLocal(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
    289: @*/
    290: PetscErrorCode DMDASNESSetObjectiveLocal(DM dm,DMDASNESObjective func,void *ctx)
    291: {
    292:   DMSNES         sdm;
    293:   DMSNES_DA      *dmdasnes;

    296:   DMGetDMSNESWrite(dm,&sdm);
    297:   DMDASNESGetContext(dm,sdm,&dmdasnes);

    299:   dmdasnes->objectivelocal    = func;
    300:   dmdasnes->objectivelocalctx = ctx;

    302:   DMSNESSetObjective(dm,SNESComputeObjective_DMDA,dmdasnes);
    303:   return 0;
    304: }

    306: static PetscErrorCode SNESComputePicard_DMDA(SNES snes,Vec X,Vec F,void *ctx)
    307: {
    308:   DM             dm;
    309:   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
    310:   DMDALocalInfo  info;
    311:   Vec            Xloc;
    312:   void           *x,*f;

    318:   SNESGetDM(snes,&dm);
    319:   DMGetLocalVector(dm,&Xloc);
    320:   DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
    321:   DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
    322:   DMDAGetLocalInfo(dm,&info);
    323:   DMDAVecGetArray(dm,Xloc,&x);
    324:   switch (dmdasnes->residuallocalimode) {
    325:   case INSERT_VALUES: {
    326:     DMDAVecGetArray(dm,F,&f);
    327:     CHKMEMQ;
    328:     (*dmdasnes->rhsplocal)(&info,x,f,dmdasnes->picardlocalctx);
    329:     CHKMEMQ;
    330:     DMDAVecRestoreArray(dm,F,&f);
    331:   } break;
    332:   case ADD_VALUES: {
    333:     Vec Floc;
    334:     DMGetLocalVector(dm,&Floc);
    335:     VecZeroEntries(Floc);
    336:     DMDAVecGetArray(dm,Floc,&f);
    337:     CHKMEMQ;
    338:     (*dmdasnes->rhsplocal)(&info,x,f,dmdasnes->picardlocalctx);
    339:     CHKMEMQ;
    340:     DMDAVecRestoreArray(dm,Floc,&f);
    341:     VecZeroEntries(F);
    342:     DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);
    343:     DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);
    344:     DMRestoreLocalVector(dm,&Floc);
    345:   } break;
    346:   default: SETERRQ(PetscObjectComm((PetscObject)snes),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdasnes->residuallocalimode);
    347:   }
    348:   DMDAVecRestoreArray(dm,Xloc,&x);
    349:   DMRestoreLocalVector(dm,&Xloc);
    350:   return 0;
    351: }

    353: static PetscErrorCode SNESComputePicardJacobian_DMDA(SNES snes,Vec X,Mat A,Mat B,void *ctx)
    354: {
    355:   DM             dm;
    356:   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
    357:   DMDALocalInfo  info;
    358:   Vec            Xloc;
    359:   void           *x;

    362:   SNESGetDM(snes,&dm);

    364:   DMGetLocalVector(dm,&Xloc);
    365:   DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);
    366:   DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);
    367:   DMDAGetLocalInfo(dm,&info);
    368:   DMDAVecGetArray(dm,Xloc,&x);
    369:   CHKMEMQ;
    370:   (*dmdasnes->jacobianplocal)(&info,x,A,B,dmdasnes->picardlocalctx);
    371:   CHKMEMQ;
    372:   DMDAVecRestoreArray(dm,Xloc,&x);
    373:   DMRestoreLocalVector(dm,&Xloc);
    374:   if (A != B) {
    375:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
    376:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
    377:   }
    378:   return 0;
    379: }

    381: /*@C
    382:    DMDASNESSetPicardLocal - set a local right hand side and matrix evaluation function for Picard iteration

    384:    Logically Collective

    386:    Input Parameters:
    387: +  dm - DM to associate callback with
    388: .  imode - INSERT_VALUES if local function computes owned part, ADD_VALUES if it contributes to ghosted part
    389: .  func - local residual evaluation
    390: -  ctx - optional context for local residual evaluation

    392:    Calling sequence for func:
    393: +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
    394: .  x - dimensional pointer to state at which to evaluate residual
    395: .  f - dimensional pointer to residual, write the residual here
    396: -  ctx - optional context passed above

    398:    Notes:
    399:     The user must use
    400:     SNESSetFunction(snes,NULL,SNESPicardComputeFunction,&user);
    401:     in their code before calling this routine.

    403:    Level: beginner

    405: .seealso: DMSNESSetFunction(), DMDASNESSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
    406: @*/
    407: PetscErrorCode DMDASNESSetPicardLocal(DM dm,InsertMode imode,PetscErrorCode (*func)(DMDALocalInfo*,void*,void*,void*),
    408:                                       PetscErrorCode (*jac)(DMDALocalInfo*,void*,Mat,Mat,void*),void *ctx)
    409: {
    410:   DMSNES         sdm;
    411:   DMSNES_DA      *dmdasnes;

    414:   DMGetDMSNESWrite(dm,&sdm);
    415:   DMDASNESGetContext(dm,sdm,&dmdasnes);

    417:   dmdasnes->residuallocalimode = imode;
    418:   dmdasnes->rhsplocal          = func;
    419:   dmdasnes->jacobianplocal     = jac;
    420:   dmdasnes->picardlocalctx     = ctx;

    422:   DMSNESSetPicard(dm,SNESComputePicard_DMDA,SNESComputePicardJacobian_DMDA,dmdasnes);
    423:   DMSNESSetMFFunction(dm,SNESComputeFunction_DMDA,dmdasnes);
    424:   return 0;
    425: }