Actual source code: matadic.c

      1: #define PETSCMAT_DLL

      3: /*
      4:     ADIC matrix-free matrix implementation
      5: */

     7:  #include src/mat/matimpl.h
     8:  #include petscda.h
     9:  #include petscsnes.h
     10:  #include petscsys.h
     12: #include "adic/ad_utils.h"

     15: typedef struct {
     16:   DA         da;
     17:   Vec        localu;         /* point at which Jacobian is evaluated */
     18:   void       *ctx;
     19:   SNES       snes;
     20:   Vec        diagonal;       /* current matrix diagonal */
     21:   PetscTruth diagonalvalid;  /* indicates if diagonal matches current base vector */
     22: } Mat_DAAD;

     26: PetscErrorCode MatAssemblyEnd_DAAD(Mat A,MatAssemblyType atype)
     27: {
     28:   Mat_DAAD *a = (Mat_DAAD*)A->data;
     30:   Vec      u;

     33:   a->diagonalvalid = PETSC_FALSE;
     34:   if (a->snes) {
     35:     SNESGetSolution(a->snes,&u);
     36:     DAGlobalToLocalBegin(a->da,u,INSERT_VALUES,a->localu);
     37:     DAGlobalToLocalEnd(a->da,u,INSERT_VALUES,a->localu);
     38:   }
     39:   return(0);
     40: }

     44: PetscErrorCode MatMult_DAAD(Mat A,Vec xx,Vec yy)
     45: {
     46:   Mat_DAAD *a = (Mat_DAAD*)A->data;
     47:   Vec      localxx;

     51:   DAGetLocalVector(a->da,&localxx);
     52:   DAGlobalToLocalBegin(a->da,xx,INSERT_VALUES,localxx);
     53:   DAGlobalToLocalEnd(a->da,xx,INSERT_VALUES,localxx);
     54:   DAMultiplyByJacobian1WithAD(a->da,a->localu,localxx,yy,a->ctx);
     55:   DARestoreLocalVector(a->da,&localxx);
     56:   return(0);
     57: }

     59:  #include src/dm/da/daimpl.h

     63: PetscErrorCode MatGetDiagonal_DAAD(Mat A,Vec dd)
     64: {
     65:   Mat_DAAD      *a = (Mat_DAAD*)A->data;
     67:   int j,nI,gI,gtdof;
     68:   PetscScalar   *avu,*ad_vustart,ad_f[2],*d;
     69:   DALocalInfo   info;
     70:   MatStencil    stencil;
     71:   void*         *ad_vu;

     75:   /* get space for derivative object.  */
     76:   DAGetAdicMFArray(a->da,PETSC_TRUE,(void **)&ad_vu,(void**)&ad_vustart,&gtdof);

     78:   /* copy input vector into derivative object */
     79:   VecGetArray(a->localu,&avu);
     80:   for (j=0; j<gtdof; j++) {
     81:     ad_vustart[2*j]   = avu[j];
     82:     ad_vustart[2*j+1] = 0.0;
     83:   }
     84:   VecRestoreArray(a->localu,&avu);

     86:   PetscADResetIndep();
     87:   PetscADIncrementTotalGradSize(1);
     88:   PetscADSetIndepDone();

     90:   VecGetArray(dd,&d);

     92:   DAGetLocalInfo(a->da,&info);
     93:   nI = 0;
     94:   for (stencil.k = info.zs; stencil.k<info.zs+info.zm; stencil.k++) {
     95:     for (stencil.j = info.ys; stencil.j<info.ys+info.ym; stencil.j++) {
     96:       for (stencil.i = info.xs; stencil.i<info.xs+info.xm; stencil.i++) {
     97:         for (stencil.c = 0; stencil.c<info.dof; stencil.c++) {
     98:           gI   = stencil.c + (stencil.i - info.gxs)*info.dof + (stencil.j - info.gys)*info.dof*info.gxm + (stencil.k - info.gzs)*info.dof*info.gxm*info.gym;
     99:           ad_vustart[1+2*gI] = 1.0;
    100:           (*a->da->adicmf_lfi)(&info,&stencil,ad_vu,ad_f,a->ctx);
    101:           d[nI] = ad_f[1];
    102:           ad_vustart[1+2*gI] = 0.0;
    103:           nI++;
    104:         }
    105:       }
    106:     }
    107:   }

    109:   VecRestoreArray(dd,&d);
    110:   DARestoreAdicMFArray(a->da,PETSC_TRUE,(void **)&ad_vu,(void**)&ad_vustart,&gtdof);
    111:   return(0);
    112: }

    117: PetscErrorCode MatRelax_DAAD(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,int its,int lits,Vec xx)
    118: {
    119:   Mat_DAAD      *a = (Mat_DAAD*)A->data;
    121:   int j,gtdof,nI,gI;
    122:   PetscScalar   *avu,*av,*ad_vustart,ad_f[2],zero = 0.0,*d,*b;
    123:   Vec           localxx,dd;
    124:   DALocalInfo   info;
    125:   MatStencil    stencil;
    126:   void*         *ad_vu;

    129:   if (omega != 1.0) SETERRQ(PETSC_ERR_ARG_WRONG,"Currently only support omega of 1.0");
    130:   if (fshift)       SETERRQ(PETSC_ERR_ARG_WRONG,"Currently do not support fshift");
    131:   if (its <= 0 || lits <= 0) SETERRQ2(PETSC_ERR_ARG_WRONG,"Relaxation requires global its %D and local its %D both positive",its,lits);

    133:   if (!a->diagonal) {
    134:     DACreateGlobalVector(a->da,&a->diagonal);
    135:   }
    136:   if (!a->diagonalvalid) {
    137:     MatGetDiagonal(A,a->diagonal);
    138:     a->diagonalvalid = PETSC_TRUE;
    139:   }
    140:   dd   = a->diagonal;

    143:   DAGetLocalVector(a->da,&localxx);
    144:   if (flag & SOR_ZERO_INITIAL_GUESS) {
    145:     VecSet(localxx,zero);
    146:   } else {
    147:     DAGlobalToLocalBegin(a->da,xx,INSERT_VALUES,localxx);
    148:     DAGlobalToLocalEnd(a->da,xx,INSERT_VALUES,localxx);
    149:   }

    151:   /* get space for derivative object.  */
    152:   DAGetAdicMFArray(a->da,PETSC_TRUE,(void **)&ad_vu,(void**)&ad_vustart,&gtdof);

    154:   /* copy input vector into derivative object */
    155:   VecGetArray(a->localu,&avu);
    156:   VecGetArray(localxx,&av);
    157:   for (j=0; j<gtdof; j++) {
    158:     ad_vustart[2*j]   = avu[j];
    159:     ad_vustart[2*j+1] = av[j];
    160:   }
    161:   VecRestoreArray(a->localu,&avu);
    162:   VecRestoreArray(localxx,&av);

    164:   PetscADResetIndep();
    165:   PetscADIncrementTotalGradSize(1);
    166:   PetscADSetIndepDone();

    168:   VecGetArray(dd,&d);
    169:   VecGetArray(bb,&b);

    171:   DAGetLocalInfo(a->da,&info);
    172:   while (its--) {
    173:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP){
    174:       nI = 0;
    175:       for (stencil.k = info.zs; stencil.k<info.zs+info.zm; stencil.k++) {
    176:         for (stencil.j = info.ys; stencil.j<info.ys+info.ym; stencil.j++) {
    177:           for (stencil.i = info.xs; stencil.i<info.xs+info.xm; stencil.i++) {
    178:             for (stencil.c = 0; stencil.c<info.dof; stencil.c++) {
    179:               (*a->da->adicmf_lfi)(&info,&stencil,ad_vu,ad_f,a->ctx);
    180:               gI   = stencil.c + (stencil.i - info.gxs)*info.dof + (stencil.j - info.gys)*info.dof*info.gxm + (stencil.k - info.gzs)*info.dof*info.gxm*info.gym;
    181:               ad_vustart[1+2*gI] += (b[nI] - ad_f[1])/d[nI];
    182:               nI++;
    183:             }
    184:           }
    185:         }
    186:       }
    187:     }
    188:     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP){
    189:       nI = info.dof*info.xm*info.ym*info.zm - 1;
    190:       for (stencil.k = info.zs+info.zm-1; stencil.k>=info.zs; stencil.k--) {
    191:         for (stencil.j = info.ys+info.ym-1; stencil.j>=info.ys; stencil.j--) {
    192:           for (stencil.i = info.xs+info.xm-1; stencil.i>=info.xs; stencil.i--) {
    193:             for (stencil.c = info.dof-1; stencil.c>=0; stencil.c--) {
    194:               (*a->da->adicmf_lfi)(&info,&stencil,ad_vu,ad_f,a->ctx);
    195:               gI   = stencil.c + (stencil.i - info.gxs)*info.dof + (stencil.j - info.gys)*info.dof*info.gxm + (stencil.k - info.gzs)*info.dof*info.gxm*info.gym;
    196:               ad_vustart[1+2*gI] += (b[nI] - ad_f[1])/d[nI];
    197:               nI--;
    198:             }
    199:           }
    200:         }
    201:       }
    202:     }
    203:   }

    205:   VecRestoreArray(dd,&d);
    206:   VecRestoreArray(bb,&b);

    208:   VecGetArray(localxx,&av);
    209:   for (j=0; j<gtdof; j++) {
    210:     av[j] = ad_vustart[2*j+1];
    211:   }
    212:   VecRestoreArray(localxx,&av);
    213:   DALocalToGlobal(a->da,localxx,INSERT_VALUES,xx);

    215:   DARestoreLocalVector(a->da,&localxx);
    216:   DARestoreAdicMFArray(a->da,PETSC_TRUE,(void **)&ad_vu,(void**)&ad_vustart,&gtdof);
    217:   return(0);
    218: }

    224: PetscErrorCode MatDestroy_DAAD(Mat A)
    225: {
    226:   Mat_DAAD *a = (Mat_DAAD*)A->data;

    230:   DADestroy(a->da);
    231:   VecDestroy(a->localu);
    232:   if (a->diagonal) {VecDestroy(a->diagonal);}
    233:   PetscFree(a);
    234:   PetscObjectComposeFunction((PetscObject)A,"MatSNESMFSetBase_C","",PETSC_NULL);
    235:   PetscObjectComposeFunction((PetscObject)A,"MatDAADSetDA_C","",PETSC_NULL);
    236:   PetscObjectComposeFunction((PetscObject)A,"MatDAADSetSNES_C","",PETSC_NULL);
    237:   PetscObjectComposeFunction((PetscObject)A,"MatDAADSetCtx_C","",PETSC_NULL);
    238:   return(0);
    239: }

    241: /* -------------------------------------------------------------------*/
    242: static struct _MatOps MatOps_Values = {0,
    243:        0,
    244:        0,
    245:        MatMult_DAAD,
    246: /* 4*/ 0,
    247:        0,
    248:        0,
    249:        0,
    250:        0,
    251:        0,
    252: /*10*/ 0,
    253:        0,
    254:        0,
    255:        MatRelax_DAAD,
    256:        0,
    257: /*15*/ 0,
    258:        0,
    259:        MatGetDiagonal_DAAD,
    260:        0,
    261:        0,
    262: /*20*/ 0,
    263:        MatAssemblyEnd_DAAD,
    264:        0,
    265:        0,
    266:        0,
    267: /*25*/ 0,
    268:        0,
    269:        0,
    270:        0,
    271:        0,
    272: /*30*/ 0,
    273:        0,
    274:        0,
    275:        0,
    276:        0,
    277: /*35*/ 0,
    278:        0,
    279:        0,
    280:        0,
    281:        0,
    282: /*40*/ 0,
    283:        0,
    284:        0,
    285:        0,
    286:        0,
    287: /*45*/ 0,
    288:        0,
    289:        0,
    290:        0,
    291:        0,
    292: /*50*/ 0,
    293:        0,
    294:        0,
    295:        0,
    296:        0,
    297: /*55*/ 0,
    298:        0,
    299:        0,
    300:        0,
    301:        0,
    302: /*60*/ 0,
    303:        MatDestroy_DAAD,
    304:        0,
    305:        0,
    306:        0,
    307: /*65*/ 0,
    308:        0,
    309:        0,
    310:        0,
    311:        0,
    312: /*70*/ 0,
    313:        0,
    314:        0,
    315:        0,
    316:        0,
    317: /*75*/ 0,
    318:        0,
    319:        0,
    320:        0,
    321:        0,
    322: /*80*/ 0,
    323:        0,
    324:        0,
    325:        0,
    326:        0,
    327: /*85*/ 0,
    328:        0,
    329:        0,
    330:        0,
    331:        0,
    332: /*90*/ 0,
    333:        0,
    334:        0,
    335:        0,
    336:        0,
    337: /*95*/ 0,
    338:        0,
    339:        0,
    340:        0};

    342: /* --------------------------------------------------------------------------------*/

    347: PetscErrorCode PETSCMAT_DLLEXPORT MatSNESMFSetBase_AD(Mat J,Vec U)
    348: {
    350:   Mat_DAAD *a = (Mat_DAAD*)J->data;

    353:   a->diagonalvalid = PETSC_FALSE;
    354:   DAGlobalToLocalBegin(a->da,U,INSERT_VALUES,a->localu);
    355:   DAGlobalToLocalEnd(a->da,U,INSERT_VALUES,a->localu);
    356:   return(0);
    357: }

    363: PetscErrorCode PETSCMAT_DLLEXPORT MatDAADSetDA_AD(Mat A,DA da)
    364: {
    365:   Mat_DAAD *a = (Mat_DAAD*)A->data;
    367:   int nc,nx,ny,nz,Nx,Ny,Nz;

    370:   a->da = da;
    371:   PetscObjectReference((PetscObject)da);
    372:   DAGetInfo(da,0,&Nx,&Ny,&Nz,0,0,0,&nc,0,0,0);
    373:   DAGetCorners(da,0,0,0,&nx,&ny,&nz);
    374:   A->m  = A->n = nc*nx*ny*nz;
    375:   A->M  = A->N = nc*Nx*Ny*Nz;
    376:   DACreateLocalVector(da,&a->localu);
    377:   return(0);
    378: }

    384: PetscErrorCode PETSCMAT_DLLEXPORT MatDAADSetSNES_AD(Mat A,SNES snes)
    385: {
    386:   Mat_DAAD *a = (Mat_DAAD*)A->data;

    389:   a->snes = snes;
    390:   return(0);
    391: }

    397: PetscErrorCode PETSCMAT_DLLEXPORT MatDAADSetCtx_AD(Mat A,void *ctx)
    398: {
    399:   Mat_DAAD *a = (Mat_DAAD*)A->data;

    402:   a->ctx = ctx;
    403:   return(0);
    404: }

    407: /*MC
    408:   MATDAAD - MATDAAD = "daad" - A matrix type that can do matrix-vector products using a local function that
    409:   is differentiated with ADIFOR or ADIC.

    411:   Level: intermediate

    413: .seealso: MatCreateDAAD
    414: M*/

    419: PetscErrorCode PETSCMAT_DLLEXPORT MatCreate_DAAD(Mat B)
    420: {
    421:   Mat_DAAD *b;

    425:   PetscNew(Mat_DAAD,&b);
    426:   B->data = (void*)b;
    427:   PetscMemcpy(B->ops,&MatOps_Values,sizeof(struct _MatOps));
    428:
    429:   PetscMapCreateMPI(B->comm,B->m,B->m,&B->rmap);
    430:   PetscMapCreateMPI(B->comm,B->n,B->n,&B->cmap);

    432:   PetscObjectChangeTypeName((PetscObject)B,MATDAAD);
    433:   PetscObjectComposeFunctionDynamic((PetscObject)B,"MatSNESMFSetBase_C","MatSNESMFSetBase_AD",MatSNESMFSetBase_AD);
    434:   PetscObjectComposeFunctionDynamic((PetscObject)B,"MatDAADSetDA_C","MatDAADSetDA_AD",MatDAADSetDA_AD);
    435:   PetscObjectComposeFunctionDynamic((PetscObject)B,"MatDAADSetSNES_C","MatDAADSetSNES_AD",MatDAADSetSNES_AD);
    436:   PetscObjectComposeFunctionDynamic((PetscObject)B,"MatDAADSetCtx_C","MatDAADSetCtx_AD",MatDAADSetCtx_AD);
    437:   return(0);
    438: }

    444: /*@C
    445:    MatDAADSetDA - Tells the matrix what DA it is using for layout and Jacobian.

    447:    Collective on Mat and DA

    449:    Input Parameters:
    450: +  mat - the matrix
    451: -  da - the DA

    453:    Level: intermediate

    455: .seealso: MatCreate(), DASetLocalAdicMFFunction(), MatCreateDAAD()

    457: @*/
    458: PetscErrorCode PETSCMAT_DLLEXPORT MatDAADSetDA(Mat A,DA da)
    459: {
    460:   PetscErrorCode ierr,(*f)(Mat,void*);

    465:   PetscObjectQueryFunction((PetscObject)A,"MatDAADSetDA_C",(void (**)(void))&f);
    466:   if (f) {
    467:     (*f)(A,da);
    468:   }
    469:   return(0);
    470: }

    474: /*@C
    475:    MatDAADSetSNES - Tells the matrix what SNES it is using for the base U.

    477:    Collective on Mat and SNES

    479:    Input Parameters:
    480: +  mat - the matrix
    481: -  snes - the SNES

    483:    Level: intermediate

    485: .seealso: MatCreate(), DASetLocalAdicMFFunction(), MatCreateDAAD(), MatDAADSetDA()

    487: @*/
    488: PetscErrorCode PETSCMAT_DLLEXPORT MatDAADSetSNES(Mat A,SNES snes)
    489: {
    490:   PetscErrorCode ierr,(*f)(Mat,void*);

    495:   PetscObjectQueryFunction((PetscObject)A,"MatDAADSetSNES_C",(void (**)(void))&f);
    496:   if (f) {
    497:     (*f)(A,snes);
    498:   }
    499:   return(0);
    500: }

    504: /*@C
    505:    MatDAADSetCtx - Sets the user context for a DAAD (ADIC matrix-free) matrix.

    507:    Collective on Mat

    509:    Input Parameters:
    510: +  mat - the matrix
    511: -  ctx - the context

    513:    Level: intermediate

    515: .seealso: MatCreate(), DASetLocalAdicMFFunction(), MatCreateDAAD(), MatDAADSetDA()

    517: @*/
    518: PetscErrorCode PETSCMAT_DLLEXPORT MatDAADSetCtx(Mat A,void *ctx)
    519: {
    520:   PetscErrorCode ierr,(*f)(Mat,void*);

    524:   PetscObjectQueryFunction((PetscObject)A,"MatDAADSetCtx_C",(void (**)(void))&f);
    525:   if (f) {
    526:     (*f)(A,ctx);
    527:   }
    528:   return(0);
    529: }

    533: /*@C
    534:    MatCreateDAAD - Creates a matrix that can do matrix-vector products using a local
    535:    function that is differentiated with ADIFOR or ADIC.

    537:    Collective on DA

    539:    Input Parameters:
    540: .  da - the DA that defines the distribution of the vectors

    542:    Output Parameter:
    543: .  A - the matrix

    545:    Level: intermediate

    547: .seealso: MatCreate(), DASetLocalAdicMFFunction()

    549: @*/
    550: PetscErrorCode PETSCMAT_DLLEXPORT MatCreateDAAD(DA da,Mat *A)
    551: {
    553:   MPI_Comm comm;

    556:   PetscObjectGetComm((PetscObject)da,&comm);
    557:   MatCreate(comm,A);
    558:   MatSetType(*A,MATDAAD);
    559:   MatDAADSetDA(*A,da);
    560:   return(0);
    561: }

    565: PetscErrorCode PETSCMAT_DLLEXPORT MatRegisterDAAD(void)
    566: {
    569:   MatRegisterDynamic(MATDAAD,PETSC_NULL,"MatCreate_DAAD",MatCreate_DAAD);
    570:   return(0);
    571: }