Date: 2026-03-01T05:24:59+00:00

Actual source code: diagonal.c

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

      3: typedef struct {
      4:   Vec              diag;
      5:   PetscBool        diag_valid;
      6:   Vec              inv_diag;
      7:   PetscBool        inv_diag_valid;
      8:   PetscObjectState diag_state, inv_diag_state;
      9:   PetscInt        *col;
     10:   PetscScalar     *val;
     11: } Mat_Diagonal;

     13: static PetscErrorCode MatDiagonalSetUpDiagonal(Mat A)
     14: {
     15:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

     17:   PetscFunctionBegin;
     18:   if (!ctx->diag_valid) {
     19:     PetscAssert(ctx->inv_diag_valid, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Neither diagonal nor inverse diagonal is in a valid state");
     20:     PetscCall(VecCopy(ctx->inv_diag, ctx->diag));
     21:     PetscCall(VecReciprocal(ctx->diag));
     22:     ctx->diag_valid = PETSC_TRUE;
     23:   }
     24:   PetscFunctionReturn(PETSC_SUCCESS);
     25: }

     27: static PetscErrorCode MatDiagonalSetUpInverseDiagonal(Mat A)
     28: {
     29:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

     31:   PetscFunctionBegin;
     32:   if (!ctx->inv_diag_valid) {
     33:     PetscAssert(ctx->diag_valid, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Neither diagonal nor inverse diagonal is in a valid state");
     34:     PetscCall(VecCopy(ctx->diag, ctx->inv_diag));
     35:     PetscCall(VecReciprocal(ctx->inv_diag));
     36:     ctx->inv_diag_valid = PETSC_TRUE;
     37:   }
     38:   PetscFunctionReturn(PETSC_SUCCESS);
     39: }

     41: static PetscErrorCode MatAXPY_Diagonal(Mat Y, PetscScalar a, Mat X, MatStructure str)
     42: {
     43:   Mat_Diagonal *yctx = (Mat_Diagonal *)Y->data;
     44:   Mat_Diagonal *xctx = (Mat_Diagonal *)X->data;

     46:   PetscFunctionBegin;
     47:   PetscCall(MatDiagonalSetUpDiagonal(Y));
     48:   PetscCall(MatDiagonalSetUpDiagonal(X));
     49:   PetscCall(VecAXPY(yctx->diag, a, xctx->diag));
     50:   yctx->inv_diag_valid = PETSC_FALSE;
     51:   PetscFunctionReturn(PETSC_SUCCESS);
     52: }

     54: static PetscErrorCode MatGetRow_Diagonal(Mat A, PetscInt row, PetscInt *ncols, PetscInt **cols, PetscScalar **vals)
     55: {
     56:   Mat_Diagonal *mat    = (Mat_Diagonal *)A->data;
     57:   PetscInt      rstart = A->rmap->rstart, rend = A->rmap->rend;

     59:   PetscFunctionBegin;
     60:   PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Only local rows");
     61:   if (ncols) *ncols = 1;
     62:   if (cols) {
     63:     if (!mat->col) PetscCall(PetscMalloc1(1, &mat->col));
     64:     *mat->col = row;
     65:     *cols     = mat->col;
     66:   }
     67:   if (vals) {
     68:     const PetscScalar *v;

     70:     if (!mat->val) PetscCall(PetscMalloc1(1, &mat->val));
     71:     PetscCall(VecGetArrayRead(mat->diag, &v));
     72:     *mat->val = v[row - rstart];
     73:     *vals     = mat->val;
     74:     PetscCall(VecRestoreArrayRead(mat->diag, &v));
     75:   }
     76:   PetscFunctionReturn(PETSC_SUCCESS);
     77: }

     79: static PetscErrorCode MatMult_Diagonal(Mat A, Vec x, Vec y)
     80: {
     81:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

     83:   PetscFunctionBegin;
     84:   PetscCall(MatDiagonalSetUpDiagonal(A));
     85:   PetscCall(VecPointwiseMult(y, ctx->diag, x));
     86:   PetscFunctionReturn(PETSC_SUCCESS);
     87: }

     89: static PetscErrorCode MatMultAdd_Diagonal(Mat mat, Vec v1, Vec v2, Vec v3)
     90: {
     91:   Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data;

     93:   PetscFunctionBegin;
     94:   PetscCall(MatDiagonalSetUpDiagonal(mat));
     95:   if (v2 != v3) {
     96:     PetscCall(VecPointwiseMult(v3, ctx->diag, v1));
     97:     PetscCall(VecAXPY(v3, 1.0, v2));
     98:   } else {
     99:     Vec w;
    100:     PetscCall(VecDuplicate(v3, &w));
    101:     PetscCall(VecPointwiseMult(w, ctx->diag, v1));
    102:     PetscCall(VecAXPY(v3, 1.0, w));
    103:     PetscCall(VecDestroy(&w));
    104:   }
    105:   PetscFunctionReturn(PETSC_SUCCESS);
    106: }

    108: static PetscErrorCode MatNorm_Diagonal(Mat A, NormType type, PetscReal *nrm)
    109: {
    110:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    112:   PetscFunctionBegin;
    113:   PetscCall(MatDiagonalSetUpDiagonal(A));
    114:   type = (type == NORM_FROBENIUS) ? NORM_2 : NORM_INFINITY;
    115:   PetscCall(VecNorm(ctx->diag, type, nrm));
    116:   PetscFunctionReturn(PETSC_SUCCESS);
    117: }

    119: static PetscErrorCode MatDuplicate_Diagonal(Mat A, MatDuplicateOption op, Mat *B)
    120: {
    121:   Mat_Diagonal *actx = (Mat_Diagonal *)A->data;

    123:   PetscFunctionBegin;
    124:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
    125:   PetscCall(MatSetSizes(*B, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
    126:   PetscCall(MatSetBlockSizesFromMats(*B, A, A));
    127:   PetscCall(MatSetType(*B, MATDIAGONAL));
    128:   PetscCall(PetscLayoutReference(A->rmap, &(*B)->rmap));
    129:   PetscCall(PetscLayoutReference(A->cmap, &(*B)->cmap));
    130:   if (op == MAT_COPY_VALUES) {
    131:     Mat_Diagonal *bctx = (Mat_Diagonal *)(*B)->data;

    133:     PetscCall(MatSetUp(A));
    134:     PetscCall(MatSetUp(*B));
    135:     PetscCall(VecCopy(actx->diag, bctx->diag));
    136:     PetscCall(VecCopy(actx->inv_diag, bctx->inv_diag));
    137:     bctx->diag_valid     = actx->diag_valid;
    138:     bctx->inv_diag_valid = actx->inv_diag_valid;
    139:   }
    140:   PetscFunctionReturn(PETSC_SUCCESS);
    141: }

    143: /*@
    144:   MatDiagonalGetDiagonal - Get the diagonal of a `MATDIAGONAL`

    146:   Input Parameter:
    147: . A - the `MATDIAGONAL`

    149:   Output Parameter:
    150: . diag - the `Vec` that defines the diagonal

    152:   Level: developer

    154:   Note:
    155:   The user must call
    156:   `MatDiagonalRestoreDiagonal()` before using the matrix again.

    158:   For a copy of the diagonal values, rather than a reference, use `MatGetDiagonal()`

    160:   Any changes to the obtained vector immediately change the action of the `Mat`. The matrix can be changed more efficiently by accessing this vector and changing its values, instead of filling a work vector and using `MatDiagonalSet()`

    162: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`, `MatGetDiagonal()`
    163: @*/
    164: PetscErrorCode MatDiagonalGetDiagonal(Mat A, Vec *diag)
    165: {
    166:   PetscFunctionBegin;
    168:   PetscAssertPointer(diag, 2);
    169:   *diag = NULL;
    170:   PetscUseMethod(A, "MatDiagonalGetDiagonal_C", (Mat, Vec *), (A, diag));
    171:   PetscFunctionReturn(PETSC_SUCCESS);
    172: }

    174: static PetscErrorCode MatDiagonalGetDiagonal_Diagonal(Mat A, Vec *diag)
    175: {
    176:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    178:   PetscFunctionBegin;
    179:   PetscCall(MatSetUp(A));
    180:   PetscCall(MatDiagonalSetUpDiagonal(A));
    181:   *diag = ctx->diag;
    182:   PetscCall(PetscObjectStateGet((PetscObject)*diag, &ctx->diag_state));
    183:   PetscFunctionReturn(PETSC_SUCCESS);
    184: }

    186: /*@
    187:   MatDiagonalRestoreDiagonal - Restore the diagonal of a `MATDIAGONAL`

    189:   Input Parameters:
    190: + A    - the `MATDIAGONAL`
    191: - diag - the `Vec` obtained from `MatDiagonalGetDiagonal()`

    193:   Level: developer

    195:   Note:
    196:   Use `MatDiagonalSet()` to change the values by copy, rather than reference.

    198: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetDiagonal()`
    199: @*/
    200: PetscErrorCode MatDiagonalRestoreDiagonal(Mat A, Vec *diag)
    201: {
    202:   PetscFunctionBegin;
    204:   PetscAssertPointer(diag, 2);
    205:   PetscUseMethod(A, "MatDiagonalRestoreDiagonal_C", (Mat, Vec *), (A, diag));
    206:   PetscFunctionReturn(PETSC_SUCCESS);
    207: }

    209: static PetscErrorCode MatDiagonalRestoreDiagonal_Diagonal(Mat A, Vec *diag)
    210: {
    211:   Mat_Diagonal    *ctx = (Mat_Diagonal *)A->data;
    212:   PetscObjectState diag_state;

    214:   PetscFunctionBegin;
    215:   PetscCheck(ctx->diag == *diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector");
    216:   ctx->diag_valid = PETSC_TRUE;
    217:   PetscCall(PetscObjectStateGet((PetscObject)*diag, &diag_state));
    218:   if (ctx->diag_state != diag_state) {
    219:     PetscCall(PetscObjectStateIncrease((PetscObject)A));
    220:     ctx->inv_diag_valid = PETSC_FALSE;
    221:   }
    222:   *diag = NULL;
    223:   PetscFunctionReturn(PETSC_SUCCESS);
    224: }

    226: /*@
    227:   MatDiagonalGetInverseDiagonal - Get the inverse diagonal of a `MATDIAGONAL`

    229:   Input Parameter:
    230: . A - the `MATDIAGONAL`

    232:   Output Parameter:
    233: . inv_diag - the `Vec` that defines the inverse diagonal

    235:   Level: developer

    237:   Note:
    238:   The user must call
    239:   `MatDiagonalRestoreInverseDiagonal()` before using the matrix again.

    241:   If a matrix is created only to call `MatSolve()` (which happens for `MATLMVMDIAGBROYDEN`),
    242:   using `MatDiagonalGetInverseDiagonal()` and `MatDiagonalRestoreInverseDiagonal()` avoids copies
    243:   and avoids any call to `VecReciprocal()`.

    245: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MATLMVMBROYDEN`, `MatSolve()`
    246: @*/
    247: PetscErrorCode MatDiagonalGetInverseDiagonal(Mat A, Vec *inv_diag)
    248: {
    249:   PetscFunctionBegin;
    251:   PetscAssertPointer(inv_diag, 2);
    252:   *inv_diag = NULL;
    253:   PetscUseMethod(A, "MatDiagonalGetInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
    254:   PetscFunctionReturn(PETSC_SUCCESS);
    255: }

    257: static PetscErrorCode MatDiagonalGetInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
    258: {
    259:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    261:   PetscFunctionBegin;
    262:   PetscCall(MatSetUp(A));
    263:   PetscCall(MatDiagonalSetUpInverseDiagonal(A));
    264:   *inv_diag = ctx->inv_diag;
    265:   PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &ctx->inv_diag_state));
    266:   PetscFunctionReturn(PETSC_SUCCESS);
    267: }

    269: /*@
    270:   MatDiagonalRestoreInverseDiagonal - Restore the inverse diagonal of a `MATDIAGONAL`

    272:   Input Parameters:
    273: + A        - the `MATDIAGONAL`
    274: - inv_diag - the `Vec` obtained from `MatDiagonalGetInverseDiagonal()`

    276:   Level: developer

    278: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetInverseDiagonal()`
    279: @*/
    280: PetscErrorCode MatDiagonalRestoreInverseDiagonal(Mat A, Vec *inv_diag)
    281: {
    282:   PetscFunctionBegin;
    284:   PetscAssertPointer(inv_diag, 2);
    285:   PetscUseMethod(A, "MatDiagonalRestoreInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
    286:   PetscFunctionReturn(PETSC_SUCCESS);
    287: }

    289: static PetscErrorCode MatDiagonalRestoreInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
    290: {
    291:   Mat_Diagonal    *ctx = (Mat_Diagonal *)A->data;
    292:   PetscObjectState inv_diag_state;

    294:   PetscFunctionBegin;
    295:   PetscCheck(ctx->inv_diag == *inv_diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector");
    296:   ctx->inv_diag_valid = PETSC_TRUE;
    297:   PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &inv_diag_state));
    298:   if (ctx->inv_diag_state != inv_diag_state) {
    299:     PetscCall(PetscObjectStateIncrease((PetscObject)A));
    300:     ctx->diag_valid = PETSC_FALSE;
    301:   }
    302:   *inv_diag = NULL;
    303:   PetscFunctionReturn(PETSC_SUCCESS);
    304: }

    306: static PetscErrorCode MatPermute_Diagonal(Mat A, IS rowp, IS colp, Mat *B)
    307: {
    308:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
    309:   Vec           v;

    311:   PetscFunctionBegin;
    312:   PetscCheck(rowp == colp, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Row permutation and column permutation must be the same");
    313:   PetscCall(VecDuplicate(ctx->diag, &v));
    314:   PetscCall(VecCopy(ctx->diag, v));
    315:   PetscCall(VecPermute(v, rowp, PETSC_FALSE));
    316:   PetscCall(MatCreateDiagonal(v, B));
    317:   PetscCall(VecDestroy(&v));
    318:   PetscFunctionReturn(PETSC_SUCCESS);
    319: }

    321: static PetscErrorCode MatSetRandom_Diagonal(Mat A, PetscRandom rand)
    322: {
    323:   Vec d;

    325:   PetscFunctionBegin;
    326:   PetscCall(MatDiagonalGetDiagonal(A, &d));
    327:   PetscCall(VecSetRandom(d, rand));
    328:   PetscCall(MatDiagonalRestoreDiagonal(A, &d));
    329:   PetscFunctionReturn(PETSC_SUCCESS);
    330: }

    332: static PetscErrorCode MatDestroy_Diagonal(Mat mat)
    333: {
    334:   Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data;

    336:   PetscFunctionBegin;
    337:   PetscCall(VecDestroy(&ctx->diag));
    338:   PetscCall(VecDestroy(&ctx->inv_diag));
    339:   PetscCall(PetscFree(ctx->col));
    340:   PetscCall(PetscFree(ctx->val));
    341:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetDiagonal_C", NULL));
    342:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreDiagonal_C", NULL));
    343:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetInverseDiagonal_C", NULL));
    344:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreInverseDiagonal_C", NULL));
    345:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_seqdense_C", NULL));
    346:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_mpidense_C", NULL));
    347:   PetscCall(PetscFree(mat->data));
    348:   mat->structural_symmetry_eternal = PETSC_FALSE;
    349:   mat->symmetry_eternal            = PETSC_FALSE;
    350:   PetscFunctionReturn(PETSC_SUCCESS);
    351: }

    353: static PetscErrorCode MatView_Diagonal(Mat J, PetscViewer viewer)
    354: {
    355:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
    356:   PetscBool     isascii;

    358:   PetscFunctionBegin;
    359:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    360:   if (isascii) {
    361:     PetscViewerFormat format;

    363:     PetscCall(MatDiagonalSetUpDiagonal(J));
    364:     PetscCall(PetscViewerGetFormat(viewer, &format));
    365:     if (format == PETSC_VIEWER_ASCII_FACTOR_INFO || format == PETSC_VIEWER_ASCII_INFO) {
    366:       PetscCall(PetscObjectPrintClassNamePrefixType((PetscObject)ctx->diag, viewer));
    367:       PetscFunctionReturn(PETSC_SUCCESS);
    368:     }
    369:     PetscCall(VecView(ctx->diag, viewer));
    370:   }
    371:   PetscFunctionReturn(PETSC_SUCCESS);
    372: }

    374: static PetscErrorCode MatGetDiagonal_Diagonal(Mat J, Vec x)
    375: {
    376:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;

    378:   PetscFunctionBegin;
    379:   PetscCall(MatDiagonalSetUpDiagonal(J));
    380:   PetscCall(VecCopy(ctx->diag, x));
    381:   PetscFunctionReturn(PETSC_SUCCESS);
    382: }

    384: static PetscErrorCode MatDiagonalSet_Diagonal(Mat J, Vec D, InsertMode is)
    385: {
    386:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;

    388:   PetscFunctionBegin;
    389:   switch (is) {
    390:   case ADD_VALUES:
    391:   case ADD_ALL_VALUES:
    392:   case ADD_BC_VALUES:
    393:     PetscCall(MatDiagonalSetUpDiagonal(J));
    394:     PetscCall(VecAXPY(ctx->diag, 1.0, D));
    395:     ctx->inv_diag_valid = PETSC_FALSE;
    396:     break;
    397:   case INSERT_VALUES:
    398:   case INSERT_BC_VALUES:
    399:   case INSERT_ALL_VALUES:
    400:     PetscCall(MatSetUp(J));
    401:     PetscCall(VecCopy(D, ctx->diag));
    402:     ctx->diag_valid     = PETSC_TRUE;
    403:     ctx->inv_diag_valid = PETSC_FALSE;
    404:     break;
    405:   case MAX_VALUES:
    406:     PetscCall(MatDiagonalSetUpDiagonal(J));
    407:     PetscCall(VecPointwiseMax(ctx->diag, D, ctx->diag));
    408:     ctx->inv_diag_valid = PETSC_FALSE;
    409:     break;
    410:   case MIN_VALUES:
    411:     PetscCall(MatDiagonalSetUpDiagonal(J));
    412:     PetscCall(VecPointwiseMin(ctx->diag, D, ctx->diag));
    413:     ctx->inv_diag_valid = PETSC_FALSE;
    414:     break;
    415:   case NOT_SET_VALUES:
    416:     break;
    417:   }
    418:   PetscFunctionReturn(PETSC_SUCCESS);
    419: }

    421: static PetscErrorCode MatShift_Diagonal(Mat Y, PetscScalar a)
    422: {
    423:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    425:   PetscFunctionBegin;
    426:   PetscCall(MatDiagonalSetUpDiagonal(Y));
    427:   PetscCall(VecShift(ctx->diag, a));
    428:   ctx->inv_diag_valid = PETSC_FALSE;
    429:   PetscFunctionReturn(PETSC_SUCCESS);
    430: }

    432: static PetscErrorCode MatScale_Diagonal(Mat Y, PetscScalar a)
    433: {
    434:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    436:   PetscFunctionBegin;
    437:   PetscCall(MatSetUp(Y));
    438:   PetscCall(MatDiagonalSetUpDiagonal(Y));
    439:   PetscCall(VecScale(ctx->diag, a));
    440:   ctx->inv_diag_valid = PETSC_FALSE;
    441:   PetscFunctionReturn(PETSC_SUCCESS);
    442: }

    444: static PetscErrorCode MatDiagonalScale_Diagonal(Mat Y, Vec l, Vec r)
    445: {
    446:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    448:   PetscFunctionBegin;
    449:   PetscCall(MatDiagonalSetUpDiagonal(Y));
    450:   if (l) {
    451:     PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, l));
    452:     ctx->inv_diag_valid = PETSC_FALSE;
    453:   }
    454:   if (r) {
    455:     PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, r));
    456:     ctx->inv_diag_valid = PETSC_FALSE;
    457:   }
    458:   PetscFunctionReturn(PETSC_SUCCESS);
    459: }

    461: static PetscErrorCode MatConjugate_Diagonal(Mat Y)
    462: {
    463:   PetscFunctionBegin;
    464:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    466:   if (ctx->diag_valid) {
    467:     PetscCall(VecConjugate(ctx->diag));
    468:     PetscCall(PetscObjectStateGet((PetscObject)ctx->diag, &ctx->diag_state));
    469:   }
    470:   if (ctx->inv_diag_valid) {
    471:     PetscCall(VecConjugate(ctx->inv_diag));
    472:     PetscCall(PetscObjectStateGet((PetscObject)ctx->inv_diag, &ctx->inv_diag_state));
    473:   }
    474:   PetscFunctionReturn(PETSC_SUCCESS);
    475: }

    477: static PetscErrorCode MatTranspose_Diagonal(Mat A, MatReuse reuse, Mat *matout)
    478: {
    479:   PetscFunctionBegin;
    480:   if (reuse == MAT_INPLACE_MATRIX) {
    481:     PetscLayout tmplayout = A->rmap;

    483:     A->rmap = A->cmap;
    484:     A->cmap = tmplayout;
    485:   } else {
    486:     Vec diag, newdiag;
    487:     if (reuse == MAT_INITIAL_MATRIX) {
    488:       PetscCall(MatDiagonalGetDiagonal(A, &diag));
    489:       PetscCall(VecDuplicate(diag, &newdiag));
    490:       PetscCall(VecCopy(diag, newdiag));
    491:       PetscCall(MatDiagonalRestoreDiagonal(A, &diag));
    492:       PetscCall(MatCreateDiagonal(newdiag, matout));
    493:       PetscCall(VecDestroy(&newdiag));
    494:     } else {
    495:       PetscCall(MatDiagonalGetDiagonal(A, &diag));
    496:       PetscCall(MatDiagonalGetDiagonal(*matout, &newdiag));
    497:       PetscCall(VecCopy(diag, newdiag));
    498:       PetscCall(MatDiagonalRestoreDiagonal(*matout, &newdiag));
    499:       PetscCall(MatDiagonalRestoreDiagonal(A, &diag));
    500:     }
    501:   }
    502:   PetscFunctionReturn(PETSC_SUCCESS);
    503: }

    505: static PetscErrorCode MatSetUp_Diagonal(Mat A)
    506: {
    507:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    509:   PetscFunctionBegin;
    510:   if (!ctx->diag) {
    511:     PetscCall(PetscLayoutSetUp(A->cmap));
    512:     PetscCall(PetscLayoutSetUp(A->rmap));
    513:     PetscCall(MatCreateVecs(A, &ctx->diag, NULL));
    514:     PetscCall(VecDuplicate(ctx->diag, &ctx->inv_diag));
    515:     PetscCall(VecZeroEntries(ctx->diag));
    516:     ctx->diag_valid = PETSC_TRUE;
    517:   }
    518:   A->assembled = PETSC_TRUE;
    519:   PetscFunctionReturn(PETSC_SUCCESS);
    520: }

    522: static PetscErrorCode MatZeroEntries_Diagonal(Mat Y)
    523: {
    524:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    526:   PetscFunctionBegin;
    527:   PetscCall(MatSetUp(Y));
    528:   PetscCall(VecZeroEntries(ctx->diag));
    529:   ctx->diag_valid     = PETSC_TRUE;
    530:   ctx->inv_diag_valid = PETSC_FALSE;
    531:   PetscFunctionReturn(PETSC_SUCCESS);
    532: }

    534: static PetscErrorCode MatSolve_Diagonal(Mat matin, Vec b, Vec x)
    535: {
    536:   Mat_Diagonal *ctx = (Mat_Diagonal *)matin->data;

    538:   PetscFunctionBegin;
    539:   PetscCall(MatDiagonalSetUpInverseDiagonal(matin));
    540:   PetscCall(VecPointwiseMult(x, b, ctx->inv_diag));
    541:   PetscFunctionReturn(PETSC_SUCCESS);
    542: }

    544: static PetscErrorCode MatGetInfo_Diagonal(Mat A, MatInfoType flag, MatInfo *info)
    545: {
    546:   PetscFunctionBegin;
    547:   info->block_size        = 1.0;
    548:   info->nz_allocated      = A->cmap->N;
    549:   info->nz_used           = A->cmap->N;
    550:   info->nz_unneeded       = 0.0;
    551:   info->assemblies        = A->num_ass;
    552:   info->mallocs           = 0.0;
    553:   info->memory            = 0;
    554:   info->fill_ratio_given  = 0;
    555:   info->fill_ratio_needed = 0;
    556:   info->factor_mallocs    = 0;
    557:   PetscFunctionReturn(PETSC_SUCCESS);
    558: }

    560: /*@
    561:   MatCreateDiagonal - Creates a matrix defined by a given vector along its diagonal.

    563:   Collective

    565:   Input Parameter:
    566: . diag - vector for the diagonal

    568:   Output Parameter:
    569: . J - the diagonal matrix

    571:   Level: advanced

    573:   Notes:
    574:   Only supports square matrices with the same number of local rows and columns.

    576:   The input vector `diag` will be referenced internally: any changes to `diag`
    577:   will affect the matrix `J`.

    579: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MATCONSTANTDIAGONAL`, `MatScale()`, `MatShift()`, `MatMult()`, `MatGetDiagonal()`, `MatSolve()`
    580:           `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`
    581: @*/
    582: PetscErrorCode MatCreateDiagonal(Vec diag, Mat *J)
    583: {
    584:   PetscFunctionBegin;
    586:   PetscCall(MatCreate(PetscObjectComm((PetscObject)diag), J));
    587:   PetscInt m, M;
    588:   PetscCall(VecGetLocalSize(diag, &m));
    589:   PetscCall(VecGetSize(diag, &M));
    590:   PetscCall(MatSetSizes(*J, m, m, M, M));
    591:   PetscCall(MatSetType(*J, MATDIAGONAL));

    593:   PetscLayout map;
    594:   PetscCall(VecGetLayout(diag, &map));
    595:   PetscCall(MatSetLayouts(*J, map, map));
    596:   Mat_Diagonal *ctx = (Mat_Diagonal *)(*J)->data;
    597:   PetscCall(PetscObjectReference((PetscObject)diag));
    598:   PetscCall(VecDestroy(&ctx->diag));
    599:   PetscCall(VecDestroy(&ctx->inv_diag));
    600:   ctx->diag           = diag;
    601:   ctx->diag_valid     = PETSC_TRUE;
    602:   ctx->inv_diag_valid = PETSC_FALSE;
    603:   VecType type;
    604:   PetscCall(VecDuplicate(ctx->diag, &ctx->inv_diag));
    605:   PetscCall(VecGetType(diag, &type));
    606:   PetscCall(PetscFree((*J)->defaultvectype));
    607:   PetscCall(PetscStrallocpy(type, &(*J)->defaultvectype));
    608:   PetscCall(MatSetUp(*J));
    609:   ctx->col = NULL;
    610:   ctx->val = NULL;
    611:   PetscFunctionReturn(PETSC_SUCCESS);
    612: }

    614: static PetscErrorCode MatProductNumeric_Diagonal_Dense(Mat C)
    615: {
    616:   Mat                A, B;
    617:   Mat_Diagonal      *a;
    618:   PetscScalar       *c;
    619:   const PetscScalar *b, *alpha;
    620:   PetscInt           ldb, ldc;

    622:   PetscFunctionBegin;
    623:   MatCheckProduct(C, 1);
    624:   A = C->product->A;
    625:   B = C->product->B;
    626:   a = (Mat_Diagonal *)A->data;
    627:   PetscCall(VecGetArrayRead(a->diag, &alpha));
    628:   PetscCall(MatDenseGetLDA(B, &ldb));
    629:   PetscCall(MatDenseGetLDA(C, &ldc));
    630:   PetscCall(MatDenseGetArrayRead(B, &b));
    631:   PetscCall(MatDenseGetArrayWrite(C, &c));
    632:   for (PetscInt j = 0; j < B->cmap->N; j++)
    633:     for (PetscInt i = 0; i < B->rmap->n; i++) c[i + j * ldc] = alpha[i] * b[i + j * ldb];
    634:   PetscCall(MatDenseRestoreArrayWrite(C, &c));
    635:   PetscCall(MatDenseRestoreArrayRead(B, &b));
    636:   PetscCall(VecRestoreArrayRead(a->diag, &alpha));
    637:   PetscCall(PetscLogFlops(B->cmap->N * B->rmap->n));
    638:   PetscFunctionReturn(PETSC_SUCCESS);
    639: }

    641: static PetscErrorCode MatProductSymbolic_Diagonal_Dense(Mat C)
    642: {
    643:   Mat      A, B;
    644:   PetscInt n, N, m, M;

    646:   PetscFunctionBegin;
    647:   MatCheckProduct(C, 1);
    648:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
    649:   A = C->product->A;
    650:   B = C->product->B;
    651:   PetscCall(MatDiagonalSetUpDiagonal(A));
    652:   PetscCall(MatGetLocalSize(C, &m, &n));
    653:   PetscCall(MatGetSize(C, &M, &N));
    654:   if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) {
    655:     PetscCall(MatGetLocalSize(B, NULL, &n));
    656:     PetscCall(MatGetSize(B, NULL, &N));
    657:     PetscCall(MatGetLocalSize(A, &m, NULL));
    658:     PetscCall(MatGetSize(A, &M, NULL));
    659:     PetscCall(MatSetSizes(C, m, n, M, N));
    660:   }
    661:   PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
    662:   PetscCall(MatSetUp(C));
    663:   C->ops->productnumeric = MatProductNumeric_Diagonal_Dense;
    664:   PetscFunctionReturn(PETSC_SUCCESS);
    665: }

    667: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense_AB(Mat C)
    668: {
    669:   PetscFunctionBegin;
    670:   C->ops->productsymbolic = MatProductSymbolic_Diagonal_Dense;
    671:   PetscFunctionReturn(PETSC_SUCCESS);
    672: }

    674: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense(Mat C)
    675: {
    676:   Mat_Product *product = C->product;

    678:   PetscFunctionBegin;
    679:   if (product->type == MATPRODUCT_AB || product->type == MATPRODUCT_AtB) PetscCall(MatProductSetFromOptions_Diagonal_Dense_AB(C));
    680:   PetscFunctionReturn(PETSC_SUCCESS);
    681: }

    683: /*MC
    684:    MATDIAGONAL - MATDIAGONAL = "diagonal" - A diagonal matrix type with the diagonal implemented as a `Vec`.  Useful for
    685:    cases where `VecPointwiseMult()` or `VecPointwiseDivide()` should be thought of as the actions of a linear operator.

    687:   Level: advanced

    689: .seealso: [](ch_matrices), `Mat`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`
    690: M*/
    691: PETSC_INTERN PetscErrorCode MatCreate_Diagonal(Mat A)
    692: {
    693:   Mat_Diagonal *ctx;

    695:   PetscFunctionBegin;
    696:   PetscCall(PetscNew(&ctx));
    697:   A->data = (void *)ctx;

    699:   A->structurally_symmetric      = PETSC_BOOL3_TRUE;
    700:   A->structural_symmetry_eternal = PETSC_TRUE;
    701:   A->symmetry_eternal            = PETSC_TRUE;
    702:   A->symmetric                   = PETSC_BOOL3_TRUE;
    703:   if (!PetscDefined(USE_COMPLEX)) A->hermitian = PETSC_BOOL3_TRUE;

    705:   A->ops->getrow           = MatGetRow_Diagonal;
    706:   A->ops->mult             = MatMult_Diagonal;
    707:   A->ops->multadd          = MatMultAdd_Diagonal;
    708:   A->ops->multtranspose    = MatMult_Diagonal;
    709:   A->ops->multtransposeadd = MatMultAdd_Diagonal;
    710:   A->ops->norm             = MatNorm_Diagonal;
    711:   A->ops->duplicate        = MatDuplicate_Diagonal;
    712:   A->ops->solve            = MatSolve_Diagonal;
    713:   A->ops->solvetranspose   = MatSolve_Diagonal;
    714:   A->ops->shift            = MatShift_Diagonal;
    715:   A->ops->scale            = MatScale_Diagonal;
    716:   A->ops->diagonalscale    = MatDiagonalScale_Diagonal;
    717:   A->ops->getdiagonal      = MatGetDiagonal_Diagonal;
    718:   A->ops->diagonalset      = MatDiagonalSet_Diagonal;
    719:   A->ops->view             = MatView_Diagonal;
    720:   A->ops->zeroentries      = MatZeroEntries_Diagonal;
    721:   A->ops->destroy          = MatDestroy_Diagonal;
    722:   A->ops->getinfo          = MatGetInfo_Diagonal;
    723:   A->ops->axpy             = MatAXPY_Diagonal;
    724:   A->ops->setup            = MatSetUp_Diagonal;
    725:   A->ops->permute          = MatPermute_Diagonal;
    726:   A->ops->setrandom        = MatSetRandom_Diagonal;
    727:   A->ops->conjugate        = MatConjugate_Diagonal;
    728:   A->ops->transpose        = MatTranspose_Diagonal;

    730:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetDiagonal_C", MatDiagonalGetDiagonal_Diagonal));
    731:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreDiagonal_C", MatDiagonalRestoreDiagonal_Diagonal));
    732:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetInverseDiagonal_C", MatDiagonalGetInverseDiagonal_Diagonal));
    733:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreInverseDiagonal_C", MatDiagonalRestoreInverseDiagonal_Diagonal));
    734:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_seqdense_C", MatProductSetFromOptions_Diagonal_Dense));
    735:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_mpidense_C", MatProductSetFromOptions_Diagonal_Dense));
    736:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATDIAGONAL));
    737:   PetscFunctionReturn(PETSC_SUCCESS);
    738: }