Date: 2026-06-01T23:13:38+00:00

Actual source code: axpy.c

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

      3: static PetscErrorCode MatTransposeAXPY_Private(Mat Y, PetscScalar a, Mat X, MatStructure str, Mat T)
      4: {
      5:   Mat         A, F;
      6:   PetscScalar vshift, vscale;
      7:   PetscErrorCode (*f)(Mat, Mat *);

      9:   PetscFunctionBegin;
     10:   if (T == X) PetscCall(MatShellGetScalingShifts(T, &vshift, &vscale, (Vec *)MAT_SHELL_NOT_ALLOWED, (Vec *)MAT_SHELL_NOT_ALLOWED, (Vec *)MAT_SHELL_NOT_ALLOWED, (Mat *)MAT_SHELL_NOT_ALLOWED, (IS *)MAT_SHELL_NOT_ALLOWED, (IS *)MAT_SHELL_NOT_ALLOWED));
     11:   else {
     12:     vshift = 0.0;
     13:     vscale = 1.0;
     14:   }
     15:   PetscCall(PetscObjectQueryFunction((PetscObject)T, "MatTransposeGetMat_C", &f));
     16:   if (f) {
     17:     PetscCall(MatTransposeGetMat(T, &A));
     18:     if (T == X) {
     19:       PetscCall(PetscInfo(NULL, "Explicitly transposing X of type MATTRANSPOSEVIRTUAL to perform MatAXPY()\n"));
     20:       PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &F));
     21:       A = Y;
     22:     } else {
     23:       PetscCall(PetscInfo(NULL, "Transposing X because Y of type MATTRANSPOSEVIRTUAL to perform MatAXPY()\n"));
     24:       PetscCall(MatTranspose(X, MAT_INITIAL_MATRIX, &F));
     25:     }
     26:   } else {
     27:     PetscCall(MatHermitianTransposeGetMat(T, &A));
     28:     if (T == X) {
     29:       PetscCall(PetscInfo(NULL, "Explicitly Hermitian transposing X of type MATHERMITIANTRANSPOSEVIRTUAL to perform MatAXPY()\n"));
     30:       PetscCall(MatHermitianTranspose(A, MAT_INITIAL_MATRIX, &F));
     31:       A = Y;
     32:     } else {
     33:       PetscCall(PetscInfo(NULL, "Hermitian transposing X because Y of type MATHERMITIANTRANSPOSEVIRTUAL to perform MatAXPY()\n"));
     34:       PetscCall(MatHermitianTranspose(X, MAT_INITIAL_MATRIX, &F));
     35:     }
     36:   }
     37:   PetscCall(MatAXPY(A, a * vscale, F, str));
     38:   PetscCall(MatShift(A, a * vshift));
     39:   PetscCall(MatDestroy(&F));
     40:   PetscFunctionReturn(PETSC_SUCCESS);
     41: }

     43: static PetscErrorCode MatAXPY_BasicWithTypeCompare(Mat Y, PetscScalar a, Mat X, MatStructure str)
     44: {
     45:   PetscBool flg;

     47:   PetscFunctionBegin;
     48:   PetscCall(MatIsShell(Y, &flg));
     49:   if (flg) { /* MatShell has special support for AXPY */
     50:     PetscErrorCode (*f)(Mat, PetscScalar, Mat, MatStructure);

     52:     PetscCall(MatGetOperation(Y, MATOP_AXPY, (PetscErrorCodeFn **)&f));
     53:     if (f) {
     54:       PetscCall((*f)(Y, a, X, str));
     55:       PetscFunctionReturn(PETSC_SUCCESS);
     56:     }
     57:   } else {
     58:     /* no need to preallocate if Y is dense */
     59:     PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)Y, &flg, MATSEQDENSE, MATMPIDENSE, ""));
     60:     if (flg) {
     61:       PetscCall(PetscObjectTypeCompare((PetscObject)X, MATNEST, &flg));
     62:       if (flg) {
     63:         PetscCall(MatAXPY_Dense_Nest(Y, a, X));
     64:         PetscFunctionReturn(PETSC_SUCCESS);
     65:       }
     66:     }
     67:     PetscCall(PetscObjectTypeCompareAny((PetscObject)X, &flg, MATSCALAPACK, MATELEMENTAL, ""));
     68:     if (flg) { /* Avoid MatAXPY_Basic() due to missing MatGetRow() */
     69:       Mat C;

     71:       PetscCall(MatConvert(X, ((PetscObject)Y)->type_name, MAT_INITIAL_MATRIX, &C));
     72:       PetscCall(MatAXPY(Y, a, C, str));
     73:       PetscCall(MatDestroy(&C));
     74:       PetscFunctionReturn(PETSC_SUCCESS);
     75:     }
     76:     PetscCall(PetscObjectTypeCompare((PetscObject)X, MATCONSTANTDIAGONAL, &flg));
     77:     if (flg) {
     78:       PetscScalar b;

     80:       PetscCall(MatConstantDiagonalGetConstant(X, &b));
     81:       PetscCall(MatShift(Y, a * b));
     82:       PetscFunctionReturn(PETSC_SUCCESS);
     83:     }
     84:     PetscCall(PetscObjectTypeCompare((PetscObject)X, MATDIAGONAL, &flg));
     85:     if (flg) {
     86:       Vec d_X;

     88:       PetscCall(MatDiagonalGetDiagonal(X, &d_X));
     89:       if (a == 1.0) {
     90:         PetscCall(MatDiagonalSet(Y, d_X, ADD_VALUES));
     91:       } else {
     92:         Vec d_Y;

     94:         PetscCall(PetscObjectQuery((PetscObject)Y, "__MatAXPY_BasicWithTypeCompare_Diagonal", (PetscObject *)&d_Y));
     95:         if (!d_Y) {
     96:           Vec _d_Y;
     97:           PetscCall(VecDuplicate(d_X, &_d_Y));
     98:           PetscCall(PetscObjectCompose((PetscObject)Y, "__MatAXPY_BasicWithTypeCompare_Diagonal", (PetscObject)_d_Y));
     99:           d_Y = _d_Y;
    100:           PetscCall(VecDestroy(&_d_Y));
    101:         }
    102:         PetscCall(VecAXPBY(d_Y, a, 0.0, d_X));
    103:         PetscCall(MatDiagonalSet(Y, d_Y, ADD_VALUES));
    104:       }
    105:       PetscCall(MatDiagonalRestoreDiagonal(X, &d_X));
    106:       PetscFunctionReturn(PETSC_SUCCESS);
    107:     }
    108:   }
    109:   PetscCall(MatAXPY_Basic(Y, a, X, str));
    110:   PetscFunctionReturn(PETSC_SUCCESS);
    111: }

    113: /*@
    114:   MatAXPY - Computes Y = a*X + Y.

    116:   Logically Collective

    118:   Input Parameters:
    119: + a   - the scalar multiplier
    120: . X   - the first matrix
    121: . Y   - the second matrix
    122: - str - either `SAME_NONZERO_PATTERN`, `DIFFERENT_NONZERO_PATTERN`, `UNKNOWN_NONZERO_PATTERN`, or `SUBSET_NONZERO_PATTERN` (nonzeros of `X` is a subset of `Y`'s)

    124:   Level: intermediate

    126: .seealso: [](ch_matrices), `Mat`, `MatAYPX()`
    127:  @*/
    128: PetscErrorCode MatAXPY(Mat Y, PetscScalar a, Mat X, MatStructure str)
    129: {
    130:   PetscInt  M1, M2, N1, N2;
    131:   PetscInt  m1, m2, n1, n2;
    132:   PetscBool sametype, transpose;

    134:   PetscFunctionBegin;
    138:   PetscCheckSameComm(Y, 1, X, 3);
    139:   PetscCall(MatGetSize(X, &M1, &N1));
    140:   PetscCall(MatGetSize(Y, &M2, &N2));
    141:   PetscCall(MatGetLocalSize(X, &m1, &n1));
    142:   PetscCall(MatGetLocalSize(Y, &m2, &n2));
    143:   PetscCheck(M1 == M2 && N1 == N2, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_SIZ, "Non conforming matrix add: global sizes X %" PetscInt_FMT " x %" PetscInt_FMT ", Y %" PetscInt_FMT " x %" PetscInt_FMT, M1, N1, M2, N2);
    144:   PetscCheck(m1 == m2 && n1 == n2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Non conforming matrix add: local sizes X %" PetscInt_FMT " x %" PetscInt_FMT ", Y %" PetscInt_FMT " x %" PetscInt_FMT, m1, n1, m2, n2);
    145:   PetscCheck(Y->assembled, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix (Y)");
    146:   PetscCheck(X->assembled, PetscObjectComm((PetscObject)X), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix (X)");
    147:   if (a == (PetscScalar)0.0) PetscFunctionReturn(PETSC_SUCCESS);
    148:   if (Y == X) {
    149:     PetscCall(MatScale(Y, 1.0 + a));
    150:     PetscFunctionReturn(PETSC_SUCCESS);
    151:   }
    152:   PetscCall(PetscObjectObjectTypeCompare((PetscObject)X, (PetscObject)Y, &sametype));
    153:   PetscCall(PetscLogEventBegin(MAT_AXPY, Y, 0, 0, 0));
    154:   if (Y->ops->axpy && (sametype || X->ops->axpy == Y->ops->axpy)) {
    155:     PetscUseTypeMethod(Y, axpy, a, X, str);
    156:   } else {
    157:     PetscCall(PetscObjectTypeCompareAny((PetscObject)X, &transpose, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
    158:     if (transpose) {
    159:       PetscCall(MatTransposeAXPY_Private(Y, a, X, str, X));
    160:     } else {
    161:       PetscCall(PetscObjectTypeCompareAny((PetscObject)Y, &transpose, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
    162:       if (transpose) {
    163:         PetscCall(MatTransposeAXPY_Private(Y, a, X, str, Y));
    164:       } else {
    165:         PetscCall(MatAXPY_BasicWithTypeCompare(Y, a, X, str));
    166:       }
    167:     }
    168:   }
    169:   PetscCall(PetscLogEventEnd(MAT_AXPY, Y, 0, 0, 0));
    170:   PetscFunctionReturn(PETSC_SUCCESS);
    171: }

    173: PetscErrorCode MatAXPY_Basic_Preallocate(Mat Y, Mat X, Mat *B)
    174: {
    175:   PetscErrorCode (*preall)(Mat, Mat, Mat *) = NULL;

    177:   PetscFunctionBegin;
    178:   /* look for any available faster alternative to the general preallocator */
    179:   PetscCall(PetscObjectQueryFunction((PetscObject)Y, "MatAXPYGetPreallocation_C", &preall));
    180:   if (preall) {
    181:     PetscCall((*preall)(Y, X, B));
    182:   } else { /* Use MatPrellocator, assumes same row-col distribution */
    183:     Mat      preallocator;
    184:     PetscInt r, rstart, rend;
    185:     PetscInt m, n, M, N;

    187:     PetscCall(MatGetRowUpperTriangular(Y));
    188:     PetscCall(MatGetRowUpperTriangular(X));
    189:     PetscCall(MatGetSize(Y, &M, &N));
    190:     PetscCall(MatGetLocalSize(Y, &m, &n));
    191:     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), &preallocator));
    192:     PetscCall(MatSetType(preallocator, MATPREALLOCATOR));
    193:     PetscCall(MatSetLayouts(preallocator, Y->rmap, Y->cmap));
    194:     PetscCall(MatSetUp(preallocator));
    195:     PetscCall(MatGetOwnershipRange(preallocator, &rstart, &rend));
    196:     for (r = rstart; r < rend; ++r) {
    197:       PetscInt           ncols;
    198:       const PetscInt    *row;
    199:       const PetscScalar *vals;

    201:       PetscCall(MatGetRow(Y, r, &ncols, &row, &vals));
    202:       PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
    203:       PetscCall(MatRestoreRow(Y, r, &ncols, &row, &vals));
    204:       PetscCall(MatGetRow(X, r, &ncols, &row, &vals));
    205:       PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
    206:       PetscCall(MatRestoreRow(X, r, &ncols, &row, &vals));
    207:     }
    208:     PetscCall(MatSetOption(preallocator, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    209:     PetscCall(MatAssemblyBegin(preallocator, MAT_FINAL_ASSEMBLY));
    210:     PetscCall(MatAssemblyEnd(preallocator, MAT_FINAL_ASSEMBLY));
    211:     PetscCall(MatRestoreRowUpperTriangular(Y));
    212:     PetscCall(MatRestoreRowUpperTriangular(X));

    214:     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), B));
    215:     PetscCall(MatSetType(*B, ((PetscObject)Y)->type_name));
    216:     PetscCall(MatSetLayouts(*B, Y->rmap, Y->cmap));
    217:     PetscCall(MatPreallocatorPreallocate(preallocator, PETSC_FALSE, *B));
    218:     PetscCall(MatDestroy(&preallocator));
    219:   }
    220:   PetscFunctionReturn(PETSC_SUCCESS);
    221: }

    223: PetscErrorCode MatAXPY_Basic(Mat Y, PetscScalar a, Mat X, MatStructure str)
    224: {
    225:   PetscFunctionBegin;
    226:   if (str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN) {
    227:     PetscBool isdense;

    229:     /* no need to preallocate if Y is dense */
    230:     PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)Y, &isdense, MATSEQDENSE, MATMPIDENSE, ""));
    231:     if (isdense) str = SUBSET_NONZERO_PATTERN;
    232:   }
    233:   if (str != DIFFERENT_NONZERO_PATTERN && str != UNKNOWN_NONZERO_PATTERN) {
    234:     PetscInt           i, start, end, j, ncols, m, n;
    235:     const PetscInt    *row;
    236:     PetscScalar       *val;
    237:     const PetscScalar *vals;
    238:     PetscBool          option;

    240:     PetscCall(MatGetSize(X, &m, &n));
    241:     PetscCall(MatGetOwnershipRange(X, &start, &end));
    242:     PetscCall(MatGetRowUpperTriangular(X));
    243:     if (a == 1.0) {
    244:       for (i = start; i < end; i++) {
    245:         PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
    246:         PetscCall(MatSetValues(Y, 1, &i, ncols, row, vals, ADD_VALUES));
    247:         PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
    248:       }
    249:     } else {
    250:       PetscInt vs = 100;
    251:       /* realloc if needed, as this function may be used in parallel */
    252:       PetscCall(PetscMalloc1(vs, &val));
    253:       for (i = start; i < end; i++) {
    254:         PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
    255:         if (vs < ncols) {
    256:           vs = PetscMin(2 * ncols, n);
    257:           PetscCall(PetscRealloc(vs * sizeof(*val), &val));
    258:         }
    259:         for (j = 0; j < ncols; j++) val[j] = a * vals[j];
    260:         PetscCall(MatSetValues(Y, 1, &i, ncols, row, val, ADD_VALUES));
    261:         PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
    262:       }
    263:       PetscCall(PetscFree(val));
    264:     }
    265:     PetscCall(MatRestoreRowUpperTriangular(X));
    266:     PetscCall(MatGetOption(Y, MAT_NO_OFF_PROC_ENTRIES, &option));
    267:     PetscCall(MatSetOption(Y, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    268:     PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
    269:     PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
    270:     PetscCall(MatSetOption(Y, MAT_NO_OFF_PROC_ENTRIES, option));
    271:   } else {
    272:     Mat B;

    274:     PetscCall(MatAXPY_Basic_Preallocate(Y, X, &B));
    275:     PetscCall(MatAXPY_BasicWithPreallocation(B, Y, a, X, str));
    276:     PetscCall(MatHeaderMerge(Y, &B));
    277:   }
    278:   PetscFunctionReturn(PETSC_SUCCESS);
    279: }

    281: PetscErrorCode MatAXPY_BasicWithPreallocation(Mat B, Mat Y, PetscScalar a, Mat X, MatStructure str)
    282: {
    283:   PetscInt           i, start, end, j, ncols, m, n;
    284:   const PetscInt    *row;
    285:   PetscScalar       *val;
    286:   const PetscScalar *vals;
    287:   PetscBool          option;

    289:   PetscFunctionBegin;
    290:   PetscCall(MatGetSize(X, &m, &n));
    291:   PetscCall(MatGetOwnershipRange(X, &start, &end));
    292:   PetscCall(MatGetRowUpperTriangular(Y));
    293:   PetscCall(MatGetRowUpperTriangular(X));
    294:   if (a == 1.0) {
    295:     for (i = start; i < end; i++) {
    296:       PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
    297:       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
    298:       PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));

    300:       PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
    301:       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
    302:       PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
    303:     }
    304:   } else {
    305:     PetscInt vs = 100;
    306:     /* realloc if needed, as this function may be used in parallel */
    307:     PetscCall(PetscMalloc1(vs, &val));
    308:     for (i = start; i < end; i++) {
    309:       PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
    310:       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
    311:       PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));

    313:       PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
    314:       if (vs < ncols) {
    315:         vs = PetscMin(2 * ncols, n);
    316:         PetscCall(PetscRealloc(vs * sizeof(*val), &val));
    317:       }
    318:       for (j = 0; j < ncols; j++) val[j] = a * vals[j];
    319:       PetscCall(MatSetValues(B, 1, &i, ncols, row, val, ADD_VALUES));
    320:       PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
    321:     }
    322:     PetscCall(PetscFree(val));
    323:   }
    324:   PetscCall(MatRestoreRowUpperTriangular(Y));
    325:   PetscCall(MatRestoreRowUpperTriangular(X));
    326:   PetscCall(MatGetOption(B, MAT_NO_OFF_PROC_ENTRIES, &option));
    327:   PetscCall(MatSetOption(B, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    328:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
    329:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
    330:   PetscCall(MatSetOption(B, MAT_NO_OFF_PROC_ENTRIES, option));
    331:   PetscFunctionReturn(PETSC_SUCCESS);
    332: }

    334: /*@
    335:   MatShift - Computes `Y =  Y + a I`, where `a` is a `PetscScalar`

    337:   Neighbor-wise Collective

    339:   Input Parameters:
    340: + Y - the matrix
    341: - a - the `PetscScalar`

    343:   Level: intermediate

    345:   Notes:
    346:   If `Y` is a rectangular matrix, the shift is done on the main diagonal of the matrix (https://en.wikipedia.org/wiki/Main_diagonal)

    348:   If the matrix `Y` is missing some diagonal entries this routine can be very slow. To make it fast one should initially
    349:   fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
    350:   entry). No operation is performed when a is zero.

    352:   To form Y = Y + diag(V) use `MatDiagonalSet()`

    354: .seealso: [](ch_matrices), `Mat`, `MatDiagonalSet()`, `MatScale()`, `MatDiagonalScale()`
    355:  @*/
    356: PetscErrorCode MatShift(Mat Y, PetscScalar a)
    357: {
    358:   PetscFunctionBegin;
    360:   PetscCheck(Y->assembled, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix");
    361:   PetscCheck(!Y->factortype, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
    362:   MatCheckPreallocated(Y, 1);
    363:   if (a == 0.0) PetscFunctionReturn(PETSC_SUCCESS);

    365:   if (Y->ops->shift) PetscUseTypeMethod(Y, shift, a);
    366:   else PetscCall(MatShift_Basic(Y, a));

    368:   PetscCall(PetscObjectStateIncrease((PetscObject)Y));
    369:   PetscFunctionReturn(PETSC_SUCCESS);
    370: }

    372: PetscErrorCode MatDiagonalSet_Default(Mat Y, Vec D, InsertMode is)
    373: {
    374:   PetscInt           i, start, end;
    375:   const PetscScalar *v;

    377:   PetscFunctionBegin;
    378:   PetscCall(MatGetOwnershipRange(Y, &start, &end));
    379:   PetscCall(VecGetArrayRead(D, &v));
    380:   for (i = start; i < end; i++) PetscCall(MatSetValues(Y, 1, &i, 1, &i, v + i - start, is));
    381:   PetscCall(VecRestoreArrayRead(D, &v));
    382:   PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
    383:   PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
    384:   PetscFunctionReturn(PETSC_SUCCESS);
    385: }

    387: /*@
    388:   MatDiagonalSet - Computes `Y` = `Y` + `D`, where `D` is a diagonal matrix
    389:   that is represented as a vector. Or Y[i,i] = D[i] if `InsertMode` is
    390:   `INSERT_VALUES`.

    392:   Neighbor-wise Collective

    394:   Input Parameters:
    395: + Y  - the input matrix
    396: . D  - the diagonal matrix, represented as a vector
    397: - is - `INSERT_VALUES` or `ADD_VALUES`

    399:   Level: intermediate

    401:   Note:
    402:   If the matrix `Y` is missing some diagonal entries this routine can be very slow. To make it fast one should initially
    403:   fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
    404:   entry).

    406: .seealso: [](ch_matrices), `Mat`, `MatShift()`, `MatScale()`, `MatDiagonalScale()`
    407: @*/
    408: PetscErrorCode MatDiagonalSet(Mat Y, Vec D, InsertMode is)
    409: {
    410:   PetscInt matlocal, veclocal;

    412:   PetscFunctionBegin;
    415:   MatCheckPreallocated(Y, 1);
    416:   PetscCall(MatGetLocalSize(Y, &matlocal, NULL));
    417:   PetscCall(VecGetLocalSize(D, &veclocal));
    418:   PetscCheck(matlocal == veclocal, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Number local rows of matrix %" PetscInt_FMT " does not match that of vector for diagonal %" PetscInt_FMT, matlocal, veclocal);
    419:   if (Y->ops->diagonalset) PetscUseTypeMethod(Y, diagonalset, D, is);
    420:   else PetscCall(MatDiagonalSet_Default(Y, D, is));
    421:   PetscCall(PetscObjectStateIncrease((PetscObject)Y));
    422:   PetscFunctionReturn(PETSC_SUCCESS);
    423: }

    425: /*@
    426:   MatAYPX - Computes Y = a*Y + X.

    428:   Logically Collective

    430:   Input Parameters:
    431: + a   - the `PetscScalar` multiplier
    432: . Y   - the first matrix
    433: . X   - the second matrix
    434: - str - either `SAME_NONZERO_PATTERN`, `DIFFERENT_NONZERO_PATTERN`, `UNKNOWN_NONZERO_PATTERN`, or `SUBSET_NONZERO_PATTERN` (nonzeros of `X` is a subset of `Y`'s)

    436:   Level: intermediate

    438: .seealso: [](ch_matrices), `Mat`, `MatAXPY()`
    439:  @*/
    440: PetscErrorCode MatAYPX(Mat Y, PetscScalar a, Mat X, MatStructure str)
    441: {
    442:   PetscFunctionBegin;
    443:   PetscCall(MatScale(Y, a));
    444:   PetscCall(MatAXPY(Y, 1.0, X, str));
    445:   PetscFunctionReturn(PETSC_SUCCESS);
    446: }

    448: /*@
    449:   MatComputeOperator - Computes the explicit matrix

    451:   Collective

    453:   Input Parameters:
    454: + inmat   - the matrix
    455: - mattype - the matrix type for the explicit operator

    457:   Output Parameter:
    458: . mat - the explicit  operator

    460:   Level: advanced

    462:   Note:
    463:   This computation is done by applying the operator to columns of the identity matrix.
    464:   This routine is costly in general, and is recommended for use only with relatively small systems.
    465:   Currently, this routine uses a dense matrix format if `mattype` == `NULL`.

    467: .seealso: [](ch_matrices), `Mat`, `MatConvert()`, `MatMult()`, `MatComputeOperatorTranspose()`
    468: @*/
    469: PetscErrorCode MatComputeOperator(Mat inmat, MatType mattype, Mat *mat)
    470: {
    471:   PetscFunctionBegin;
    473:   PetscAssertPointer(mat, 3);
    474:   PetscCall(MatConvert_Shell(inmat, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
    475:   PetscFunctionReturn(PETSC_SUCCESS);
    476: }

    478: /*@
    479:   MatComputeOperatorTranspose - Computes the explicit matrix representation of
    480:   a give matrix that can apply `MatMultTranspose()`

    482:   Collective

    484:   Input Parameters:
    485: + inmat   - the matrix
    486: - mattype - the matrix type for the explicit operator

    488:   Output Parameter:
    489: . mat - the explicit  operator transposed

    491:   Level: advanced

    493:   Note:
    494:   This computation is done by applying the transpose of the operator to columns of the identity matrix.
    495:   This routine is costly in general, and is recommended for use only with relatively small systems.
    496:   Currently, this routine uses a dense matrix format if `mattype` == `NULL`.

    498: .seealso: [](ch_matrices), `Mat`, `MatConvert()`, `MatMult()`, `MatComputeOperator()`
    499: @*/
    500: PetscErrorCode MatComputeOperatorTranspose(Mat inmat, MatType mattype, Mat *mat)
    501: {
    502:   Mat A;

    504:   PetscFunctionBegin;
    506:   PetscAssertPointer(mat, 3);
    507:   PetscCall(MatCreateTranspose(inmat, &A));
    508:   PetscCall(MatConvert_Shell(A, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
    509:   PetscCall(MatDestroy(&A));
    510:   PetscFunctionReturn(PETSC_SUCCESS);
    511: }

    513: /*@
    514:   MatFilter - Set all values in the matrix with an absolute value less than or equal to the tolerance to zero, and optionally compress the underlying storage

    516:   Input Parameters:
    517: + A        - The matrix
    518: . tol      - The zero tolerance
    519: . compress - Whether the storage from the input matrix `A` should be compressed once values less than or equal to `tol` are set to zero
    520: - keep     - If `compress` is true and for a given row of `A`, the diagonal coefficient is less than or equal to `tol`, indicates whether it should be left in the structure or eliminated as well

    522:   Level: intermediate

    524: .seealso: [](ch_matrices), `Mat`, `MatCreate()`, `MatZeroEntries()`, `MatEliminateZeros()`, `VecFilter()`
    525:  @*/
    526: PetscErrorCode MatFilter(Mat A, PetscReal tol, PetscBool compress, PetscBool keep)
    527: {
    528:   Mat          a;
    529:   PetscScalar *newVals;
    530:   PetscInt    *newCols, rStart, rEnd, maxRows, r, colMax = 0, nnz0 = 0, nnz1 = 0;
    531:   PetscBool    flg;

    533:   PetscFunctionBegin;
    534:   PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)A, &flg, MATSEQDENSE, MATMPIDENSE, ""));
    535:   if (flg) {
    536:     PetscCall(MatDenseGetLocalMatrix(A, &a));
    537:     PetscCall(MatDenseGetLDA(a, &r));
    538:     PetscCall(MatGetSize(a, &rStart, &rEnd));
    539:     PetscCall(MatDenseGetArray(a, &newVals));
    540:     for (; colMax < rEnd; ++colMax) {
    541:       for (maxRows = 0; maxRows < rStart; ++maxRows) newVals[maxRows + colMax * r] = PetscAbsScalar(newVals[maxRows + colMax * r]) <= tol ? 0.0 : newVals[maxRows + colMax * r];
    542:     }
    543:     PetscCall(MatDenseRestoreArray(a, &newVals));
    544:   } else {
    545:     const PetscInt *ranges;
    546:     PetscMPIInt     rank, size;

    548:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));
    549:     PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
    550:     PetscCall(MatGetOwnershipRanges(A, &ranges));
    551:     rStart = ranges[rank];
    552:     rEnd   = ranges[rank + 1];
    553:     PetscCall(MatGetRowUpperTriangular(A));
    554:     for (r = rStart; r < rEnd; ++r) {
    555:       PetscInt ncols;

    557:       PetscCall(MatGetRow(A, r, &ncols, NULL, NULL));
    558:       colMax = PetscMax(colMax, ncols);
    559:       PetscCall(MatRestoreRow(A, r, &ncols, NULL, NULL));
    560:     }
    561:     maxRows = 0;
    562:     for (r = 0; r < size; ++r) maxRows = PetscMax(maxRows, ranges[r + 1] - ranges[r]);
    563:     PetscCall(PetscCalloc2(colMax, &newCols, colMax, &newVals));
    564:     PetscCall(MatGetOption(A, MAT_NO_OFF_PROC_ENTRIES, &flg)); /* cache user-defined value */
    565:     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    566:     /* short-circuit code in MatAssemblyBegin() and MatAssemblyEnd()             */
    567:     /* that are potentially called many times depending on the distribution of A */
    568:     for (r = rStart; r < rStart + maxRows; ++r) {
    569:       if (r < rEnd) {
    570:         const PetscScalar *vals;
    571:         const PetscInt    *cols;
    572:         PetscInt           ncols, newcols = 0, c;

    574:         PetscCall(MatGetRow(A, r, &ncols, &cols, &vals));
    575:         nnz0 += ncols - 1;
    576:         for (c = 0; c < ncols; ++c) {
    577:           if (PetscUnlikely(PetscAbsScalar(vals[c]) <= tol)) newCols[newcols++] = cols[c];
    578:         }
    579:         nnz1 += ncols - newcols - 1;
    580:         PetscCall(MatRestoreRow(A, r, &ncols, &cols, &vals));
    581:         PetscCall(MatSetValues(A, 1, &r, newcols, newCols, newVals, INSERT_VALUES));
    582:       }
    583:       PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    584:       PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    585:     }
    586:     PetscCall(MatRestoreRowUpperTriangular(A));
    587:     PetscCall(PetscFree2(newCols, newVals));
    588:     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, flg)); /* reset option to its user-defined value */
    589:     if (nnz0 > 0) PetscCall(PetscInfo(NULL, "Filtering left %g%% edges in graph\n", 100 * (double)nnz1 / (double)nnz0));
    590:     else PetscCall(PetscInfo(NULL, "Warning: %" PetscInt_FMT " edges to filter with %" PetscInt_FMT " rows\n", nnz0, maxRows));
    591:   }
    592:   if (compress && A->ops->eliminatezeros) {
    593:     Mat       B;
    594:     PetscBool flg;

    596:     PetscCall(PetscObjectTypeCompareAny((PetscObject)A, &flg, MATSEQAIJHIPSPARSE, MATMPIAIJHIPSPARSE, ""));
    597:     if (!flg) {
    598:       PetscCall(MatEliminateZeros(A, keep));
    599:       PetscCall(MatDuplicate(A, MAT_COPY_VALUES, &B));
    600:       PetscCall(MatHeaderReplace(A, &B));
    601:     }
    602:   }
    603:   PetscFunctionReturn(PETSC_SUCCESS);
    604: }