Date: 2025-03-28T21:11:15+00:00

Actual source code: adamat.c

      1: #include <petsc/private/matimpl.h>

      3: static 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:   PetscFunctionBegin;
     22:   PetscCall(MatShellGetContext(mat, &ctx));
     23:   PetscCall(MatMult(ctx->A, a, ctx->W));
     24:   if (ctx->D1) PetscCall(VecPointwiseMult(ctx->W, ctx->D1, ctx->W));
     25:   PetscCall(MatMultTranspose(ctx->A, ctx->W, y));
     26:   if (ctx->D2) {
     27:     PetscCall(VecPointwiseMult(ctx->W2, ctx->D2, a));
     28:     PetscCall(VecAXPY(y, one, ctx->W2));
     29:   }
     30:   PetscFunctionReturn(PETSC_SUCCESS);
     31: }

     33: static PetscErrorCode MatMultTranspose_ADA(Mat mat, Vec a, Vec y)
     34: {
     35:   PetscFunctionBegin;
     36:   PetscCall(MatMult_ADA(mat, a, y));
     37:   PetscFunctionReturn(PETSC_SUCCESS);
     38: }

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

     45:   PetscFunctionBegin;
     46:   PetscCheck(mode != INSERT_VALUES, PetscObjectComm((PetscObject)M), PETSC_ERR_SUP, "Cannot insert diagonal entries of this matrix type, can only add");
     47:   PetscCall(MatShellGetContext(M, &ctx));
     48:   if (!ctx->D2) {
     49:     PetscCall(VecDuplicate(D, &ctx->D2));
     50:     PetscCall(VecSet(ctx->D2, zero));
     51:   }
     52:   PetscCall(VecAXPY(ctx->D2, one, D));
     53:   PetscFunctionReturn(PETSC_SUCCESS);
     54: }

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

     60:   PetscFunctionBegin;
     61:   PetscCall(MatShellGetContext(mat, &ctx));
     62:   PetscCall(VecDestroy(&ctx->W));
     63:   PetscCall(VecDestroy(&ctx->W2));
     64:   PetscCall(VecDestroy(&ctx->ADADiag));
     65:   PetscCall(MatDestroy(&ctx->A));
     66:   PetscCall(VecDestroy(&ctx->D1));
     67:   PetscCall(VecDestroy(&ctx->D2));
     68:   PetscCall(PetscFree(ctx));
     69:   PetscFunctionReturn(PETSC_SUCCESS);
     70: }

     72: static PetscErrorCode MatView_ADA(Mat mat, PetscViewer viewer)
     73: {
     74:   PetscFunctionBegin;
     75:   PetscFunctionReturn(PETSC_SUCCESS);
     76: }

     78: static PetscErrorCode MatShift_ADA(Mat Y, PetscReal a)
     79: {
     80:   TaoMatADACtx ctx;

     82:   PetscFunctionBegin;
     83:   PetscCall(MatShellGetContext(Y, &ctx));
     84:   PetscCall(VecShift(ctx->D2, a));
     85:   PetscFunctionReturn(PETSC_SUCCESS);
     86: }

     88: static PetscErrorCode MatDuplicate_ADA(Mat mat, MatDuplicateOption op, Mat *M)
     89: {
     90:   TaoMatADACtx ctx;
     91:   Mat          A2;
     92:   Vec          D1b = NULL, D2b;

     94:   PetscFunctionBegin;
     95:   PetscCall(MatShellGetContext(mat, &ctx));
     96:   PetscCall(MatDuplicate(ctx->A, op, &A2));
     97:   if (ctx->D1) {
     98:     PetscCall(VecDuplicate(ctx->D1, &D1b));
     99:     PetscCall(VecCopy(ctx->D1, D1b));
    100:   }
    101:   PetscCall(VecDuplicate(ctx->D2, &D2b));
    102:   PetscCall(VecCopy(ctx->D2, D2b));
    103:   PetscCall(MatCreateADA(A2, D1b, D2b, M));
    104:   if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1b));
    105:   PetscCall(PetscObjectDereference((PetscObject)D2b));
    106:   PetscCall(PetscObjectDereference((PetscObject)A2));
    107:   PetscFunctionReturn(PETSC_SUCCESS);
    108: }

    110: static PetscErrorCode MatEqual_ADA(Mat A, Mat B, PetscBool *flg)
    111: {
    112:   TaoMatADACtx ctx1, ctx2;

    114:   PetscFunctionBegin;
    115:   PetscCall(MatShellGetContext(A, &ctx1));
    116:   PetscCall(MatShellGetContext(B, &ctx2));
    117:   PetscCall(VecEqual(ctx1->D2, ctx2->D2, flg));
    118:   if (*flg == PETSC_TRUE) PetscCall(VecEqual(ctx1->D1, ctx2->D1, flg));
    119:   if (*flg == PETSC_TRUE) PetscCall(MatEqual(ctx1->A, ctx2->A, flg));
    120:   PetscFunctionReturn(PETSC_SUCCESS);
    121: }

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

    127:   PetscFunctionBegin;
    128:   PetscCall(MatShellGetContext(mat, &ctx));
    129:   PetscCall(VecScale(ctx->D1, a));
    130:   if (ctx->D2) PetscCall(VecScale(ctx->D2, a));
    131:   PetscFunctionReturn(PETSC_SUCCESS);
    132: }

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

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

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

    155:   PetscFunctionBegin;
    156:   PetscCall(MatShellGetContext(mat, &ctx));
    157:   PetscCall(MatGetOwnershipRange(mat, &low, &high));
    158:   PetscCall(MatGetSize(mat, &m, &n));

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

    168:   PetscCall(VecGetArray(ctx->ADADiag, &dptr));
    169:   for (i = low; i < high; i++) dptr[i - low] = dtemp[i];
    170:   PetscCall(VecRestoreArray(ctx->ADADiag, &dptr));
    171:   PetscCall(PetscFree(dtemp));
    172:   PetscFunctionReturn(PETSC_SUCCESS);
    173: }

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

    180:   PetscFunctionBegin;
    181:   PetscCall(MatShellGetContext(mat, &ctx));
    182:   PetscCall(MatADAComputeDiagonal(mat));
    183:   PetscCall(VecCopy(ctx->ADADiag, v));
    184:   if (ctx->D2) PetscCall(VecAXPY(v, one, ctx->D2));
    185:   PetscFunctionReturn(PETSC_SUCCESS);
    186: }

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

    197:   PetscFunctionBegin;
    198:   PetscCall(ISEqual(isrow, iscol, &isequal));
    199:   PetscCheck(isequal, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only for identical column and row indices");
    200:   PetscCall(MatShellGetContext(mat, &ctx));

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

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

    214:   if (ctx->D2) {
    215:     Vec D2sub;

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

    225:   PetscCall(MatCreateADA(Atemp, D1, D2, newmat));
    226:   PetscCall(MatShellGetContext(*newmat, &ctx));
    227:   PetscCall(PetscObjectDereference((PetscObject)Atemp));
    228:   if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1));
    229:   if (ctx->D2) PetscCall(PetscObjectDereference((PetscObject)D2));
    230:   PetscFunctionReturn(PETSC_SUCCESS);
    231: }

    233: static PetscErrorCode MatCreateSubMatrices_ADA(Mat A, PetscInt n, IS *irow, IS *icol, MatReuse scall, Mat **B)
    234: {
    235:   PetscInt i;

    237:   PetscFunctionBegin;
    238:   if (scall == MAT_INITIAL_MATRIX) PetscCall(PetscCalloc1(n + 1, B));
    239:   for (i = 0; i < n; i++) PetscCall(MatCreateSubMatrix_ADA(A, irow[i], icol[i], scall, &(*B)[i]));
    240:   PetscFunctionReturn(PETSC_SUCCESS);
    241: }

    243: static PetscErrorCode MatGetColumnVector_ADA(Mat mat, Vec Y, PetscInt col)
    244: {
    245:   PetscInt    low, high;
    246:   PetscScalar zero = 0.0, one = 1.0;

    248:   PetscFunctionBegin;
    249:   PetscCall(VecSet(Y, zero));
    250:   PetscCall(VecGetOwnershipRange(Y, &low, &high));
    251:   if (col >= low && col < high) PetscCall(VecSetValue(Y, col, one, INSERT_VALUES));
    252:   PetscCall(VecAssemblyBegin(Y));
    253:   PetscCall(VecAssemblyEnd(Y));
    254:   PetscCall(MatMult_ADA(mat, Y, Y));
    255:   PetscFunctionReturn(PETSC_SUCCESS);
    256: }

    258: PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat, MatType newtype, Mat *NewMat)
    259: {
    260:   PetscMPIInt  size;
    261:   PetscBool    sametype, issame, isdense, isseqdense;
    262:   TaoMatADACtx ctx;

    264:   PetscFunctionBegin;
    265:   PetscCall(MatShellGetContext(mat, &ctx));
    266:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));

    268:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, newtype, &sametype));
    269:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSAME, &issame));
    270:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSE, &isdense));
    271:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSEQDENSE, &isseqdense));

    273:   if (sametype || issame) {
    274:     PetscCall(MatDuplicate(mat, MAT_COPY_VALUES, NewMat));
    275:   } else if (isdense) {
    276:     PetscInt           i, j, low, high, m, n, M, N;
    277:     const PetscScalar *dptr;
    278:     Vec                X;

    280:     PetscCall(VecDuplicate(ctx->D2, &X));
    281:     PetscCall(MatGetSize(mat, &M, &N));
    282:     PetscCall(MatGetLocalSize(mat, &m, &n));
    283:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)mat), m, m, N, N, NULL, NewMat));
    284:     PetscCall(MatGetOwnershipRange(*NewMat, &low, &high));
    285:     for (i = 0; i < M; i++) {
    286:       PetscCall(MatGetColumnVector_ADA(mat, X, i));
    287:       PetscCall(VecGetArrayRead(X, &dptr));
    288:       for (j = 0; j < high - low; j++) PetscCall(MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES));
    289:       PetscCall(VecRestoreArrayRead(X, &dptr));
    290:     }
    291:     PetscCall(MatAssemblyBegin(*NewMat, MAT_FINAL_ASSEMBLY));
    292:     PetscCall(MatAssemblyEnd(*NewMat, MAT_FINAL_ASSEMBLY));
    293:     PetscCall(VecDestroy(&X));
    294:   } else if (isseqdense && size == 1) {
    295:     PetscInt           i, j, low, high, m, n, M, N;
    296:     const PetscScalar *dptr;
    297:     Vec                X;

    299:     PetscCall(VecDuplicate(ctx->D2, &X));
    300:     PetscCall(MatGetSize(mat, &M, &N));
    301:     PetscCall(MatGetLocalSize(mat, &m, &n));
    302:     PetscCall(MatCreateSeqDense(PetscObjectComm((PetscObject)mat), N, N, NULL, NewMat));
    303:     PetscCall(MatGetOwnershipRange(*NewMat, &low, &high));
    304:     for (i = 0; i < M; i++) {
    305:       PetscCall(MatGetColumnVector_ADA(mat, X, i));
    306:       PetscCall(VecGetArrayRead(X, &dptr));
    307:       for (j = 0; j < high - low; j++) PetscCall(MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES));
    308:       PetscCall(VecRestoreArrayRead(X, &dptr));
    309:     }
    310:     PetscCall(MatAssemblyBegin(*NewMat, MAT_FINAL_ASSEMBLY));
    311:     PetscCall(MatAssemblyEnd(*NewMat, MAT_FINAL_ASSEMBLY));
    312:     PetscCall(VecDestroy(&X));
    313:   } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "No support to convert objects to that type");
    314:   PetscFunctionReturn(PETSC_SUCCESS);
    315: }

    317: static PetscErrorCode MatNorm_ADA(Mat mat, NormType type, PetscReal *norm)
    318: {
    319:   TaoMatADACtx ctx;

    321:   PetscFunctionBegin;
    322:   PetscCall(MatShellGetContext(mat, &ctx));
    323:   if (type == NORM_FROBENIUS) {
    324:     *norm = 1.0;
    325:   } else if (type == NORM_1 || type == NORM_INFINITY) {
    326:     *norm = 1.0;
    327:   } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No two norm");
    328:   PetscFunctionReturn(PETSC_SUCCESS);
    329: }

    331: /*
    332:    MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal

    334:    Collective

    336:    Input Parameters:
    337: +  mat - matrix of arbitrary type
    338: .  d1 - A vector defining a diagonal matrix
    339: -  d2 - A vector defining a diagonal matrix

    341:    Output Parameter:
    342: .  J - New matrix whose operations are defined in terms of mat, D1, and D2.

    344:    Level: developer

    346:    Note:
    347:    The user provides the input data and is responsible for destroying
    348:    this data after matrix `J` has been destroyed.

    350: .seealso: `Mat`, `MatCreate()`
    351: */
    352: static PetscErrorCode MatCreateADA(Mat mat, Vec d1, Vec d2, Mat *J)
    353: {
    354:   MPI_Comm     comm = PetscObjectComm((PetscObject)mat);
    355:   TaoMatADACtx ctx;
    356:   PetscInt     nloc, n;

    358:   PetscFunctionBegin;
    359:   PetscCall(PetscNew(&ctx));
    360:   ctx->A  = mat;
    361:   ctx->D1 = d1;
    362:   ctx->D2 = d2;
    363:   if (d1) {
    364:     PetscCall(VecDuplicate(d1, &ctx->W));
    365:     PetscCall(PetscObjectReference((PetscObject)d1));
    366:   } else {
    367:     ctx->W = NULL;
    368:   }
    369:   if (d2) {
    370:     PetscCall(VecDuplicate(d2, &ctx->W2));
    371:     PetscCall(VecDuplicate(d2, &ctx->ADADiag));
    372:     PetscCall(PetscObjectReference((PetscObject)d2));
    373:   } else {
    374:     ctx->W2      = NULL;
    375:     ctx->ADADiag = NULL;
    376:   }

    378:   ctx->GotDiag = 0;
    379:   PetscCall(PetscObjectReference((PetscObject)mat));

    381:   PetscCall(VecGetLocalSize(d2, &nloc));
    382:   PetscCall(VecGetSize(d2, &n));

    384:   PetscCall(MatCreateShell(comm, nloc, nloc, n, n, ctx, J));
    385:   PetscCall(MatShellSetManageScalingShifts(*J));
    386:   PetscCall(MatShellSetOperation(*J, MATOP_MULT, (void (*)(void))MatMult_ADA));
    387:   PetscCall(MatShellSetOperation(*J, MATOP_DESTROY, (void (*)(void))MatDestroy_ADA));
    388:   PetscCall(MatShellSetOperation(*J, MATOP_VIEW, (void (*)(void))MatView_ADA));
    389:   PetscCall(MatShellSetOperation(*J, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_ADA));
    390:   PetscCall(MatShellSetOperation(*J, MATOP_DIAGONAL_SET, (void (*)(void))MatDiagonalSet_ADA));
    391:   PetscCall(MatShellSetOperation(*J, MATOP_SHIFT, (void (*)(void))MatShift_ADA));
    392:   PetscCall(MatShellSetOperation(*J, MATOP_EQUAL, (void (*)(void))MatEqual_ADA));
    393:   PetscCall(MatShellSetOperation(*J, MATOP_SCALE, (void (*)(void))MatScale_ADA));
    394:   PetscCall(MatShellSetOperation(*J, MATOP_TRANSPOSE, (void (*)(void))MatTranspose_ADA));
    395:   PetscCall(MatShellSetOperation(*J, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiagonal_ADA));
    396:   PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRICES, (void (*)(void))MatCreateSubMatrices_ADA));
    397:   PetscCall(MatShellSetOperation(*J, MATOP_NORM, (void (*)(void))MatNorm_ADA));
    398:   PetscCall(MatShellSetOperation(*J, MATOP_DUPLICATE, (void (*)(void))MatDuplicate_ADA));
    399:   PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRIX, (void (*)(void))MatCreateSubMatrix_ADA));

    401:   PetscCall(MatSetOption(*J, MAT_SYMMETRIC, PETSC_TRUE));
    402:   PetscFunctionReturn(PETSC_SUCCESS);
    403: }