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 `MatDiagonalRestoreDiagonal()` before using the matrix again.
148: For a copy of the diagonal values, rather than a reference, use `MatGetDiagonal()`.
150: Any changes to the obtained vector immediately change the action of the `Mat`.
151: The matrix can be changed more efficiently by accessing this vector and changing its values, instead of filling a work vector and using `MatDiagonalSet()`
153: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`, `MatGetDiagonal()`
154: @*/
155: PetscErrorCode MatDiagonalGetDiagonal(Mat A, Vec *diag)
156: {
157: PetscFunctionBegin;
159: PetscAssertPointer(diag, 2);
160: *diag = NULL;
161: PetscUseMethod(A, "MatDiagonalGetDiagonal_C", (Mat, Vec *), (A, diag));
162: PetscFunctionReturn(PETSC_SUCCESS);
163: }
165: static PetscErrorCode MatDiagonalGetDiagonal_Diagonal(Mat A, Vec *diag)
166: {
167: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
169: PetscFunctionBegin;
170: PetscCall(MatSetUp(A));
171: PetscCall(MatDiagonalSetUpDiagonal(A));
172: *diag = ctx->diag;
173: PetscCall(PetscObjectStateGet((PetscObject)*diag, &ctx->diag_state));
174: PetscFunctionReturn(PETSC_SUCCESS);
175: }
177: /*@
178: MatDiagonalRestoreDiagonal - Restore the diagonal of a `MATDIAGONAL`
180: Input Parameters:
181: + A - the `MATDIAGONAL`
182: - diag - the `Vec` obtained from `MatDiagonalGetDiagonal()`
184: Level: developer
186: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetDiagonal()`
187: @*/
188: PetscErrorCode MatDiagonalRestoreDiagonal(Mat A, Vec *diag)
189: {
190: PetscFunctionBegin;
192: PetscAssertPointer(diag, 2);
193: PetscUseMethod(A, "MatDiagonalRestoreDiagonal_C", (Mat, Vec *), (A, diag));
194: PetscFunctionReturn(PETSC_SUCCESS);
195: }
197: static PetscErrorCode MatDiagonalRestoreDiagonal_Diagonal(Mat A, Vec *diag)
198: {
199: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
200: PetscObjectState diag_state;
202: PetscFunctionBegin;
203: PetscCheck(ctx->diag == *diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector");
204: ctx->diag_valid = PETSC_TRUE;
205: PetscCall(PetscObjectStateGet((PetscObject)*diag, &diag_state));
206: if (ctx->diag_state != diag_state) {
207: PetscCall(PetscObjectStateIncrease((PetscObject)A));
208: ctx->inv_diag_valid = PETSC_FALSE;
209: }
210: *diag = NULL;
211: PetscFunctionReturn(PETSC_SUCCESS);
212: }
214: /*@
215: MatDiagonalGetInverseDiagonal - Get the inverse diagonal of a `MATDIAGONAL`
217: Input Parameter:
218: . A - the `MATDIAGONAL`
220: Output Parameter:
221: . inv_diag - the `Vec` that defines the inverse diagonal
223: Level: developer
225: Note:
226: The user must call
227: `MatDiagonalRestoreInverseDiagonal()` before using the matrix again.
229: If a matrix is created only to call `MatSolve()` (which happens for `MATLMVMDIAGBROYDEN`),
230: using `MatDiagonalGetInverseDiagonal()` and `MatDiagonalRestoreInverseDiagonal()` avoids copies
231: and avoids any call to `VecReciprocal()`.
233: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MATLMVMBROYDEN`, `MatSolve()`
234: @*/
235: PetscErrorCode MatDiagonalGetInverseDiagonal(Mat A, Vec *inv_diag)
236: {
237: PetscFunctionBegin;
239: PetscAssertPointer(inv_diag, 2);
240: *inv_diag = NULL;
241: PetscUseMethod(A, "MatDiagonalGetInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
242: PetscFunctionReturn(PETSC_SUCCESS);
243: }
245: static PetscErrorCode MatDiagonalGetInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
246: {
247: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
249: PetscFunctionBegin;
250: PetscCall(MatSetUp(A));
251: PetscCall(MatDiagonalSetUpInverseDiagonal(A));
252: *inv_diag = ctx->inv_diag;
253: PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &ctx->inv_diag_state));
254: PetscFunctionReturn(PETSC_SUCCESS);
255: }
257: /*@
258: MatDiagonalRestoreInverseDiagonal - Restore the inverse diagonal of a `MATDIAGONAL`
260: Input Parameters:
261: + A - the `MATDIAGONAL`
262: - inv_diag - the `Vec` obtained from `MatDiagonalGetInverseDiagonal()`
264: Level: developer
266: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetInverseDiagonal()`
267: @*/
268: PetscErrorCode MatDiagonalRestoreInverseDiagonal(Mat A, Vec *inv_diag)
269: {
270: PetscFunctionBegin;
272: PetscAssertPointer(inv_diag, 2);
273: PetscUseMethod(A, "MatDiagonalRestoreInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
274: PetscFunctionReturn(PETSC_SUCCESS);
275: }
277: static PetscErrorCode MatDiagonalRestoreInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
278: {
279: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
280: PetscObjectState inv_diag_state;
282: PetscFunctionBegin;
283: PetscCheck(ctx->inv_diag == *inv_diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector");
284: ctx->inv_diag_valid = PETSC_TRUE;
285: PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &inv_diag_state));
286: if (ctx->inv_diag_state != inv_diag_state) {
287: PetscCall(PetscObjectStateIncrease((PetscObject)A));
288: ctx->diag_valid = PETSC_FALSE;
289: }
290: *inv_diag = NULL;
291: PetscFunctionReturn(PETSC_SUCCESS);
292: }
294: static PetscErrorCode MatPermute_Diagonal(Mat A, IS rowp, IS colp, Mat *B)
295: {
296: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
297: Vec v;
299: PetscFunctionBegin;
300: PetscCheck(rowp == colp, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Row permutation and column permutation must be the same");
301: PetscCall(VecDuplicate(ctx->diag, &v));
302: PetscCall(VecCopy(ctx->diag, v));
303: PetscCall(VecPermute(v, rowp, PETSC_FALSE));
304: PetscCall(MatCreateDiagonal(v, B));
305: PetscCall(VecDestroy(&v));
306: PetscFunctionReturn(PETSC_SUCCESS);
307: }
309: static PetscErrorCode MatSetRandom_Diagonal(Mat A, PetscRandom rand)
310: {
311: Vec d;
313: PetscFunctionBegin;
314: PetscCall(MatDiagonalGetDiagonal(A, &d));
315: PetscCall(VecSetRandom(d, rand));
316: PetscCall(MatDiagonalRestoreDiagonal(A, &d));
317: PetscFunctionReturn(PETSC_SUCCESS);
318: }
320: static PetscErrorCode MatDestroy_Diagonal(Mat mat)
321: {
322: Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data;
324: PetscFunctionBegin;
325: PetscCall(VecDestroy(&ctx->diag));
326: PetscCall(VecDestroy(&ctx->inv_diag));
327: PetscCall(PetscFree(ctx->col));
328: PetscCall(PetscFree(ctx->val));
329: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetDiagonal_C", NULL));
330: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreDiagonal_C", NULL));
331: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetInverseDiagonal_C", NULL));
332: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreInverseDiagonal_C", NULL));
333: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_seqdense_C", NULL));
334: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_mpidense_C", NULL));
335: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_diagonal_C", NULL));
336: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_anytype_C", NULL));
337: PetscCall(PetscFree(mat->data));
338: PetscFunctionReturn(PETSC_SUCCESS);
339: }
341: static PetscErrorCode MatView_Diagonal(Mat J, PetscViewer viewer)
342: {
343: Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
344: PetscBool isascii;
346: PetscFunctionBegin;
347: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
348: if (isascii) {
349: PetscViewerFormat format;
351: PetscCall(MatDiagonalSetUpDiagonal(J));
352: PetscCall(PetscViewerGetFormat(viewer, &format));
353: if (format == PETSC_VIEWER_ASCII_FACTOR_INFO || format == PETSC_VIEWER_ASCII_INFO) {
354: PetscCall(PetscObjectPrintClassNamePrefixType((PetscObject)ctx->diag, viewer));
355: PetscFunctionReturn(PETSC_SUCCESS);
356: }
357: PetscCall(VecView(ctx->diag, viewer));
358: }
359: PetscFunctionReturn(PETSC_SUCCESS);
360: }
362: static PetscErrorCode MatGetDiagonal_Diagonal(Mat J, Vec x)
363: {
364: Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
366: PetscFunctionBegin;
367: PetscCall(MatDiagonalSetUpDiagonal(J));
368: PetscCall(VecCopy(ctx->diag, x));
369: PetscFunctionReturn(PETSC_SUCCESS);
370: }
372: static PetscErrorCode MatADot_Diagonal_Local(Mat A, Vec x, Vec y, PetscScalar *val)
373: {
374: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
375: PetscInt n = x->map->n;
376: const PetscScalar *ya, *xa, *wa;
377: PetscScalar sum = 0;
379: PetscFunctionBegin;
380: PetscCheckSameComm(x, 2, ctx->diag, 1);
381: PetscCheckSameComm(y, 3, ctx->diag, 1);
382: PetscCall(VecGetArrayRead(x, &xa));
383: PetscCall(VecGetArrayRead(y, &ya));
384: PetscCall(VecGetArrayRead(ctx->diag, &wa));
385: for (PetscInt i = 0; i < n; i++) {
386: sum += PetscConj(ya[i]) * wa[i] * xa[i];
387: }
388: if (n > 0) PetscCall(PetscLogFlops(3.0 * n));
389: PetscCall(VecRestoreArrayRead(x, &xa));
390: PetscCall(VecRestoreArrayRead(y, &ya));
391: PetscCall(VecRestoreArrayRead(ctx->diag, &wa));
392: *val = sum;
393: PetscFunctionReturn(PETSC_SUCCESS);
394: }
396: static PetscErrorCode MatADot_Diagonal_MPI(Mat A, Vec x, Vec y, PetscScalar *val)
397: {
398: PetscFunctionBegin;
399: PetscCall(MatDiagonalSetUpDiagonal(A));
400: PetscUseTypeMethod(A, adot_local, x, y, val);
401: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, val, 1, MPIU_SCALAR, MPIU_SUM, PetscObjectComm((PetscObject)A)));
402: PetscFunctionReturn(PETSC_SUCCESS);
403: }
405: static PetscErrorCode MatANorm_Diagonal_Local(Mat A, Vec x, PetscReal *val)
406: {
407: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
408: PetscInt n = x->map->n;
409: const PetscScalar *xa, *wa;
410: PetscScalar sum = 0;
412: PetscFunctionBegin;
413: PetscCheckSameComm(x, 2, ctx->diag, 1);
414: PetscCall(VecGetArrayRead(x, &xa));
415: PetscCall(VecGetArrayRead(ctx->diag, &wa));
416: for (PetscInt i = 0; i < n; i++) {
417: sum += PetscConj(xa[i]) * wa[i] * xa[i];
418: }
419: if (n > 0) PetscCall(PetscLogFlops(3.0 * n));
420: PetscCall(VecRestoreArrayRead(x, &xa));
421: PetscCall(VecRestoreArrayRead(ctx->diag, &wa));
422: PetscCheck(PetscAbsReal(PetscImaginaryPart(sum)) <= 100 * PETSC_MACHINE_EPSILON * PetscMax(1.0, PetscAbsScalar(sum)), PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Matrix argument is not Hermitian (diagonal has nonzero imaginary parts)");
423: *val = PetscRealPart(sum);
424: PetscFunctionReturn(PETSC_SUCCESS);
425: }
427: static PetscErrorCode MatANorm_Diagonal_MPI(Mat A, Vec x, PetscReal *val)
428: {
429: PetscFunctionBegin;
430: PetscCall(MatDiagonalSetUpDiagonal(A));
431: PetscUseTypeMethod(A, anorm_local, x, val);
432: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, val, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
433: *val = PetscSqrtReal(*val);
434: PetscFunctionReturn(PETSC_SUCCESS);
435: }
437: static PetscErrorCode MatANorm_Diagonal_Seq(Mat A, Vec x, PetscReal *val)
438: {
439: PetscFunctionBegin;
440: PetscCall(MatDiagonalSetUpDiagonal(A));
441: PetscUseTypeMethod(A, anorm_local, x, val);
442: PetscCheck(*val >= 0.0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Matrix argument is not positive definite (diagonal has negative entries)");
443: *val = PetscSqrtReal(*val);
444: PetscFunctionReturn(PETSC_SUCCESS);
445: }
447: /*@
448: MatDiagonalSetDiagonal - Sets the diagonal for a `MATDIAGONAL`
450: Collective
452: Input Parameter:
453: + J - the `MATDIAGONAL` matrix
454: - diag - the vector for the diagonal
456: Level: advanced
458: Notes:
459: The input vector `diag` will be referenced internally: any changes to `diag`
460: will affect the matrix `J`.
462: This routine can only be called once for the given `J` matrix.
464: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MatCreateDiagonal()`, `MATDIAGONAL`, `MatScale()`, `MatShift()`, `MatMult()`, `MatGetDiagonal()`, `MatSolve()`,
465: `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`,
466: `MATCONSTANTDIAGONAL`
467: @*/
468: PetscErrorCode MatDiagonalSetDiagonal(Mat J, Vec diag)
469: {
470: Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
471: PetscMPIInt comm_size;
472: PetscLayout map;
473: VecType type;
474: PetscBool iskokkos, ismpi, iscuda, iship;
476: PetscFunctionBegin;
479: PetscCheckSameComm(J, 1, diag, 2);
480: PetscCheck(!ctx->diag, PetscObjectComm((PetscObject)J), PETSC_ERR_ARG_WRONGSTATE, "MATDIAGONAL already has a diagonal");
481: ctx->diag = diag;
482: PetscCall(PetscObjectReference((PetscObject)ctx->diag));
483: PetscCall(VecDuplicate(ctx->diag, &ctx->inv_diag));
484: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)J), &comm_size));
485: PetscCall(VecGetLayout(ctx->diag, &map));
486: PetscCall(MatSetLayouts(J, map, map));
487: PetscCall(PetscFree(J->defaultvectype));
488: PetscCall(VecGetType(ctx->diag, &type));
489: PetscCall(PetscStrallocpy(type, &J->defaultvectype));
491: PetscCall(PetscObjectTypeCompareAny((PetscObject)ctx->diag, &iscuda, VECCUDA, VECMPICUDA, VECSEQCUDA, ""));
492: PetscCall(PetscObjectTypeCompareAny((PetscObject)ctx->diag, &iship, VECHIP, VECMPIHIP, VECSEQHIP, ""));
493: PetscCall(PetscObjectTypeCompareAny((PetscObject)ctx->diag, &iskokkos, VECKOKKOS, VECMPIKOKKOS, VECSEQKOKKOS, ""));
494: PetscCall(PetscObjectBaseTypeCompare((PetscObject)ctx->diag, VECMPI, &ismpi));
495: ismpi = ismpi || comm_size > 1;
496: J->ops->adot_local = MatADot_Diagonal_Local;
497: J->ops->adot = MatADot_Diagonal_Local;
498: J->ops->anorm_local = MatANorm_Diagonal_Local;
499: if (iskokkos) {
500: PetscCheck(PetscDefined(HAVE_KOKKOS_KERNELS), PetscObjectComm((PetscObject)ctx->diag), PETSC_ERR_SUP, "Reconfigure using KOKKOS kernels support");
501: #if PetscDefined(HAVE_KOKKOS_KERNELS)
502: J->ops->adot_local = MatADot_Diagonal_SeqKokkos;
503: J->ops->adot = MatADot_Diagonal_SeqKokkos;
504: J->ops->anorm_local = MatANormSq_Diagonal_SeqKokkos;
505: #endif
506: }
507: if (iscuda) {
508: PetscCheck(PetscDefined(HAVE_CUDA), PetscObjectComm((PetscObject)ctx->diag), PETSC_ERR_SUP, "Reconfigure using CUDA support");
509: #if PetscDefined(HAVE_CUDA)
510: J->ops->adot_local = MatADot_Diagonal_SeqCUDA;
511: J->ops->adot = MatADot_Diagonal_SeqCUDA;
512: J->ops->anorm_local = MatANormSq_Diagonal_SeqCUDA;
513: #endif
514: }
515: if (iship) {
516: PetscCheck(PetscDefined(HAVE_HIP), PetscObjectComm((PetscObject)ctx->diag), PETSC_ERR_SUP, "Reconfigure using HIP support");
517: #if PetscDefined(HAVE_HIP)
518: J->ops->adot_local = MatADot_Diagonal_SeqHIP;
519: J->ops->adot = MatADot_Diagonal_SeqHIP;
520: J->ops->anorm_local = MatANormSq_Diagonal_SeqHIP;
521: #endif
522: }
523: if (ismpi) {
524: J->ops->adot = MatADot_Diagonal_MPI;
525: J->ops->anorm = MatANorm_Diagonal_MPI;
526: } else {
527: J->ops->anorm = MatANorm_Diagonal_Seq;
528: }
529: ctx->col = NULL;
530: ctx->val = NULL;
531: ctx->diag_valid = PETSC_TRUE;
532: J->assembled = PETSC_TRUE;
533: PetscFunctionReturn(PETSC_SUCCESS);
534: }
536: static PetscErrorCode MatDiagonalSet_Diagonal(Mat J, Vec D, InsertMode is)
537: {
538: Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
540: PetscFunctionBegin;
541: switch (is) {
542: case ADD_VALUES:
543: case ADD_ALL_VALUES:
544: case ADD_BC_VALUES:
545: PetscCall(MatDiagonalSetUpDiagonal(J));
546: PetscCall(VecAXPY(ctx->diag, 1.0, D));
547: ctx->inv_diag_valid = PETSC_FALSE;
548: break;
549: case INSERT_VALUES:
550: case INSERT_BC_VALUES:
551: case INSERT_ALL_VALUES:
552: PetscCall(VecCopy(D, ctx->diag));
553: ctx->diag_valid = PETSC_TRUE;
554: ctx->inv_diag_valid = PETSC_FALSE;
555: break;
556: case MAX_VALUES:
557: PetscCall(MatDiagonalSetUpDiagonal(J));
558: PetscCall(VecPointwiseMax(ctx->diag, D, ctx->diag));
559: ctx->inv_diag_valid = PETSC_FALSE;
560: break;
561: case MIN_VALUES:
562: PetscCall(MatDiagonalSetUpDiagonal(J));
563: PetscCall(VecPointwiseMin(ctx->diag, D, ctx->diag));
564: ctx->inv_diag_valid = PETSC_FALSE;
565: break;
566: case NOT_SET_VALUES:
567: break;
568: }
569: PetscFunctionReturn(PETSC_SUCCESS);
570: }
572: static PetscErrorCode MatShift_Diagonal(Mat Y, PetscScalar a)
573: {
574: Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;
576: PetscFunctionBegin;
577: PetscCall(MatDiagonalSetUpDiagonal(Y));
578: PetscCall(VecShift(ctx->diag, a));
579: ctx->inv_diag_valid = PETSC_FALSE;
580: PetscFunctionReturn(PETSC_SUCCESS);
581: }
583: static PetscErrorCode MatScale_Diagonal(Mat Y, PetscScalar a)
584: {
585: Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;
587: PetscFunctionBegin;
588: PetscCall(MatDiagonalSetUpDiagonal(Y));
589: PetscCall(VecScale(ctx->diag, a));
590: ctx->inv_diag_valid = PETSC_FALSE;
591: PetscFunctionReturn(PETSC_SUCCESS);
592: }
594: static PetscErrorCode MatDiagonalScale_Diagonal(Mat Y, Vec l, Vec r)
595: {
596: Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;
598: PetscFunctionBegin;
599: PetscCall(MatDiagonalSetUpDiagonal(Y));
600: if (l) {
601: PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, l));
602: ctx->inv_diag_valid = PETSC_FALSE;
603: }
604: if (r) {
605: PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, r));
606: ctx->inv_diag_valid = PETSC_FALSE;
607: }
608: PetscFunctionReturn(PETSC_SUCCESS);
609: }
611: static PetscErrorCode MatConjugate_Diagonal(Mat Y)
612: {
613: PetscFunctionBegin;
614: Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;
616: if (ctx->diag_valid) {
617: PetscCall(VecConjugate(ctx->diag));
618: PetscCall(PetscObjectStateGet((PetscObject)ctx->diag, &ctx->diag_state));
619: }
620: if (ctx->inv_diag_valid) {
621: PetscCall(VecConjugate(ctx->inv_diag));
622: PetscCall(PetscObjectStateGet((PetscObject)ctx->inv_diag, &ctx->inv_diag_state));
623: }
624: PetscFunctionReturn(PETSC_SUCCESS);
625: }
627: static PetscErrorCode MatTranspose_Diagonal(Mat A, MatReuse reuse, Mat *matout)
628: {
629: PetscFunctionBegin;
630: if (reuse == MAT_INPLACE_MATRIX) {
631: PetscLayout tmplayout = A->rmap;
633: A->rmap = A->cmap;
634: A->cmap = tmplayout;
635: } else {
636: Vec diag, newdiag;
637: if (reuse == MAT_INITIAL_MATRIX) {
638: PetscCall(MatDiagonalGetDiagonal(A, &diag));
639: PetscCall(VecDuplicate(diag, &newdiag));
640: PetscCall(VecCopy(diag, newdiag));
641: PetscCall(MatDiagonalRestoreDiagonal(A, &diag));
642: PetscCall(MatCreateDiagonal(newdiag, matout));
643: PetscCall(VecDestroy(&newdiag));
644: } else {
645: PetscCall(MatDiagonalGetDiagonal(A, &diag));
646: PetscCall(MatDiagonalGetDiagonal(*matout, &newdiag));
647: PetscCall(VecCopy(diag, newdiag));
648: PetscCall(MatDiagonalRestoreDiagonal(*matout, &newdiag));
649: PetscCall(MatDiagonalRestoreDiagonal(A, &diag));
650: }
651: }
652: PetscFunctionReturn(PETSC_SUCCESS);
653: }
655: static PetscErrorCode MatSetUp_Diagonal(Mat A)
656: {
657: Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;
659: PetscFunctionBegin;
660: if (!ctx->diag) {
661: Vec diag;
663: PetscCall(PetscLayoutSetUp(A->cmap));
664: PetscCall(PetscLayoutSetUp(A->rmap));
665: PetscCall(MatCreateVecs(A, &diag, NULL));
666: PetscCall(MatDiagonalSetDiagonal(A, diag));
667: PetscCall(VecDestroy(&diag));
668: }
669: A->assembled = PETSC_TRUE;
670: PetscFunctionReturn(PETSC_SUCCESS);
671: }
673: static PetscErrorCode MatZeroEntries_Diagonal(Mat Y)
674: {
675: Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;
677: PetscFunctionBegin;
678: PetscCall(VecZeroEntries(ctx->diag));
679: ctx->diag_valid = PETSC_TRUE;
680: ctx->inv_diag_valid = PETSC_FALSE;
681: PetscFunctionReturn(PETSC_SUCCESS);
682: }
684: static PetscErrorCode MatSolve_Diagonal(Mat matin, Vec b, Vec x)
685: {
686: Mat_Diagonal *ctx = (Mat_Diagonal *)matin->data;
688: PetscFunctionBegin;
689: PetscCall(MatDiagonalSetUpInverseDiagonal(matin));
690: PetscCall(VecPointwiseMult(x, b, ctx->inv_diag));
691: PetscFunctionReturn(PETSC_SUCCESS);
692: }
694: static PetscErrorCode MatGetInfo_Diagonal(Mat A, MatInfoType flag, MatInfo *info)
695: {
696: PetscFunctionBegin;
697: info->block_size = 1.0;
698: info->nz_allocated = A->cmap->N;
699: info->nz_used = A->cmap->N;
700: info->nz_unneeded = 0.0;
701: info->assemblies = A->num_ass;
702: info->mallocs = 0.0;
703: info->memory = 0;
704: info->fill_ratio_given = 0;
705: info->fill_ratio_needed = 0;
706: info->factor_mallocs = 0;
707: PetscFunctionReturn(PETSC_SUCCESS);
708: }
710: /*@
711: MatCreateDiagonal - Creates a matrix defined by a given vector along its diagonal.
713: Collective
715: Input Parameter:
716: . diag - vector for the diagonal
718: Output Parameter:
719: . J - the diagonal matrix
721: Level: advanced
723: Notes:
724: Only supports square matrices with the same number of local rows and columns.
726: The input vector `diag` will be referenced internally: any changes to `diag`
727: will affect the matrix `J`.
729: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MATDIAGONAL`, `MatScale()`, `MatShift()`, `MatMult()`, `MatGetDiagonal()`, `MatSolve()`,
730: `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`,
731: `MATCONSTANTDIAGONAL`
732: @*/
733: PetscErrorCode MatCreateDiagonal(Vec diag, Mat *J)
734: {
735: PetscFunctionBegin;
737: PetscCall(MatCreate(PetscObjectComm((PetscObject)diag), J));
738: PetscCall(MatSetType(*J, MATDIAGONAL));
739: PetscCall(MatDiagonalSetDiagonal(*J, diag));
740: PetscFunctionReturn(PETSC_SUCCESS);
741: }
743: static PetscErrorCode MatProductNumeric_Diagonal_Dense(Mat C)
744: {
745: Mat A, B;
746: Mat_Diagonal *a;
747: PetscScalar *c;
748: const PetscScalar *b, *alpha;
749: PetscInt ldb, ldc;
751: PetscFunctionBegin;
752: MatCheckProduct(C, 1);
753: A = C->product->A;
754: B = C->product->B;
755: PetscCall(MatDiagonalSetUpDiagonal(A));
756: a = (Mat_Diagonal *)A->data;
757: PetscCall(VecGetArrayRead(a->diag, &alpha));
758: PetscCall(MatDenseGetLDA(B, &ldb));
759: PetscCall(MatDenseGetLDA(C, &ldc));
760: PetscCall(MatDenseGetArrayRead(B, &b));
761: PetscCall(MatDenseGetArrayWrite(C, &c));
762: for (PetscInt j = 0; j < B->cmap->N; j++)
763: for (PetscInt i = 0; i < B->rmap->n; i++) c[i + j * ldc] = alpha[i] * b[i + j * ldb];
764: PetscCall(MatDenseRestoreArrayWrite(C, &c));
765: PetscCall(MatDenseRestoreArrayRead(B, &b));
766: PetscCall(VecRestoreArrayRead(a->diag, &alpha));
767: PetscCall(PetscLogFlops(B->cmap->N * B->rmap->n));
768: PetscFunctionReturn(PETSC_SUCCESS);
769: }
771: static PetscErrorCode MatProductSymbolic_Diagonal_Dense(Mat C)
772: {
773: Mat A, B;
774: PetscInt n, N, m, M;
776: PetscFunctionBegin;
777: MatCheckProduct(C, 1);
778: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
779: A = C->product->A;
780: B = C->product->B;
781: PetscCall(MatGetLocalSize(B, NULL, &n));
782: PetscCall(MatGetSize(B, NULL, &N));
783: PetscCall(MatGetLocalSize(A, &m, NULL));
784: PetscCall(MatGetSize(A, &M, NULL));
785: PetscCall(MatSetSizes(C, m, n, M, N));
786: PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
787: PetscCall(MatSetVecType(C, B->defaultvectype));
788: PetscCall(MatSetUp(C));
789: C->ops->productnumeric = MatProductNumeric_Diagonal_Dense;
790: PetscFunctionReturn(PETSC_SUCCESS);
791: }
793: /* PtAP for diagonal * diagonal: C_ii = d_A_i * d_P_i^2 */
794: static PetscErrorCode MatProductNumeric_PtAP_Diagonal_Diagonal(Mat C)
795: {
796: Mat A = C->product->A, P = C->product->B;
797: Mat_Diagonal *a = (Mat_Diagonal *)A->data, *p = (Mat_Diagonal *)P->data, *c = (Mat_Diagonal *)C->data;
799: PetscFunctionBegin;
800: MatCheckProduct(C, 1);
801: PetscCall(MatDiagonalSetUpDiagonal(A));
802: PetscCall(MatDiagonalSetUpDiagonal(P));
803: PetscCall(VecPointwiseMult(c->diag, a->diag, p->diag));
804: PetscCall(VecPointwiseMult(c->diag, c->diag, p->diag));
805: c->diag_valid = PETSC_TRUE;
806: c->inv_diag_valid = PETSC_FALSE;
807: PetscFunctionReturn(PETSC_SUCCESS);
808: }
810: static PetscErrorCode MatProductSymbolic_PtAP_Diagonal_Diagonal(Mat C)
811: {
812: Mat P = C->product->B;
813: Mat_Diagonal *p = (Mat_Diagonal *)P->data;
814: Vec cdiag;
816: PetscFunctionBegin;
817: MatCheckProduct(C, 1);
818: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
819: PetscCall(MatSetSizes(C, P->cmap->n, P->cmap->n, P->cmap->N, P->cmap->N));
820: PetscCall(MatSetType(C, MATDIAGONAL));
821: /* Duplicate P's diagonal vec so C inherits the correct vec type (e.g., Kokkos, CUDA, HIP) */
822: PetscCall(MatDiagonalSetUpDiagonal(P));
823: PetscCall(VecDuplicate(p->diag, &cdiag));
824: PetscCall(MatDiagonalSetDiagonal(C, cdiag));
825: PetscCall(VecDestroy(&cdiag));
826: C->assembled = PETSC_TRUE;
827: C->ops->productnumeric = MatProductNumeric_PtAP_Diagonal_Diagonal;
828: PetscFunctionReturn(PETSC_SUCCESS);
829: }
831: /* PtAP for any * diagonal: C = D * A * D (bilateral diagonal scaling) */
832: static PetscErrorCode MatProductNumeric_PtAP_Anytype_Diagonal(Mat C)
833: {
834: Mat A = C->product->A, P = C->product->B;
835: Mat_Diagonal *p = (Mat_Diagonal *)P->data;
837: PetscFunctionBegin;
838: MatCheckProduct(C, 1);
839: PetscCall(MatDiagonalSetUpDiagonal(P));
840: PetscCall(MatCopy(A, C, SAME_NONZERO_PATTERN));
841: PetscCall(MatDiagonalScale(C, p->diag, p->diag));
842: PetscFunctionReturn(PETSC_SUCCESS);
843: }
845: static PetscErrorCode MatProductSymbolic_PtAP_Anytype_Diagonal(Mat C)
846: {
847: Mat A = C->product->A;
848: Mat_Product *product = C->product;
849: Mat Cwork;
851: PetscFunctionBegin;
852: MatCheckProduct(C, 1);
853: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
854: PetscCall(MatDuplicate(A, MAT_DO_NOT_COPY_VALUES, &Cwork));
855: C->product = NULL;
856: PetscCall(MatHeaderReplace(C, &Cwork));
857: C->product = product;
858: C->ops->productnumeric = MatProductNumeric_PtAP_Anytype_Diagonal;
859: PetscFunctionReturn(PETSC_SUCCESS);
860: }
862: /* PtAP for diagonal A and any non-diagonal P: C = P^T * A * P
863: Computed as P^T * AP where AP = A*P is a diagonal-scaled copy of P.
864: The Unsafe decomposition is not used because it requires the old-style
865: transposematmultnumeric pointer, which GPU types do not set. */
866: typedef struct {
867: Mat AP; /* AP = A*P (scaled copy of P) */
868: Mat PtAP; /* P^T * AP result via MatProduct AtB */
869: PetscObjectState pnnzstate; /* P's nonzero state when inner symbolic was last built */
870: } MatProductCtx_PtAP_Diagonal_Anytype;
872: static PetscErrorCode MatProductCtxDestroy_PtAP_Diagonal_Anytype(PetscCtxRt data)
873: {
874: MatProductCtx_PtAP_Diagonal_Anytype *ctx = *(MatProductCtx_PtAP_Diagonal_Anytype **)data;
876: PetscFunctionBegin;
877: PetscCall(MatDestroy(&ctx->AP));
878: PetscCall(MatDestroy(&ctx->PtAP));
879: PetscCall(PetscFree(ctx));
880: PetscFunctionReturn(PETSC_SUCCESS);
881: }
883: static PetscErrorCode MatProductNumeric_PtAP_Diagonal_Anytype(Mat C)
884: {
885: Mat_Product *product = C->product;
886: Mat A = product->A, P = product->B;
887: MatProductCtx_PtAP_Diagonal_Anytype *ctx = (MatProductCtx_PtAP_Diagonal_Anytype *)product->data;
888: Mat_Diagonal *a = (Mat_Diagonal *)A->data;
889: PetscObjectState pnnzstate;
891: PetscFunctionBegin;
892: MatCheckProduct(C, 1);
893: /* Update AP = A * P */
894: PetscCall(MatDiagonalSetUpDiagonal(A));
895: PetscCall(MatCopy(P, ctx->AP, SAME_NONZERO_PATTERN));
896: PetscCall(MatDiagonalScale(ctx->AP, a->diag, NULL));
897: /* Rebuild inner symbolic if P's nonzero structure has changed */
898: PetscCall(MatGetNonzeroState(P, &pnnzstate));
899: if (pnnzstate != ctx->pnnzstate) {
900: PetscCall(MatDestroy(&ctx->PtAP));
901: PetscCall(MatProductCreate(P, ctx->AP, NULL, &ctx->PtAP));
902: PetscCall(MatProductSetType(ctx->PtAP, MATPRODUCT_AtB));
903: PetscCall(MatProductSetFill(ctx->PtAP, product->fill));
904: PetscCall(MatProductSetFromOptions(ctx->PtAP));
905: PetscCall(MatProductSymbolic(ctx->PtAP));
906: ctx->pnnzstate = pnnzstate;
907: }
908: /* Recompute P^T * AP */
909: PetscCall(MatProductNumeric(ctx->PtAP));
910: PetscCall(MatCopy(ctx->PtAP, C, SAME_NONZERO_PATTERN));
911: PetscFunctionReturn(PETSC_SUCCESS);
912: }
914: static PetscErrorCode MatProductSymbolic_PtAP_Diagonal_Anytype(Mat C)
915: {
916: Mat_Product *product = C->product;
917: Mat P = product->B;
918: MatProductCtx_PtAP_Diagonal_Anytype *ctx;
919: Mat Cwork;
921: PetscFunctionBegin;
922: MatCheckProduct(C, 1);
923: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
924: PetscCall(PetscNew(&ctx));
926: /* AP = A * P (structure only; numeric phase fills values) */
927: PetscCall(MatDuplicate(P, MAT_DO_NOT_COPY_VALUES, &ctx->AP));
929: /* PtAP = P^T * AP (symbolic only: D*P has same sparsity as P) */
930: PetscCall(MatProductCreate(P, ctx->AP, NULL, &ctx->PtAP));
931: PetscCall(MatProductSetType(ctx->PtAP, MATPRODUCT_AtB));
932: PetscCall(MatProductSetFill(ctx->PtAP, product->fill));
933: PetscCall(MatProductSetFromOptions(ctx->PtAP));
934: PetscCall(MatProductSymbolic(ctx->PtAP));
936: /* Record P's nonzero state so numeric phase can detect structural changes */
937: PetscCall(MatGetNonzeroState(P, &ctx->pnnzstate));
939: /* Set up C with the same structure as PtAP */
940: PetscCall(MatDuplicate(ctx->PtAP, MAT_DO_NOT_COPY_VALUES, &Cwork));
941: C->product = NULL;
942: PetscCall(MatHeaderReplace(C, &Cwork));
943: C->product = product;
944: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
945: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
946: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
947: product->data = ctx;
948: product->destroy = MatProductCtxDestroy_PtAP_Diagonal_Anytype;
949: C->ops->productnumeric = MatProductNumeric_PtAP_Diagonal_Anytype;
950: PetscFunctionReturn(PETSC_SUCCESS);
951: }
953: /* AB for MATDIAGONAL A and any non-diagonal B: C = D * B (row scaling of B) */
954: static PetscErrorCode MatProductNumeric_AB_Diagonal_Anytype(Mat C)
955: {
956: Mat A = C->product->A, B = C->product->B;
957: Mat_Diagonal *a = (Mat_Diagonal *)A->data;
959: PetscFunctionBegin;
960: MatCheckProduct(C, 1);
961: PetscCall(MatDiagonalSetUpDiagonal(A));
962: PetscCall(MatCopy(B, C, SAME_NONZERO_PATTERN));
963: PetscCall(MatDiagonalScale(C, a->diag, NULL));
964: PetscFunctionReturn(PETSC_SUCCESS);
965: }
967: static PetscErrorCode MatProductSymbolic_AB_Diagonal_Anytype(Mat C)
968: {
969: Mat B = C->product->B;
970: Mat_Product *product = C->product;
971: Mat Cwork;
973: PetscFunctionBegin;
974: MatCheckProduct(C, 1);
975: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
976: PetscCall(MatDuplicate(B, MAT_DO_NOT_COPY_VALUES, &Cwork));
977: C->product = NULL;
978: PetscCall(MatHeaderReplace(C, &Cwork));
979: C->product = product;
980: C->ops->productnumeric = MatProductNumeric_AB_Diagonal_Anytype;
981: PetscFunctionReturn(PETSC_SUCCESS);
982: }
984: /* AB for any non-diagonal A and MATDIAGONAL B: C = A * D (column scaling of A) */
985: static PetscErrorCode MatProductNumeric_AB_Anytype_Diagonal(Mat C)
986: {
987: Mat A = C->product->A, B = C->product->B;
988: Mat_Diagonal *b = (Mat_Diagonal *)B->data;
990: PetscFunctionBegin;
991: MatCheckProduct(C, 1);
992: PetscCall(MatDiagonalSetUpDiagonal(B));
993: PetscCall(MatCopy(A, C, SAME_NONZERO_PATTERN));
994: PetscCall(MatDiagonalScale(C, NULL, b->diag));
995: PetscFunctionReturn(PETSC_SUCCESS);
996: }
998: static PetscErrorCode MatProductSymbolic_AB_Anytype_Diagonal(Mat C)
999: {
1000: Mat A = C->product->A;
1001: Mat_Product *product = C->product;
1002: Mat Cwork;
1004: PetscFunctionBegin;
1005: MatCheckProduct(C, 1);
1006: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
1007: PetscCall(MatDuplicate(A, MAT_DO_NOT_COPY_VALUES, &Cwork));
1008: C->product = NULL;
1009: PetscCall(MatHeaderReplace(C, &Cwork));
1010: C->product = product;
1011: C->ops->productnumeric = MatProductNumeric_AB_Anytype_Diagonal;
1012: PetscFunctionReturn(PETSC_SUCCESS);
1013: }
1015: /* AB for MATDIAGONAL * MATDIAGONAL: C_ii = d_A_i * d_B_i */
1016: static PetscErrorCode MatProductNumeric_AB_Diagonal_Diagonal(Mat C)
1017: {
1018: Mat A = C->product->A, B = C->product->B;
1019: Mat_Diagonal *a = (Mat_Diagonal *)A->data, *b = (Mat_Diagonal *)B->data, *c = (Mat_Diagonal *)C->data;
1021: PetscFunctionBegin;
1022: MatCheckProduct(C, 1);
1023: PetscCall(MatDiagonalSetUpDiagonal(A));
1024: PetscCall(MatDiagonalSetUpDiagonal(B));
1025: PetscCall(VecPointwiseMult(c->diag, a->diag, b->diag));
1026: c->diag_valid = PETSC_TRUE;
1027: c->inv_diag_valid = PETSC_FALSE;
1028: PetscFunctionReturn(PETSC_SUCCESS);
1029: }
1031: static PetscErrorCode MatProductSymbolic_AB_Diagonal_Diagonal(Mat C)
1032: {
1033: Mat A = C->product->A, B = C->product->B;
1034: Mat_Diagonal *b = (Mat_Diagonal *)B->data;
1035: Vec cdiag;
1037: PetscFunctionBegin;
1038: MatCheckProduct(C, 1);
1039: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
1040: PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N));
1041: PetscCall(MatSetType(C, MATDIAGONAL));
1042: /* Duplicate B's diagonal Vec so C inherits the correct VecType (e.g., Kokkos, CUDA, HIP) */
1043: PetscCall(MatDiagonalSetUpDiagonal(B));
1044: PetscCall(VecDuplicate(b->diag, &cdiag));
1045: PetscCall(MatDiagonalSetDiagonal(C, cdiag));
1046: PetscCall(VecDestroy(&cdiag));
1047: C->assembled = PETSC_TRUE;
1048: C->ops->productnumeric = MatProductNumeric_AB_Diagonal_Diagonal;
1049: PetscFunctionReturn(PETSC_SUCCESS);
1050: }
1052: static PetscErrorCode MatProductSetFromOptions_Diagonal_Diagonal(Mat C)
1053: {
1054: Mat_Product *product = C->product;
1056: PetscFunctionBegin;
1057: if (product->type == MATPRODUCT_PtAP) C->ops->productsymbolic = MatProductSymbolic_PtAP_Diagonal_Diagonal;
1058: else if (product->type == MATPRODUCT_AB) C->ops->productsymbolic = MatProductSymbolic_AB_Diagonal_Diagonal;
1059: PetscFunctionReturn(PETSC_SUCCESS);
1060: }
1062: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense_AB(Mat C)
1063: {
1064: PetscFunctionBegin;
1065: C->ops->productsymbolic = MatProductSymbolic_Diagonal_Dense;
1066: PetscFunctionReturn(PETSC_SUCCESS);
1067: }
1069: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense(Mat C)
1070: {
1071: Mat_Product *product = C->product;
1073: PetscFunctionBegin;
1074: if (product->type == MATPRODUCT_AB || product->type == MATPRODUCT_AtB) PetscCall(MatProductSetFromOptions_Diagonal_Dense_AB(C));
1075: else if (product->type == MATPRODUCT_PtAP) C->ops->productsymbolic = MatProductSymbolic_PtAP_Diagonal_Anytype;
1076: PetscFunctionReturn(PETSC_SUCCESS);
1077: }
1079: /*
1080: Fallback dispatcher for matrix type combinations that have no specific registration
1081: (e.g., GPU types like seqaijkokkos). Queried via the "anytype" mechanism in
1082: MatProductSetFromOptions_Private() after specific type queries fail.
1083: */
1084: static PetscErrorCode MatProductSetFromOptions_Diagonal_Anytype(Mat C)
1085: {
1086: Mat_Product *product = C->product;
1087: PetscBool Adiag, Bdiag;
1089: PetscFunctionBegin;
1090: PetscCall(PetscObjectTypeCompare((PetscObject)product->A, MATDIAGONAL, &Adiag));
1091: PetscCall(PetscObjectTypeCompare((PetscObject)product->B, MATDIAGONAL, &Bdiag));
1092: if (Adiag && !Bdiag && (product->type == MATPRODUCT_PtAP)) C->ops->productsymbolic = MatProductSymbolic_PtAP_Diagonal_Anytype;
1093: else if (Bdiag && !Adiag && (product->type == MATPRODUCT_PtAP)) C->ops->productsymbolic = MatProductSymbolic_PtAP_Anytype_Diagonal;
1094: else if (Adiag && !Bdiag && (product->type == MATPRODUCT_AB)) C->ops->productsymbolic = MatProductSymbolic_AB_Diagonal_Anytype;
1095: else if (Bdiag && !Adiag && (product->type == MATPRODUCT_AB)) C->ops->productsymbolic = MatProductSymbolic_AB_Anytype_Diagonal;
1096: PetscFunctionReturn(PETSC_SUCCESS);
1097: }
1099: /*MC
1100: MATDIAGONAL - MATDIAGONAL = "diagonal" - A diagonal matrix type with the diagonal implemented as a `Vec`. Useful for
1101: cases where `VecPointwiseMult()` or `VecPointwiseDivide()` should be thought of as the actions of a linear operator.
1103: Options Database Key:
1104: . -mat_vec_type type - set the `VecType` of the vector defining the diagonal
1106: Level: advanced
1108: Note:
1109: The first use case is simply to call `MatCreateDiagonal()' to provide the vector containing the diagonal matrix. The input vector will be
1110: referenced internally by the matrix: any changes to it will affect the matrix. Similar changes to the matrix will affect the vector.
1112: For the second use case call `MatSetType()` with a type of `MATDIAGONAL` followed by a call to `MatDiagonalSetDiagonal()`. The input vector will be
1113: referenced internally by the matrix: any changes to it will affect the matrix. Similar changes to the matrix will affect the vector.
1115: For the third use case call `MatSetType()` with a type of `MATDIAGONAL` followed by calls to `MatSetSizes()` and `MatDiagonalSet()`
1116: (in this case the diagonal vector will not be referenced internally by the matrix, its values will be copied) or some other
1117: operation to provide the matrix entries. One can control the `VecType` of the diagonal by calling `MatSetVecType()` or using `-mat_vec_type type`.
1119: .seealso: [](ch_matrices), `Mat`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`,
1120: `MatDiagonalGetInverseDiagonal()`, `MatDiagonalSet()`, `MatDiagonalSetDiagonal()`, `MatSetVecType()`
1121: M*/
1122: PETSC_INTERN PetscErrorCode MatCreate_Diagonal(Mat A)
1123: {
1124: Mat_Diagonal *ctx;
1126: PetscFunctionBegin;
1127: PetscCall(PetscNew(&ctx));
1128: A->data = (void *)ctx;
1130: A->structurally_symmetric = PETSC_BOOL3_TRUE;
1131: A->structural_symmetry_eternal = PETSC_TRUE;
1132: A->symmetry_eternal = PETSC_TRUE;
1133: A->symmetric = PETSC_BOOL3_TRUE;
1134: if (!PetscDefined(USE_COMPLEX)) A->hermitian = PETSC_BOOL3_TRUE;
1136: A->ops->getrow = MatGetRow_Diagonal;
1137: A->ops->mult = MatMult_Diagonal;
1138: A->ops->multadd = MatMultAdd_Diagonal;
1139: A->ops->multtranspose = MatMult_Diagonal;
1140: A->ops->multtransposeadd = MatMultAdd_Diagonal;
1141: A->ops->norm = MatNorm_Diagonal;
1142: A->ops->duplicate = MatDuplicate_Diagonal;
1143: A->ops->solve = MatSolve_Diagonal;
1144: A->ops->solvetranspose = MatSolve_Diagonal;
1145: A->ops->shift = MatShift_Diagonal;
1146: A->ops->scale = MatScale_Diagonal;
1147: A->ops->diagonalscale = MatDiagonalScale_Diagonal;
1148: A->ops->getdiagonal = MatGetDiagonal_Diagonal;
1149: A->ops->diagonalset = MatDiagonalSet_Diagonal;
1150: A->ops->view = MatView_Diagonal;
1151: A->ops->zeroentries = MatZeroEntries_Diagonal;
1152: A->ops->destroy = MatDestroy_Diagonal;
1153: A->ops->getinfo = MatGetInfo_Diagonal;
1154: A->ops->axpy = MatAXPY_Diagonal;
1155: A->ops->setup = MatSetUp_Diagonal;
1156: A->ops->permute = MatPermute_Diagonal;
1157: A->ops->setrandom = MatSetRandom_Diagonal;
1158: A->ops->conjugate = MatConjugate_Diagonal;
1159: A->ops->transpose = MatTranspose_Diagonal;
1161: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetDiagonal_C", MatDiagonalGetDiagonal_Diagonal));
1162: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreDiagonal_C", MatDiagonalRestoreDiagonal_Diagonal));
1163: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetInverseDiagonal_C", MatDiagonalGetInverseDiagonal_Diagonal));
1164: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreInverseDiagonal_C", MatDiagonalRestoreInverseDiagonal_Diagonal));
1165: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_seqdense_C", MatProductSetFromOptions_Diagonal_Dense));
1166: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_mpidense_C", MatProductSetFromOptions_Diagonal_Dense));
1167: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_diagonal_C", MatProductSetFromOptions_Diagonal_Diagonal));
1168: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_anytype_C", MatProductSetFromOptions_Diagonal_Anytype));
1169: PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATDIAGONAL));
1170: PetscFunctionReturn(PETSC_SUCCESS);
1171: }