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

Actual source code: adamat.c

      1: #include <petscmat.h>

      3: PETSC_INTERN PetscErrorCode MatCreateADA(Mat,Vec, Vec, Mat*);

      5: typedef struct {
      6:   Mat      A;
      7:   Vec      D1;
      8:   Vec      D2;
      9:   Vec      W;
     10:   Vec      W2;
     11:   Vec      ADADiag;
     12:   PetscInt GotDiag;
     13: } _p_TaoMatADACtx;
     14: typedef  _p_TaoMatADACtx* TaoMatADACtx;

     16: static PetscErrorCode MatMult_ADA(Mat mat,Vec a,Vec y)
     17: {
     18:   TaoMatADACtx   ctx;
     19:   PetscReal      one = 1.0;

     21:   MatShellGetContext(mat,&ctx);
     22:   MatMult(ctx->A,a,ctx->W);
     23:   if (ctx->D1) {
     24:     VecPointwiseMult(ctx->W,ctx->D1,ctx->W);
     25:   }
     26:   MatMultTranspose(ctx->A,ctx->W,y);
     27:   if (ctx->D2) {
     28:     VecPointwiseMult(ctx->W2, ctx->D2, a);
     29:     VecAXPY(y, one, ctx->W2);
     30:   }
     31:   return 0;
     32: }

     34: static PetscErrorCode MatMultTranspose_ADA(Mat mat,Vec a,Vec y)
     35: {
     36:   MatMult_ADA(mat,a,y);
     37:   return 0;
     38: }

     40: static PetscErrorCode MatDiagonalSet_ADA(Mat M,Vec D, InsertMode mode)
     41: {
     42:   TaoMatADACtx   ctx;
     43:   PetscReal      zero=0.0,one = 1.0;

     46:   MatShellGetContext(M,&ctx);
     47:   if (!ctx->D2) {
     48:     VecDuplicate(D,&ctx->D2);
     49:     VecSet(ctx->D2, zero);
     50:   }
     51:   VecAXPY(ctx->D2, one, D);
     52:   return 0;
     53: }

     55: static PetscErrorCode MatDestroy_ADA(Mat mat)
     56: {
     57:   TaoMatADACtx   ctx;

     59:   MatShellGetContext(mat,&ctx);
     60:   VecDestroy(&ctx->W);
     61:   VecDestroy(&ctx->W2);
     62:   VecDestroy(&ctx->ADADiag);
     63:   MatDestroy(&ctx->A);
     64:   VecDestroy(&ctx->D1);
     65:   VecDestroy(&ctx->D2);
     66:   PetscFree(ctx);
     67:   return 0;
     68: }

     70: static PetscErrorCode MatView_ADA(Mat mat,PetscViewer viewer)
     71: {
     72:   return 0;
     73: }

     75: static PetscErrorCode MatShift_ADA(Mat Y, PetscReal a)
     76: {
     77:   TaoMatADACtx   ctx;

     79:   MatShellGetContext(Y,&ctx);
     80:   VecShift(ctx->D2,a);
     81:   return 0;
     82: }

     84: static PetscErrorCode MatDuplicate_ADA(Mat mat,MatDuplicateOption op,Mat *M)
     85: {
     86:   TaoMatADACtx      ctx;
     87:   Mat               A2;
     88:   Vec               D1b=NULL,D2b;

     90:   MatShellGetContext(mat,&ctx);
     91:   MatDuplicate(ctx->A,op,&A2);
     92:   if (ctx->D1) {
     93:     VecDuplicate(ctx->D1,&D1b);
     94:     VecCopy(ctx->D1,D1b);
     95:   }
     96:   VecDuplicate(ctx->D2,&D2b);
     97:   VecCopy(ctx->D2,D2b);
     98:   MatCreateADA(A2,D1b,D2b,M);
     99:   if (ctx->D1) {
    100:     PetscObjectDereference((PetscObject)D1b);
    101:   }
    102:   PetscObjectDereference((PetscObject)D2b);
    103:   PetscObjectDereference((PetscObject)A2);
    104:   return 0;
    105: }

    107: static PetscErrorCode MatEqual_ADA(Mat A,Mat B,PetscBool *flg)
    108: {
    109:   TaoMatADACtx   ctx1,ctx2;

    111:   MatShellGetContext(A,&ctx1);
    112:   MatShellGetContext(B,&ctx2);
    113:   VecEqual(ctx1->D2,ctx2->D2,flg);
    114:   if (*flg==PETSC_TRUE) {
    115:     VecEqual(ctx1->D1,ctx2->D1,flg);
    116:   }
    117:   if (*flg==PETSC_TRUE) {
    118:     MatEqual(ctx1->A,ctx2->A,flg);
    119:   }
    120:   return 0;
    121: }

    123: static PetscErrorCode MatScale_ADA(Mat mat, PetscReal a)
    124: {
    125:   TaoMatADACtx   ctx;

    127:   MatShellGetContext(mat,&ctx);
    128:   VecScale(ctx->D1,a);
    129:   if (ctx->D2) {
    130:     VecScale(ctx->D2,a);
    131:   }
    132:   return 0;
    133: }

    135: static PetscErrorCode MatTranspose_ADA(Mat mat,MatReuse reuse,Mat *B)
    136: {
    137:   TaoMatADACtx   ctx;

    139:   MatShellGetContext(mat,&ctx);
    140:   if (reuse == MAT_INITIAL_MATRIX) {
    141:     MatDuplicate(mat,MAT_COPY_VALUES,B);
    142:   } else if (reuse == MAT_REUSE_MATRIX) {
    143:     MatCopy(mat,*B,SAME_NONZERO_PATTERN);
    144:   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"Does not support inplace transpose");
    145:   return 0;
    146: }

    148: static PetscErrorCode MatADAComputeDiagonal(Mat mat)
    149: {
    150:   PetscInt       i,m,n,low,high;
    151:   PetscScalar    *dtemp,*dptr;
    152:   TaoMatADACtx   ctx;

    154:   MatShellGetContext(mat,&ctx);
    155:   MatGetOwnershipRange(mat, &low, &high);
    156:   MatGetSize(mat,&m,&n);

    158:   PetscMalloc1(n,&dtemp);
    159:   for (i=0; i<n; i++) {
    160:     MatGetColumnVector(ctx->A, ctx->W, i);
    161:     VecPointwiseMult(ctx->W,ctx->W,ctx->W);
    162:     VecDotBegin(ctx->D1, ctx->W,dtemp+i);
    163:   }
    164:   for (i=0; i<n; i++) {
    165:     VecDotEnd(ctx->D1, ctx->W,dtemp+i);
    166:   }

    168:   VecGetArray(ctx->ADADiag,&dptr);
    169:   for (i=low; i<high; i++) {
    170:     dptr[i-low]= dtemp[i];
    171:   }
    172:   VecRestoreArray(ctx->ADADiag,&dptr);
    173:   PetscFree(dtemp);
    174:   return 0;
    175: }

    177: static PetscErrorCode MatGetDiagonal_ADA(Mat mat,Vec v)
    178: {
    179:   PetscReal       one=1.0;
    180:   TaoMatADACtx    ctx;

    182:   MatShellGetContext(mat,&ctx);
    183:   MatADAComputeDiagonal(mat);
    184:   VecCopy(ctx->ADADiag,v);
    185:   if (ctx->D2) {
    186:     VecAXPY(v, one, ctx->D2);
    187:   }
    188:   return 0;
    189: }

    191: static PetscErrorCode MatCreateSubMatrix_ADA(Mat mat,IS isrow,IS iscol,MatReuse cll, Mat *newmat)
    192: {
    193:   PetscInt          low,high;
    194:   IS                ISrow;
    195:   Vec               D1,D2;
    196:   Mat               Atemp;
    197:   TaoMatADACtx      ctx;
    198:   PetscBool         isequal;

    200:   ISEqual(isrow,iscol,&isequal);
    202:   MatShellGetContext(mat,&ctx);

    204:   MatGetOwnershipRange(ctx->A,&low,&high);
    205:   ISCreateStride(PetscObjectComm((PetscObject)mat),high-low,low,1,&ISrow);
    206:   MatCreateSubMatrix(ctx->A,ISrow,iscol,cll,&Atemp);
    207:   ISDestroy(&ISrow);

    209:   if (ctx->D1) {
    210:     VecDuplicate(ctx->D1,&D1);
    211:     VecCopy(ctx->D1,D1);
    212:   } else {
    213:     D1 = NULL;
    214:   }

    216:   if (ctx->D2) {
    217:     Vec D2sub;

    219:     VecGetSubVector(ctx->D2,isrow,&D2sub);
    220:     VecDuplicate(D2sub,&D2);
    221:     VecCopy(D2sub,D2);
    222:     VecRestoreSubVector(ctx->D2,isrow,&D2sub);
    223:   } else {
    224:     D2 = NULL;
    225:   }

    227:   MatCreateADA(Atemp,D1,D2,newmat);
    228:   MatShellGetContext(*newmat,&ctx);
    229:   PetscObjectDereference((PetscObject)Atemp);
    230:   if (ctx->D1) {
    231:     PetscObjectDereference((PetscObject)D1);
    232:   }
    233:   if (ctx->D2) {
    234:     PetscObjectDereference((PetscObject)D2);
    235:   }
    236:   return 0;
    237: }

    239: static PetscErrorCode MatCreateSubMatrices_ADA(Mat A,PetscInt n, IS *irow,IS *icol,MatReuse scall,Mat **B)
    240: {
    241:   PetscInt       i;

    243:   if (scall == MAT_INITIAL_MATRIX) {
    244:     PetscCalloc1(n+1,B);
    245:   }
    246:   for (i=0; i<n; i++) {
    247:     MatCreateSubMatrix_ADA(A,irow[i],icol[i],scall,&(*B)[i]);
    248:   }
    249:   return 0;
    250: }

    252: static PetscErrorCode MatGetColumnVector_ADA(Mat mat,Vec Y, PetscInt col)
    253: {
    254:   PetscInt       low,high;
    255:   PetscScalar    zero=0.0,one=1.0;

    257:   VecSet(Y, zero);
    258:   VecGetOwnershipRange(Y,&low,&high);
    259:   if (col>=low && col<high) {
    260:     VecSetValue(Y,col,one,INSERT_VALUES);
    261:   }
    262:   VecAssemblyBegin(Y);
    263:   VecAssemblyEnd(Y);
    264:   MatMult_ADA(mat,Y,Y);
    265:   return 0;
    266: }

    268: PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat,MatType newtype,Mat *NewMat)
    269: {
    270:   PetscMPIInt    size;
    271:   PetscBool      sametype, issame, isdense, isseqdense;
    272:   TaoMatADACtx   ctx;

    274:   MatShellGetContext(mat,&ctx);
    275:   MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);

    277:   PetscObjectTypeCompare((PetscObject)mat,newtype,&sametype);
    278:   PetscObjectTypeCompare((PetscObject)mat,MATSAME,&issame);
    279:   PetscObjectTypeCompare((PetscObject)mat,MATMPIDENSE,&isdense);
    280:   PetscObjectTypeCompare((PetscObject)mat,MATSEQDENSE,&isseqdense);

    282:   if (sametype || issame) {
    283:     MatDuplicate(mat,MAT_COPY_VALUES,NewMat);
    284:   } else if (isdense) {
    285:     PetscInt          i,j,low,high,m,n,M,N;
    286:     const PetscScalar *dptr;
    287:     Vec               X;

    289:     VecDuplicate(ctx->D2,&X);
    290:     MatGetSize(mat,&M,&N);
    291:     MatGetLocalSize(mat,&m,&n);
    292:     MatCreateDense(PetscObjectComm((PetscObject)mat),m,m,N,N,NULL,NewMat);
    293:     MatGetOwnershipRange(*NewMat,&low,&high);
    294:     for (i=0;i<M;i++) {
    295:       MatGetColumnVector_ADA(mat,X,i);
    296:       VecGetArrayRead(X,&dptr);
    297:       for (j=0; j<high-low; j++) {
    298:         MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES);
    299:       }
    300:       VecRestoreArrayRead(X,&dptr);
    301:     }
    302:     MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY);
    303:     MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY);
    304:     VecDestroy(&X);
    305:   } else if (isseqdense && size==1) {
    306:     PetscInt          i,j,low,high,m,n,M,N;
    307:     const PetscScalar *dptr;
    308:     Vec               X;

    310:     VecDuplicate(ctx->D2,&X);
    311:     MatGetSize(mat,&M,&N);
    312:     MatGetLocalSize(mat,&m,&n);
    313:     MatCreateSeqDense(PetscObjectComm((PetscObject)mat),N,N,NULL,NewMat);
    314:     MatGetOwnershipRange(*NewMat,&low,&high);
    315:     for (i=0;i<M;i++) {
    316:       MatGetColumnVector_ADA(mat,X,i);
    317:       VecGetArrayRead(X,&dptr);
    318:       for (j=0; j<high-low; j++) {
    319:         MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES);
    320:       }
    321:       VecRestoreArrayRead(X,&dptr);
    322:     }
    323:     MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY);
    324:     MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY);
    325:     VecDestroy(&X);
    326:   } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"No support to convert objects to that type");
    327:   return 0;
    328: }

    330: static PetscErrorCode MatNorm_ADA(Mat mat,NormType type,PetscReal *norm)
    331: {
    332:   TaoMatADACtx   ctx;

    334:   MatShellGetContext(mat,&ctx);
    335:   if (type == NORM_FROBENIUS) {
    336:     *norm = 1.0;
    337:   } else if (type == NORM_1 || type == NORM_INFINITY) {
    338:     *norm = 1.0;
    339:   } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No two norm");
    340:   return 0;
    341: }

    343: /*@C
    344:    MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal

    346:    Collective on matrix

    348:    Input Parameters:
    349: +  mat - matrix of arbitrary type
    350: .  d1 - A vector defining a diagonal matrix
    351: -  d2 - A vector defining a diagonal matrix

    353:    Output Parameters:
    354: .  J - New matrix whose operations are defined in terms of mat, D1, and D2.

    356:    Notes:
    357:    The user provides the input data and is responsible for destroying
    358:    this data after matrix J has been destroyed.

    360:    Level: developer

    362: .seealso: MatCreate()
    363: @*/
    364: PetscErrorCode MatCreateADA(Mat mat,Vec d1, Vec d2, Mat *J)
    365: {
    366:   MPI_Comm       comm = PetscObjectComm((PetscObject)mat);
    367:   TaoMatADACtx   ctx;
    368:   PetscInt       nloc,n;

    370:   PetscNew(&ctx);
    371:   ctx->A=mat;
    372:   ctx->D1=d1;
    373:   ctx->D2=d2;
    374:   if (d1) {
    375:     VecDuplicate(d1,&ctx->W);
    376:     PetscObjectReference((PetscObject)d1);
    377:   } else {
    378:     ctx->W = NULL;
    379:   }
    380:   if (d2) {
    381:     VecDuplicate(d2,&ctx->W2);
    382:     VecDuplicate(d2,&ctx->ADADiag);
    383:     PetscObjectReference((PetscObject)d2);
    384:   } else {
    385:     ctx->W2      = NULL;
    386:     ctx->ADADiag = NULL;
    387:   }

    389:   ctx->GotDiag = 0;
    390:   PetscObjectReference((PetscObject)mat);

    392:   VecGetLocalSize(d2,&nloc);
    393:   VecGetSize(d2,&n);

    395:   MatCreateShell(comm,nloc,nloc,n,n,ctx,J);
    396:   MatShellSetManageScalingShifts(*J);
    397:   MatShellSetOperation(*J,MATOP_MULT,(void(*)(void))MatMult_ADA);
    398:   MatShellSetOperation(*J,MATOP_DESTROY,(void(*)(void))MatDestroy_ADA);
    399:   MatShellSetOperation(*J,MATOP_VIEW,(void(*)(void))MatView_ADA);
    400:   MatShellSetOperation(*J,MATOP_MULT_TRANSPOSE,(void(*)(void))MatMultTranspose_ADA);
    401:   MatShellSetOperation(*J,MATOP_DIAGONAL_SET,(void(*)(void))MatDiagonalSet_ADA);
    402:   MatShellSetOperation(*J,MATOP_SHIFT,(void(*)(void))MatShift_ADA);
    403:   MatShellSetOperation(*J,MATOP_EQUAL,(void(*)(void))MatEqual_ADA);
    404:   MatShellSetOperation(*J,MATOP_SCALE,(void(*)(void))MatScale_ADA);
    405:   MatShellSetOperation(*J,MATOP_TRANSPOSE,(void(*)(void))MatTranspose_ADA);
    406:   MatShellSetOperation(*J,MATOP_GET_DIAGONAL,(void(*)(void))MatGetDiagonal_ADA);
    407:   MatShellSetOperation(*J,MATOP_CREATE_SUBMATRICES,(void(*)(void))MatCreateSubMatrices_ADA);
    408:   MatShellSetOperation(*J,MATOP_NORM,(void(*)(void))MatNorm_ADA);
    409:   MatShellSetOperation(*J,MATOP_DUPLICATE,(void(*)(void))MatDuplicate_ADA);
    410:   MatShellSetOperation(*J,MATOP_CREATE_SUBMATRIX,(void(*)(void))MatCreateSubMatrix_ADA);

    412:   PetscLogObjectParent((PetscObject)(*J),(PetscObject)ctx->W);
    413:   PetscLogObjectParent((PetscObject)mat,(PetscObject)(*J));

    415:   MatSetOption(*J,MAT_SYMMETRIC,PETSC_TRUE);
    416:   return 0;
    417: }