Date: 2026-07-30T21:14:15+00:00

Actual source code: diagonal.c

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

      4: static PetscErrorCode MatDiagonalSetUpDiagonal(Mat A)
      5: {
      6:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

      8:   PetscFunctionBegin;
      9:   if (!ctx->diag_valid) {
     10:     PetscAssert(ctx->inv_diag_valid, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Neither diagonal nor inverse diagonal is in a valid state");
     11:     PetscCall(VecCopy(ctx->inv_diag, ctx->diag));
     12:     PetscCall(VecReciprocal(ctx->diag));
     13:     ctx->diag_valid = PETSC_TRUE;
     14:   }
     15:   PetscFunctionReturn(PETSC_SUCCESS);
     16: }

     18: static PetscErrorCode MatDiagonalSetUpInverseDiagonal(Mat A)
     19: {
     20:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

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

     32: static PetscErrorCode MatAXPY_Diagonal(Mat Y, PetscScalar a, Mat X, MatStructure str)
     33: {
     34:   Mat_Diagonal *yctx = (Mat_Diagonal *)Y->data;
     35:   Mat_Diagonal *xctx = (Mat_Diagonal *)X->data;

     37:   PetscFunctionBegin;
     38:   PetscCall(MatDiagonalSetUpDiagonal(Y));
     39:   PetscCall(MatDiagonalSetUpDiagonal(X));
     40:   PetscCall(VecAXPY(yctx->diag, a, xctx->diag));
     41:   yctx->inv_diag_valid = PETSC_FALSE;
     42:   PetscFunctionReturn(PETSC_SUCCESS);
     43: }

     45: static PetscErrorCode MatGetRow_Diagonal(Mat A, PetscInt row, PetscInt *ncols, PetscInt **cols, PetscScalar **vals)
     46: {
     47:   Mat_Diagonal *mat    = (Mat_Diagonal *)A->data;
     48:   PetscInt      rstart = A->rmap->rstart, rend = A->rmap->rend;

     50:   PetscFunctionBegin;
     51:   PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Only local rows");
     52:   if (ncols) *ncols = 1;
     53:   if (cols) {
     54:     if (!mat->col) PetscCall(PetscMalloc1(1, &mat->col));
     55:     *mat->col = row;
     56:     *cols     = mat->col;
     57:   }
     58:   if (vals) {
     59:     const PetscScalar *v;

     61:     if (!mat->val) PetscCall(PetscMalloc1(1, &mat->val));
     62:     PetscCall(VecGetArrayRead(mat->diag, &v));
     63:     *mat->val = v[row - rstart];
     64:     *vals     = mat->val;
     65:     PetscCall(VecRestoreArrayRead(mat->diag, &v));
     66:   }
     67:   PetscFunctionReturn(PETSC_SUCCESS);
     68: }

     70: static PetscErrorCode MatMult_Diagonal(Mat A, Vec x, Vec y)
     71: {
     72:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

     74:   PetscFunctionBegin;
     75:   PetscCall(MatDiagonalSetUpDiagonal(A));
     76:   PetscCall(VecPointwiseMult(y, ctx->diag, x));
     77:   PetscFunctionReturn(PETSC_SUCCESS);
     78: }

     80: static PetscErrorCode MatMultAdd_Diagonal(Mat mat, Vec v1, Vec v2, Vec v3)
     81: {
     82:   Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data;

     84:   PetscFunctionBegin;
     85:   PetscCall(MatDiagonalSetUpDiagonal(mat));
     86:   if (v2 != v3) {
     87:     PetscCall(VecPointwiseMult(v3, ctx->diag, v1));
     88:     PetscCall(VecAXPY(v3, 1.0, v2));
     89:   } else {
     90:     Vec w;
     91:     PetscCall(VecDuplicate(v3, &w));
     92:     PetscCall(VecPointwiseMult(w, ctx->diag, v1));
     93:     PetscCall(VecAXPY(v3, 1.0, w));
     94:     PetscCall(VecDestroy(&w));
     95:   }
     96:   PetscFunctionReturn(PETSC_SUCCESS);
     97: }

     99: static PetscErrorCode MatNorm_Diagonal(Mat A, NormType type, PetscReal *nrm)
    100: {
    101:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    103:   PetscFunctionBegin;
    104:   PetscCall(MatDiagonalSetUpDiagonal(A));
    105:   type = (type == NORM_FROBENIUS) ? NORM_2 : NORM_INFINITY;
    106:   PetscCall(VecNorm(ctx->diag, type, nrm));
    107:   PetscFunctionReturn(PETSC_SUCCESS);
    108: }

    110: static PetscErrorCode MatDuplicate_Diagonal(Mat A, MatDuplicateOption op, Mat *B)
    111: {
    112:   Mat_Diagonal *actx = (Mat_Diagonal *)A->data;

    114:   PetscFunctionBegin;
    115:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
    116:   PetscCall(MatSetSizes(*B, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
    117:   PetscCall(MatSetBlockSizesFromMats(*B, A, A));
    118:   PetscCall(MatSetType(*B, MATDIAGONAL));
    119:   PetscCall(PetscLayoutReference(A->rmap, &(*B)->rmap));
    120:   PetscCall(PetscLayoutReference(A->cmap, &(*B)->cmap));
    121:   if (op == MAT_COPY_VALUES) {
    122:     Mat_Diagonal *bctx = (Mat_Diagonal *)(*B)->data;

    124:     PetscCall(MatSetUp(A));
    125:     PetscCall(MatSetUp(*B));
    126:     PetscCall(VecCopy(actx->diag, bctx->diag));
    127:     PetscCall(VecCopy(actx->inv_diag, bctx->inv_diag));
    128:     bctx->diag_valid     = actx->diag_valid;
    129:     bctx->inv_diag_valid = actx->inv_diag_valid;
    130:   }
    131:   PetscFunctionReturn(PETSC_SUCCESS);
    132: }

    134: /*@
    135:   MatDiagonalGetDiagonal - Get the diagonal of a `MATDIAGONAL`

    137:   Input Parameter:
    138: . A - the `MATDIAGONAL`

    140:   Output Parameter:
    141: . diag - the `Vec` that defines the diagonal

    143:   Level: developer

    145:   Note:
    146:   The user must call
    147:   `MatDiagonalRestoreDiagonal()` before using the matrix again.

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

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

    154: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`, `MatGetDiagonal()`
    155: @*/
    156: PetscErrorCode MatDiagonalGetDiagonal(Mat A, Vec *diag)
    157: {
    158:   PetscFunctionBegin;
    160:   PetscAssertPointer(diag, 2);
    161:   *diag = NULL;
    162:   PetscUseMethod(A, "MatDiagonalGetDiagonal_C", (Mat, Vec *), (A, diag));
    163:   PetscFunctionReturn(PETSC_SUCCESS);
    164: }

    166: static PetscErrorCode MatDiagonalGetDiagonal_Diagonal(Mat A, Vec *diag)
    167: {
    168:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    170:   PetscFunctionBegin;
    171:   PetscCall(MatSetUp(A));
    172:   PetscCall(MatDiagonalSetUpDiagonal(A));
    173:   *diag = ctx->diag;
    174:   PetscCall(PetscObjectStateGet((PetscObject)*diag, &ctx->diag_state));
    175:   PetscFunctionReturn(PETSC_SUCCESS);
    176: }

    178: /*@
    179:   MatDiagonalRestoreDiagonal - Restore the diagonal of a `MATDIAGONAL`

    181:   Input Parameters:
    182: + A    - the `MATDIAGONAL`
    183: - diag - the `Vec` obtained from `MatDiagonalGetDiagonal()`

    185:   Level: developer

    187: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetDiagonal()`
    188: @*/
    189: PetscErrorCode MatDiagonalRestoreDiagonal(Mat A, Vec *diag)
    190: {
    191:   PetscFunctionBegin;
    193:   PetscAssertPointer(diag, 2);
    194:   PetscUseMethod(A, "MatDiagonalRestoreDiagonal_C", (Mat, Vec *), (A, diag));
    195:   PetscFunctionReturn(PETSC_SUCCESS);
    196: }

    198: static PetscErrorCode MatDiagonalRestoreDiagonal_Diagonal(Mat A, Vec *diag)
    199: {
    200:   Mat_Diagonal    *ctx = (Mat_Diagonal *)A->data;
    201:   PetscObjectState diag_state;

    203:   PetscFunctionBegin;
    204:   PetscCheck(ctx->diag == *diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector");
    205:   ctx->diag_valid = PETSC_TRUE;
    206:   PetscCall(PetscObjectStateGet((PetscObject)*diag, &diag_state));
    207:   if (ctx->diag_state != diag_state) {
    208:     PetscCall(PetscObjectStateIncrease((PetscObject)A));
    209:     ctx->inv_diag_valid = PETSC_FALSE;
    210:   }
    211:   *diag = NULL;
    212:   PetscFunctionReturn(PETSC_SUCCESS);
    213: }

    215: /*@
    216:   MatDiagonalGetInverseDiagonal - Get the inverse diagonal of a `MATDIAGONAL`

    218:   Input Parameter:
    219: . A - the `MATDIAGONAL`

    221:   Output Parameter:
    222: . inv_diag - the `Vec` that defines the inverse diagonal

    224:   Level: developer

    226:   Note:
    227:   The user must call
    228:   `MatDiagonalRestoreInverseDiagonal()` before using the matrix again.

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

    234: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MATLMVMBROYDEN`, `MatSolve()`
    235: @*/
    236: PetscErrorCode MatDiagonalGetInverseDiagonal(Mat A, Vec *inv_diag)
    237: {
    238:   PetscFunctionBegin;
    240:   PetscAssertPointer(inv_diag, 2);
    241:   *inv_diag = NULL;
    242:   PetscUseMethod(A, "MatDiagonalGetInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
    243:   PetscFunctionReturn(PETSC_SUCCESS);
    244: }

    246: static PetscErrorCode MatDiagonalGetInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
    247: {
    248:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    250:   PetscFunctionBegin;
    251:   PetscCall(MatSetUp(A));
    252:   PetscCall(MatDiagonalSetUpInverseDiagonal(A));
    253:   *inv_diag = ctx->inv_diag;
    254:   PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &ctx->inv_diag_state));
    255:   PetscFunctionReturn(PETSC_SUCCESS);
    256: }

    258: /*@
    259:   MatDiagonalRestoreInverseDiagonal - Restore the inverse diagonal of a `MATDIAGONAL`

    261:   Input Parameters:
    262: + A        - the `MATDIAGONAL`
    263: - inv_diag - the `Vec` obtained from `MatDiagonalGetInverseDiagonal()`

    265:   Level: developer

    267: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetInverseDiagonal()`
    268: @*/
    269: PetscErrorCode MatDiagonalRestoreInverseDiagonal(Mat A, Vec *inv_diag)
    270: {
    271:   PetscFunctionBegin;
    273:   PetscAssertPointer(inv_diag, 2);
    274:   PetscUseMethod(A, "MatDiagonalRestoreInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
    275:   PetscFunctionReturn(PETSC_SUCCESS);
    276: }

    278: static PetscErrorCode MatDiagonalRestoreInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
    279: {
    280:   Mat_Diagonal    *ctx = (Mat_Diagonal *)A->data;
    281:   PetscObjectState inv_diag_state;

    283:   PetscFunctionBegin;
    284:   PetscCheck(ctx->inv_diag == *inv_diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector");
    285:   ctx->inv_diag_valid = PETSC_TRUE;
    286:   PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &inv_diag_state));
    287:   if (ctx->inv_diag_state != inv_diag_state) {
    288:     PetscCall(PetscObjectStateIncrease((PetscObject)A));
    289:     ctx->diag_valid = PETSC_FALSE;
    290:   }
    291:   *inv_diag = NULL;
    292:   PetscFunctionReturn(PETSC_SUCCESS);
    293: }

    295: static PetscErrorCode MatPermute_Diagonal(Mat A, IS rowp, IS colp, Mat *B)
    296: {
    297:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
    298:   Vec           v;

    300:   PetscFunctionBegin;
    301:   PetscCheck(rowp == colp, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Row permutation and column permutation must be the same");
    302:   PetscCall(VecDuplicate(ctx->diag, &v));
    303:   PetscCall(VecCopy(ctx->diag, v));
    304:   PetscCall(VecPermute(v, rowp, PETSC_FALSE));
    305:   PetscCall(MatCreateDiagonal(v, B));
    306:   PetscCall(VecDestroy(&v));
    307:   PetscFunctionReturn(PETSC_SUCCESS);
    308: }

    310: static PetscErrorCode MatSetRandom_Diagonal(Mat A, PetscRandom rand)
    311: {
    312:   Vec d;

    314:   PetscFunctionBegin;
    315:   PetscCall(MatDiagonalGetDiagonal(A, &d));
    316:   PetscCall(VecSetRandom(d, rand));
    317:   PetscCall(MatDiagonalRestoreDiagonal(A, &d));
    318:   PetscFunctionReturn(PETSC_SUCCESS);
    319: }

    321: static PetscErrorCode MatDestroy_Diagonal(Mat mat)
    322: {
    323:   Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data;

    325:   PetscFunctionBegin;
    326:   PetscCall(VecDestroy(&ctx->diag));
    327:   PetscCall(VecDestroy(&ctx->inv_diag));
    328:   PetscCall(PetscFree(ctx->col));
    329:   PetscCall(PetscFree(ctx->val));
    330:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetDiagonal_C", NULL));
    331:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreDiagonal_C", NULL));
    332:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetInverseDiagonal_C", NULL));
    333:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreInverseDiagonal_C", NULL));
    334:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_seqdense_C", NULL));
    335:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_mpidense_C", NULL));
    336:   PetscCall(PetscFree(mat->data));
    337:   PetscFunctionReturn(PETSC_SUCCESS);
    338: }

    340: static PetscErrorCode MatView_Diagonal(Mat J, PetscViewer viewer)
    341: {
    342:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
    343:   PetscBool     isascii;

    345:   PetscFunctionBegin;
    346:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    347:   if (isascii) {
    348:     PetscViewerFormat format;

    350:     PetscCall(MatDiagonalSetUpDiagonal(J));
    351:     PetscCall(PetscViewerGetFormat(viewer, &format));
    352:     if (format == PETSC_VIEWER_ASCII_FACTOR_INFO || format == PETSC_VIEWER_ASCII_INFO) {
    353:       PetscCall(PetscObjectPrintClassNamePrefixType((PetscObject)ctx->diag, viewer));
    354:       PetscFunctionReturn(PETSC_SUCCESS);
    355:     }
    356:     PetscCall(VecView(ctx->diag, viewer));
    357:   }
    358:   PetscFunctionReturn(PETSC_SUCCESS);
    359: }

    361: static PetscErrorCode MatGetDiagonal_Diagonal(Mat J, Vec x)
    362: {
    363:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;

    365:   PetscFunctionBegin;
    366:   PetscCall(MatDiagonalSetUpDiagonal(J));
    367:   PetscCall(VecCopy(ctx->diag, x));
    368:   PetscFunctionReturn(PETSC_SUCCESS);
    369: }

    371: static PetscErrorCode MatADot_Diagonal_Local(Mat A, Vec x, Vec y, PetscScalar *val)
    372: {
    373:   Mat_Diagonal      *ctx = (Mat_Diagonal *)A->data;
    374:   PetscInt           n   = x->map->n;
    375:   const PetscScalar *ya, *xa, *wa;
    376:   PetscScalar        sum = 0;

    378:   PetscFunctionBegin;
    379:   PetscCheckSameTypeAndComm(x, 2, ctx->diag, 1);
    380:   PetscCheckSameTypeAndComm(y, 3, ctx->diag, 1);
    381:   PetscCall(VecGetArrayRead(x, &xa));
    382:   PetscCall(VecGetArrayRead(y, &ya));
    383:   PetscCall(VecGetArrayRead(ctx->diag, &wa));
    384:   for (PetscInt i = 0; i < n; i++) {
    385:     sum += PetscConj(ya[i]) * wa[i] * xa[i];
    386:   }
    387:   if (n > 0) PetscCall(PetscLogFlops(3.0 * n));
    388:   PetscCall(VecRestoreArrayRead(x, &xa));
    389:   PetscCall(VecRestoreArrayRead(y, &ya));
    390:   PetscCall(VecRestoreArrayRead(ctx->diag, &wa));
    391:   *val = sum;
    392:   PetscFunctionReturn(PETSC_SUCCESS);
    393: }

    395: static PetscErrorCode MatADot_Diagonal_MPI(Mat A, Vec x, Vec y, PetscScalar *val)
    396: {
    397:   PetscFunctionBegin;
    398:   PetscCall(MatDiagonalSetUpDiagonal(A));
    399:   PetscUseTypeMethod(A, adot_local, x, y, val);
    400:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, val, 1, MPIU_SCALAR, MPIU_SUM, PetscObjectComm((PetscObject)A)));
    401:   PetscFunctionReturn(PETSC_SUCCESS);
    402: }

    404: static PetscErrorCode MatANorm_Diagonal_Local(Mat A, Vec x, PetscReal *val)
    405: {
    406:   Mat_Diagonal      *ctx = (Mat_Diagonal *)A->data;
    407:   PetscInt           n   = x->map->n;
    408:   const PetscScalar *xa, *wa;
    409:   PetscScalar        sum = 0;

    411:   PetscFunctionBegin;
    412:   PetscCheckSameTypeAndComm(x, 2, ctx->diag, 1);
    413:   PetscCall(VecGetArrayRead(x, &xa));
    414:   PetscCall(VecGetArrayRead(ctx->diag, &wa));
    415:   for (PetscInt i = 0; i < n; i++) {
    416:     sum += PetscConj(xa[i]) * wa[i] * xa[i];
    417:   }
    418:   if (n > 0) PetscCall(PetscLogFlops(3.0 * n));
    419:   PetscCall(VecRestoreArrayRead(x, &xa));
    420:   PetscCall(VecRestoreArrayRead(ctx->diag, &wa));
    421:   PetscCheck(PetscAbsReal(PetscImaginaryPart(sum)) < 100 * PETSC_MACHINE_EPSILON, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Matrix argument is not Hermitian (diagonal has nonzero imaginary parts)");
    422:   *val = PetscRealPart(sum);
    423:   PetscFunctionReturn(PETSC_SUCCESS);
    424: }

    426: static PetscErrorCode MatANorm_Diagonal_MPI(Mat A, Vec x, PetscReal *val)
    427: {
    428:   PetscFunctionBegin;
    429:   PetscCall(MatDiagonalSetUpDiagonal(A));
    430:   PetscUseTypeMethod(A, anorm_local, x, val);
    431:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, val, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
    432:   *val = PetscSqrtReal(*val);
    433:   PetscFunctionReturn(PETSC_SUCCESS);
    434: }

    436: static PetscErrorCode MatANorm_Diagonal_Seq(Mat A, Vec x, PetscReal *val)
    437: {
    438:   PetscFunctionBegin;
    439:   PetscCall(MatDiagonalSetUpDiagonal(A));
    440:   PetscUseTypeMethod(A, anorm_local, x, val);
    441:   PetscCheck(*val >= 0.0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Matrix argument is not positive definite (diagonal has negative entries)");
    442:   *val = PetscSqrtReal(*val);
    443:   PetscFunctionReturn(PETSC_SUCCESS);
    444: }

    446: /*@
    447:   MatDiagonalSetDiagonal - Sets the diagonal for a `MATDIAGONAL`

    449:   Collective

    451:   Input Parameter:
    452: + J    - the `MATDIAGONAL` matrix
    453: - diag - the vector for the diagonal

    455:   Level: advanced

    457:   Notes:
    458:   The input vector `diag` will be referenced internally: any changes to `diag`
    459:   will affect the matrix `J`.

    461:   This routine can only be called once for the given `J` matrix.

    463: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MatCreateDiagonal()`, `MATDIAGONAL`, `MatScale()`, `MatShift()`, `MatMult()`, `MatGetDiagonal()`, `MatSolve()`,
    464:           `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`,
    465:           `MATCONSTANTDIAGONAL`
    466: @*/
    467: PetscErrorCode MatDiagonalSetDiagonal(Mat J, Vec diag)
    468: {
    469:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
    470:   PetscMPIInt   comm_size;
    471:   PetscLayout   map;
    472:   VecType       type;
    473:   PetscBool     iskokkos, ismpi, iscuda, iship;

    475:   PetscFunctionBegin;
    478:   PetscCheckSameComm(J, 1, diag, 2);
    479:   PetscCheck(!ctx->diag, PetscObjectComm((PetscObject)J), PETSC_ERR_ARG_WRONGSTATE, "MATDIAGONAL already has a diagonal");
    480:   ctx->diag = diag;
    481:   PetscCall(PetscObjectReference((PetscObject)ctx->diag));
    482:   PetscCall(VecDuplicate(ctx->diag, &ctx->inv_diag));
    483:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)J), &comm_size));
    484:   PetscCall(VecGetLayout(ctx->diag, &map));
    485:   PetscCall(MatSetLayouts(J, map, map));
    486:   PetscCall(PetscFree(J->defaultvectype));
    487:   PetscCall(VecGetType(ctx->diag, &type));
    488:   PetscCall(PetscStrallocpy(type, &J->defaultvectype));

    490:   PetscCall(PetscObjectTypeCompareAny((PetscObject)ctx->diag, &iscuda, VECCUDA, VECMPICUDA, VECSEQCUDA, ""));
    491:   PetscCall(PetscObjectTypeCompareAny((PetscObject)ctx->diag, &iship, VECHIP, VECMPIHIP, VECSEQHIP, ""));
    492:   PetscCall(PetscObjectTypeCompareAny((PetscObject)ctx->diag, &iskokkos, VECKOKKOS, VECMPIKOKKOS, VECSEQKOKKOS, ""));
    493:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)ctx->diag, VECMPI, &ismpi));
    494:   ismpi               = ismpi || comm_size > 1;
    495:   J->ops->adot_local  = MatADot_Diagonal_Local;
    496:   J->ops->adot        = MatADot_Diagonal_Local;
    497:   J->ops->anorm_local = MatANorm_Diagonal_Local;
    498:   if (iskokkos) {
    499:     PetscCheck(PetscDefined(HAVE_KOKKOS_KERNELS), PetscObjectComm((PetscObject)ctx->diag), PETSC_ERR_SUP, "Reconfigure using KOKKOS kernels support");
    500: #if PetscDefined(HAVE_KOKKOS_KERNELS)
    501:     J->ops->adot_local  = MatADot_Diagonal_SeqKokkos;
    502:     J->ops->adot        = MatADot_Diagonal_SeqKokkos;
    503:     J->ops->anorm_local = MatANormSq_Diagonal_SeqKokkos;
    504: #endif
    505:   }
    506:   if (iscuda) {
    507:     PetscCheck(PetscDefined(HAVE_CUDA), PetscObjectComm((PetscObject)ctx->diag), PETSC_ERR_SUP, "Reconfigure using CUDA support");
    508: #if PetscDefined(HAVE_CUDA)
    509:     J->ops->adot_local  = MatADot_Diagonal_SeqCUDA;
    510:     J->ops->adot        = MatADot_Diagonal_SeqCUDA;
    511:     J->ops->anorm_local = MatANormSq_Diagonal_SeqCUDA;
    512: #endif
    513:   }
    514:   if (iship) {
    515:     PetscCheck(PetscDefined(HAVE_HIP), PetscObjectComm((PetscObject)ctx->diag), PETSC_ERR_SUP, "Reconfigure using HIP support");
    516: #if PetscDefined(HAVE_HIP)
    517:     J->ops->adot_local  = MatADot_Diagonal_SeqHIP;
    518:     J->ops->adot        = MatADot_Diagonal_SeqHIP;
    519:     J->ops->anorm_local = MatANormSq_Diagonal_SeqHIP;
    520: #endif
    521:   }
    522:   if (ismpi) {
    523:     J->ops->adot  = MatADot_Diagonal_MPI;
    524:     J->ops->anorm = MatANorm_Diagonal_MPI;
    525:   } else {
    526:     J->ops->anorm = MatANorm_Diagonal_Seq;
    527:   }
    528:   ctx->col        = NULL;
    529:   ctx->val        = NULL;
    530:   ctx->diag_valid = PETSC_TRUE;
    531:   J->assembled    = PETSC_TRUE;
    532:   PetscFunctionReturn(PETSC_SUCCESS);
    533: }

    535: static PetscErrorCode MatDiagonalSet_Diagonal(Mat J, Vec D, InsertMode is)
    536: {
    537:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;

    539:   PetscFunctionBegin;
    540:   switch (is) {
    541:   case ADD_VALUES:
    542:   case ADD_ALL_VALUES:
    543:   case ADD_BC_VALUES:
    544:     PetscCall(MatDiagonalSetUpDiagonal(J));
    545:     PetscCall(VecAXPY(ctx->diag, 1.0, D));
    546:     ctx->inv_diag_valid = PETSC_FALSE;
    547:     break;
    548:   case INSERT_VALUES:
    549:   case INSERT_BC_VALUES:
    550:   case INSERT_ALL_VALUES:
    551:     PetscCall(VecCopy(D, ctx->diag));
    552:     ctx->diag_valid     = PETSC_TRUE;
    553:     ctx->inv_diag_valid = PETSC_FALSE;
    554:     break;
    555:   case MAX_VALUES:
    556:     PetscCall(MatDiagonalSetUpDiagonal(J));
    557:     PetscCall(VecPointwiseMax(ctx->diag, D, ctx->diag));
    558:     ctx->inv_diag_valid = PETSC_FALSE;
    559:     break;
    560:   case MIN_VALUES:
    561:     PetscCall(MatDiagonalSetUpDiagonal(J));
    562:     PetscCall(VecPointwiseMin(ctx->diag, D, ctx->diag));
    563:     ctx->inv_diag_valid = PETSC_FALSE;
    564:     break;
    565:   case NOT_SET_VALUES:
    566:     break;
    567:   }
    568:   PetscFunctionReturn(PETSC_SUCCESS);
    569: }

    571: static PetscErrorCode MatShift_Diagonal(Mat Y, PetscScalar a)
    572: {
    573:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    575:   PetscFunctionBegin;
    576:   PetscCall(MatDiagonalSetUpDiagonal(Y));
    577:   PetscCall(VecShift(ctx->diag, a));
    578:   ctx->inv_diag_valid = PETSC_FALSE;
    579:   PetscFunctionReturn(PETSC_SUCCESS);
    580: }

    582: static PetscErrorCode MatScale_Diagonal(Mat Y, PetscScalar a)
    583: {
    584:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    586:   PetscFunctionBegin;
    587:   PetscCall(MatDiagonalSetUpDiagonal(Y));
    588:   PetscCall(VecScale(ctx->diag, a));
    589:   ctx->inv_diag_valid = PETSC_FALSE;
    590:   PetscFunctionReturn(PETSC_SUCCESS);
    591: }

    593: static PetscErrorCode MatDiagonalScale_Diagonal(Mat Y, Vec l, Vec r)
    594: {
    595:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    597:   PetscFunctionBegin;
    598:   PetscCall(MatDiagonalSetUpDiagonal(Y));
    599:   if (l) {
    600:     PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, l));
    601:     ctx->inv_diag_valid = PETSC_FALSE;
    602:   }
    603:   if (r) {
    604:     PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, r));
    605:     ctx->inv_diag_valid = PETSC_FALSE;
    606:   }
    607:   PetscFunctionReturn(PETSC_SUCCESS);
    608: }

    610: static PetscErrorCode MatConjugate_Diagonal(Mat Y)
    611: {
    612:   PetscFunctionBegin;
    613:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    615:   if (ctx->diag_valid) {
    616:     PetscCall(VecConjugate(ctx->diag));
    617:     PetscCall(PetscObjectStateGet((PetscObject)ctx->diag, &ctx->diag_state));
    618:   }
    619:   if (ctx->inv_diag_valid) {
    620:     PetscCall(VecConjugate(ctx->inv_diag));
    621:     PetscCall(PetscObjectStateGet((PetscObject)ctx->inv_diag, &ctx->inv_diag_state));
    622:   }
    623:   PetscFunctionReturn(PETSC_SUCCESS);
    624: }

    626: static PetscErrorCode MatTranspose_Diagonal(Mat A, MatReuse reuse, Mat *matout)
    627: {
    628:   PetscFunctionBegin;
    629:   if (reuse == MAT_INPLACE_MATRIX) {
    630:     PetscLayout tmplayout = A->rmap;

    632:     A->rmap = A->cmap;
    633:     A->cmap = tmplayout;
    634:   } else {
    635:     Vec diag, newdiag;
    636:     if (reuse == MAT_INITIAL_MATRIX) {
    637:       PetscCall(MatDiagonalGetDiagonal(A, &diag));
    638:       PetscCall(VecDuplicate(diag, &newdiag));
    639:       PetscCall(VecCopy(diag, newdiag));
    640:       PetscCall(MatDiagonalRestoreDiagonal(A, &diag));
    641:       PetscCall(MatCreateDiagonal(newdiag, matout));
    642:       PetscCall(VecDestroy(&newdiag));
    643:     } else {
    644:       PetscCall(MatDiagonalGetDiagonal(A, &diag));
    645:       PetscCall(MatDiagonalGetDiagonal(*matout, &newdiag));
    646:       PetscCall(VecCopy(diag, newdiag));
    647:       PetscCall(MatDiagonalRestoreDiagonal(*matout, &newdiag));
    648:       PetscCall(MatDiagonalRestoreDiagonal(A, &diag));
    649:     }
    650:   }
    651:   PetscFunctionReturn(PETSC_SUCCESS);
    652: }

    654: static PetscErrorCode MatSetUp_Diagonal(Mat A)
    655: {
    656:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

    658:   PetscFunctionBegin;
    659:   if (!ctx->diag) {
    660:     Vec diag;

    662:     PetscCall(PetscLayoutSetUp(A->cmap));
    663:     PetscCall(PetscLayoutSetUp(A->rmap));
    664:     PetscCall(MatCreateVecs(A, &diag, NULL));
    665:     PetscCall(MatDiagonalSetDiagonal(A, diag));
    666:     PetscCall(VecDestroy(&diag));
    667:   }
    668:   A->assembled = PETSC_TRUE;
    669:   PetscFunctionReturn(PETSC_SUCCESS);
    670: }

    672: static PetscErrorCode MatZeroEntries_Diagonal(Mat Y)
    673: {
    674:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

    676:   PetscFunctionBegin;
    677:   PetscCall(VecZeroEntries(ctx->diag));
    678:   ctx->diag_valid     = PETSC_TRUE;
    679:   ctx->inv_diag_valid = PETSC_FALSE;
    680:   PetscFunctionReturn(PETSC_SUCCESS);
    681: }

    683: static PetscErrorCode MatSolve_Diagonal(Mat matin, Vec b, Vec x)
    684: {
    685:   Mat_Diagonal *ctx = (Mat_Diagonal *)matin->data;

    687:   PetscFunctionBegin;
    688:   PetscCall(MatDiagonalSetUpInverseDiagonal(matin));
    689:   PetscCall(VecPointwiseMult(x, b, ctx->inv_diag));
    690:   PetscFunctionReturn(PETSC_SUCCESS);
    691: }

    693: static PetscErrorCode MatGetInfo_Diagonal(Mat A, MatInfoType flag, MatInfo *info)
    694: {
    695:   PetscFunctionBegin;
    696:   info->block_size        = 1.0;
    697:   info->nz_allocated      = A->cmap->N;
    698:   info->nz_used           = A->cmap->N;
    699:   info->nz_unneeded       = 0.0;
    700:   info->assemblies        = A->num_ass;
    701:   info->mallocs           = 0.0;
    702:   info->memory            = 0;
    703:   info->fill_ratio_given  = 0;
    704:   info->fill_ratio_needed = 0;
    705:   info->factor_mallocs    = 0;
    706:   PetscFunctionReturn(PETSC_SUCCESS);
    707: }

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

    712:   Collective

    714:   Input Parameter:
    715: . diag - vector for the diagonal

    717:   Output Parameter:
    718: . J - the diagonal matrix

    720:   Level: advanced

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

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

    728: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MATDIAGONAL`, `MatScale()`, `MatShift()`, `MatMult()`, `MatGetDiagonal()`, `MatSolve()`,
    729:           `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`,
    730:           `MATCONSTANTDIAGONAL`
    731: @*/
    732: PetscErrorCode MatCreateDiagonal(Vec diag, Mat *J)
    733: {
    734:   PetscFunctionBegin;
    736:   PetscCall(MatCreate(PetscObjectComm((PetscObject)diag), J));
    737:   PetscCall(MatSetType(*J, MATDIAGONAL));
    738:   PetscCall(MatDiagonalSetDiagonal(*J, diag));
    739:   PetscFunctionReturn(PETSC_SUCCESS);
    740: }

    742: static PetscErrorCode MatProductNumeric_Diagonal_Dense(Mat C)
    743: {
    744:   Mat                A, B;
    745:   Mat_Diagonal      *a;
    746:   PetscScalar       *c;
    747:   const PetscScalar *b, *alpha;
    748:   PetscInt           ldb, ldc;

    750:   PetscFunctionBegin;
    751:   MatCheckProduct(C, 1);
    752:   A = C->product->A;
    753:   B = C->product->B;
    754:   PetscCall(MatDiagonalSetUpDiagonal(A));
    755:   a = (Mat_Diagonal *)A->data;
    756:   PetscCall(VecGetArrayRead(a->diag, &alpha));
    757:   PetscCall(MatDenseGetLDA(B, &ldb));
    758:   PetscCall(MatDenseGetLDA(C, &ldc));
    759:   PetscCall(MatDenseGetArrayRead(B, &b));
    760:   PetscCall(MatDenseGetArrayWrite(C, &c));
    761:   for (PetscInt j = 0; j < B->cmap->N; j++)
    762:     for (PetscInt i = 0; i < B->rmap->n; i++) c[i + j * ldc] = alpha[i] * b[i + j * ldb];
    763:   PetscCall(MatDenseRestoreArrayWrite(C, &c));
    764:   PetscCall(MatDenseRestoreArrayRead(B, &b));
    765:   PetscCall(VecRestoreArrayRead(a->diag, &alpha));
    766:   PetscCall(PetscLogFlops(B->cmap->N * B->rmap->n));
    767:   PetscFunctionReturn(PETSC_SUCCESS);
    768: }

    770: static PetscErrorCode MatProductSymbolic_Diagonal_Dense(Mat C)
    771: {
    772:   Mat      A, B;
    773:   PetscInt n, N, m, M;

    775:   PetscFunctionBegin;
    776:   MatCheckProduct(C, 1);
    777:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
    778:   A = C->product->A;
    779:   B = C->product->B;
    780:   PetscCall(MatGetLocalSize(C, &m, &n));
    781:   PetscCall(MatGetSize(C, &M, &N));
    782:   if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) {
    783:     PetscCall(MatGetLocalSize(B, NULL, &n));
    784:     PetscCall(MatGetSize(B, NULL, &N));
    785:     PetscCall(MatGetLocalSize(A, &m, NULL));
    786:     PetscCall(MatGetSize(A, &M, NULL));
    787:     PetscCall(MatSetSizes(C, m, n, M, N));
    788:   }
    789:   PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
    790:   PetscCall(MatSetUp(C));
    791:   C->ops->productnumeric = MatProductNumeric_Diagonal_Dense;
    792:   PetscFunctionReturn(PETSC_SUCCESS);
    793: }

    795: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense_AB(Mat C)
    796: {
    797:   PetscFunctionBegin;
    798:   C->ops->productsymbolic = MatProductSymbolic_Diagonal_Dense;
    799:   PetscFunctionReturn(PETSC_SUCCESS);
    800: }

    802: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense(Mat C)
    803: {
    804:   Mat_Product *product = C->product;

    806:   PetscFunctionBegin;
    807:   if (product->type == MATPRODUCT_AB || product->type == MATPRODUCT_AtB) PetscCall(MatProductSetFromOptions_Diagonal_Dense_AB(C));
    808:   PetscFunctionReturn(PETSC_SUCCESS);
    809: }

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

    815:   Options Database Key:
    816: . -mat_vec_type type - set the `VecType` of the vector defining the diagonal

    818:   Level: advanced

    820:   Note:
    821:   The first use case is simply to call `MatCreateDiagonal()' to provide the vector containing the diagonal matrix.  The input vector  will be
    822:   referenced internally by the matrix: any changes to it will affect the matrix. Similar changes to the matrix will affect the vector.

    824:   For the second use case call `MatSetType()` with a type of `MATDIAGONAL` followed by a call to `MatDiagonalSetDiagonal()`. The input vector  will be
    825:   referenced internally by the matrix: any changes to it will affect the matrix.  Similar changes to the matrix will affect the vector.

    827:   For the third use case call `MatSetType()` with a type of `MATDIAGONAL` followed by calls to `MatSetSizes()` and `MatDiagonalSet()`
    828:   (in this case the diagonal vector will not be referenced internally by the matrix, its values will be copied) or some other
    829:   operation to provide the matrix entries. One can control the `VecType` of the diagonal by calling `MatSetVecType()` or using `-mat_vec_type type`.

    831: .seealso: [](ch_matrices), `Mat`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`,
    832:           `MatDiagonalGetInverseDiagonal()`, `MatDiagonalSet()`, `MatDiagonalSetDiagonal()`, `MatSetVecType()`
    833: M*/
    834: PETSC_INTERN PetscErrorCode MatCreate_Diagonal(Mat A)
    835: {
    836:   Mat_Diagonal *ctx;

    838:   PetscFunctionBegin;
    839:   PetscCall(PetscNew(&ctx));
    840:   A->data = (void *)ctx;

    842:   A->structurally_symmetric      = PETSC_BOOL3_TRUE;
    843:   A->structural_symmetry_eternal = PETSC_TRUE;
    844:   A->symmetry_eternal            = PETSC_TRUE;
    845:   A->symmetric                   = PETSC_BOOL3_TRUE;
    846:   if (!PetscDefined(USE_COMPLEX)) A->hermitian = PETSC_BOOL3_TRUE;

    848:   A->ops->getrow           = MatGetRow_Diagonal;
    849:   A->ops->mult             = MatMult_Diagonal;
    850:   A->ops->multadd          = MatMultAdd_Diagonal;
    851:   A->ops->multtranspose    = MatMult_Diagonal;
    852:   A->ops->multtransposeadd = MatMultAdd_Diagonal;
    853:   A->ops->norm             = MatNorm_Diagonal;
    854:   A->ops->duplicate        = MatDuplicate_Diagonal;
    855:   A->ops->solve            = MatSolve_Diagonal;
    856:   A->ops->solvetranspose   = MatSolve_Diagonal;
    857:   A->ops->shift            = MatShift_Diagonal;
    858:   A->ops->scale            = MatScale_Diagonal;
    859:   A->ops->diagonalscale    = MatDiagonalScale_Diagonal;
    860:   A->ops->getdiagonal      = MatGetDiagonal_Diagonal;
    861:   A->ops->diagonalset      = MatDiagonalSet_Diagonal;
    862:   A->ops->view             = MatView_Diagonal;
    863:   A->ops->zeroentries      = MatZeroEntries_Diagonal;
    864:   A->ops->destroy          = MatDestroy_Diagonal;
    865:   A->ops->getinfo          = MatGetInfo_Diagonal;
    866:   A->ops->axpy             = MatAXPY_Diagonal;
    867:   A->ops->setup            = MatSetUp_Diagonal;
    868:   A->ops->permute          = MatPermute_Diagonal;
    869:   A->ops->setrandom        = MatSetRandom_Diagonal;
    870:   A->ops->conjugate        = MatConjugate_Diagonal;
    871:   A->ops->transpose        = MatTranspose_Diagonal;

    873:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetDiagonal_C", MatDiagonalGetDiagonal_Diagonal));
    874:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreDiagonal_C", MatDiagonalRestoreDiagonal_Diagonal));
    875:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetInverseDiagonal_C", MatDiagonalGetInverseDiagonal_Diagonal));
    876:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreInverseDiagonal_C", MatDiagonalRestoreInverseDiagonal_Diagonal));
    877:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_seqdense_C", MatProductSetFromOptions_Diagonal_Dense));
    878:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_mpidense_C", MatProductSetFromOptions_Diagonal_Dense));
    879:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATDIAGONAL));
    880:   PetscFunctionReturn(PETSC_SUCCESS);
    881: }