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:   Level: advanced

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

821:   For the second use case call `MatSetType()` with a type of `MATDIAGONAL` followed by a call to `MatDiagonalSetDiagonal()`. 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 third use case call `MatSetType()` with a type of `MATDIAGONAL` followed by a calls to `MatSetSizes()` and `MatDiagonalSet()`. In this case
825:   the diagonal vector will not be referenced internally by the matrix, its values will be copied.

827: .seealso: [](ch_matrices), `Mat`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`,
828:           `MatDiagonalGetInverseDiagonal()`, `MatDiagonalSet()`, `MatDiagonalSetDiagonal()`
829: M*/
830: PETSC_INTERN PetscErrorCode MatCreate_Diagonal(Mat A)
831: {
832:   Mat_Diagonal *ctx;

834:   PetscFunctionBegin;
835:   PetscCall(PetscNew(&ctx));
836:   A->data = (void *)ctx;

838:   A->structurally_symmetric      = PETSC_BOOL3_TRUE;
839:   A->structural_symmetry_eternal = PETSC_TRUE;
840:   A->symmetry_eternal            = PETSC_TRUE;
841:   A->symmetric                   = PETSC_BOOL3_TRUE;
842:   if (!PetscDefined(USE_COMPLEX)) A->hermitian = PETSC_BOOL3_TRUE;

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

869:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetDiagonal_C", MatDiagonalGetDiagonal_Diagonal));
870:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreDiagonal_C", MatDiagonalRestoreDiagonal_Diagonal));
871:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetInverseDiagonal_C", MatDiagonalGetInverseDiagonal_Diagonal));
872:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreInverseDiagonal_C", MatDiagonalRestoreInverseDiagonal_Diagonal));
873:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_seqdense_C", MatProductSetFromOptions_Diagonal_Dense));
874:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_mpidense_C", MatProductSetFromOptions_Diagonal_Dense));
875:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATDIAGONAL));
876:   PetscFunctionReturn(PETSC_SUCCESS);
877: }