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: }