Actual source code: mcomposite.c
1: #include <../src/mat/impls/shell/shell.h>
3: const char *const MatCompositeMergeTypes[] = {"left", "right", "MatCompositeMergeType", "MAT_COMPOSITE_", NULL};
5: typedef struct _Mat_CompositeLink *Mat_CompositeLink;
6: struct _Mat_CompositeLink {
7: Mat mat;
8: Vec work;
9: Mat_CompositeLink next, prev;
10: };
12: typedef struct {
13: MatCompositeType type;
14: Mat_CompositeLink head, tail;
15: Vec work;
16: PetscInt nmat;
17: PetscBool merge;
18: MatCompositeMergeType mergetype;
19: MatStructure structure;
21: PetscScalar *scalings;
22: PetscBool merge_mvctx; /* Whether need to merge mvctx of component matrices */
23: Vec *lvecs; /* [nmat] Basically, they are Mvctx->lvec of each component matrix */
24: PetscScalar *larray; /* [len] Data arrays of lvecs[] are stored consecutively in larray */
25: PetscInt len; /* Length of larray[] */
26: Vec gvec; /* Union of lvecs[] without duplicated entries */
27: PetscInt *location; /* A map that maps entries in garray[] to larray[] */
28: VecScatter Mvctx;
29: } Mat_Composite;
31: static PetscErrorCode MatDestroy_Composite(Mat mat)
32: {
33: Mat_Composite *shell;
34: Mat_CompositeLink next, oldnext;
35: PetscInt i;
37: PetscFunctionBegin;
38: PetscCall(MatShellGetContext(mat, &shell));
39: next = shell->head;
40: while (next) {
41: PetscCall(MatDestroy(&next->mat));
42: if (next->work && (!next->next || next->work != next->next->work)) PetscCall(VecDestroy(&next->work));
43: oldnext = next;
44: next = next->next;
45: PetscCall(PetscFree(oldnext));
46: }
47: PetscCall(VecDestroy(&shell->work));
49: if (shell->Mvctx) {
50: for (i = 0; i < shell->nmat; i++) PetscCall(VecDestroy(&shell->lvecs[i]));
51: PetscCall(PetscFree3(shell->location, shell->larray, shell->lvecs));
52: PetscCall(VecDestroy(&shell->gvec));
53: PetscCall(VecScatterDestroy(&shell->Mvctx));
54: }
56: PetscCall(PetscFree(shell->scalings));
57: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeAddMat_C", NULL));
58: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeSetType_C", NULL));
59: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeGetType_C", NULL));
60: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeSetMergeType_C", NULL));
61: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeSetMatStructure_C", NULL));
62: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeGetMatStructure_C", NULL));
63: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeMerge_C", NULL));
64: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeGetNumberMat_C", NULL));
65: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeGetMat_C", NULL));
66: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatCompositeSetScalings_C", NULL));
67: PetscCall(PetscFree(shell));
68: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatShellSetContext_C", NULL)); // needed to avoid a call to MatShellSetContext_Immutable()
69: PetscFunctionReturn(PETSC_SUCCESS);
70: }
72: static PetscErrorCode MatMult_Composite_Multiplicative(Mat A, Vec x, Vec y)
73: {
74: Mat_Composite *shell;
75: Mat_CompositeLink next;
76: Vec out;
78: PetscFunctionBegin;
79: PetscCall(MatShellGetContext(A, &shell));
80: next = shell->head;
81: PetscCheck(next, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Must provide at least one matrix with MatCompositeAddMat()");
82: while (next->next) {
83: if (!next->work) { /* should reuse previous work if the same size */
84: PetscCall(MatCreateVecs(next->mat, NULL, &next->work));
85: }
86: out = next->work;
87: PetscCall(MatMult(next->mat, x, out));
88: x = out;
89: next = next->next;
90: }
91: PetscCall(MatMult(next->mat, x, y));
92: if (shell->scalings) {
93: PetscScalar scale = 1.0;
94: for (PetscInt i = 0; i < shell->nmat; i++) scale *= shell->scalings[i];
95: PetscCall(VecScale(y, scale));
96: }
97: PetscFunctionReturn(PETSC_SUCCESS);
98: }
100: static PetscErrorCode MatMultTranspose_Composite_Multiplicative(Mat A, Vec x, Vec y)
101: {
102: Mat_Composite *shell;
103: Mat_CompositeLink tail;
104: Vec out;
106: PetscFunctionBegin;
107: PetscCall(MatShellGetContext(A, &shell));
108: tail = shell->tail;
109: PetscCheck(tail, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Must provide at least one matrix with MatCompositeAddMat()");
110: while (tail->prev) {
111: if (!tail->prev->work) { /* should reuse previous work if the same size */
112: PetscCall(MatCreateVecs(tail->mat, NULL, &tail->prev->work));
113: }
114: out = tail->prev->work;
115: PetscCall(MatMultTranspose(tail->mat, x, out));
116: x = out;
117: tail = tail->prev;
118: }
119: PetscCall(MatMultTranspose(tail->mat, x, y));
120: if (shell->scalings) {
121: PetscScalar scale = 1.0;
122: for (PetscInt i = 0; i < shell->nmat; i++) scale *= shell->scalings[i];
123: PetscCall(VecScale(y, scale));
124: }
125: PetscFunctionReturn(PETSC_SUCCESS);
126: }
128: static PetscErrorCode MatMult_Composite(Mat mat, Vec x, Vec y)
129: {
130: Mat_Composite *shell;
131: Mat_CompositeLink cur;
132: Vec y2, xin;
133: Mat A, B;
134: PetscInt i, j, k, n, nuniq, lo, hi, mid, *gindices, *buf, *tmp, tot;
135: const PetscScalar *vals;
136: const PetscInt *garray;
137: IS ix, iy;
138: PetscBool match;
140: PetscFunctionBegin;
141: PetscCall(MatShellGetContext(mat, &shell));
142: cur = shell->head;
143: PetscCheck(cur, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Must provide at least one matrix with MatCompositeAddMat()");
145: /* Try to merge Mvctx when instructed but not yet done. We did not do it in MatAssemblyEnd() since at that time
146: we did not know whether mat is ADDITIVE or MULTIPLICATIVE. Only now we are assured mat is ADDITIVE and
147: it is legal to merge Mvctx, because all component matrices have the same size.
148: */
149: if (shell->merge_mvctx && !shell->Mvctx) {
150: /* Currently only implemented for MATMPIAIJ */
151: for (cur = shell->head; cur; cur = cur->next) {
152: PetscCall(PetscObjectTypeCompare((PetscObject)cur->mat, MATMPIAIJ, &match));
153: if (!match) {
154: shell->merge_mvctx = PETSC_FALSE;
155: goto skip_merge_mvctx;
156: }
157: }
159: /* Go through matrices first time to count total number of nonzero off-diag columns (may have dups) */
160: tot = 0;
161: for (cur = shell->head; cur; cur = cur->next) {
162: PetscCall(MatMPIAIJGetSeqAIJ(cur->mat, NULL, &B, NULL));
163: PetscCall(MatGetLocalSize(B, NULL, &n));
164: tot += n;
165: }
166: PetscCall(PetscMalloc3(tot, &shell->location, tot, &shell->larray, shell->nmat, &shell->lvecs));
167: shell->len = tot;
169: /* Go through matrices second time to sort off-diag columns and remove dups */
170: PetscCall(PetscMalloc1(tot, &gindices)); /* No Malloc2() since we will give one to PETSc and free the other */
171: PetscCall(PetscMalloc1(tot, &buf));
172: nuniq = 0; /* Number of unique nonzero columns */
173: for (cur = shell->head; cur; cur = cur->next) {
174: PetscCall(MatMPIAIJGetSeqAIJ(cur->mat, NULL, &B, &garray));
175: PetscCall(MatGetLocalSize(B, NULL, &n));
176: /* Merge pre-sorted garray[0,n) and gindices[0,nuniq) to buf[] */
177: i = j = k = 0;
178: while (i < n && j < nuniq) {
179: if (garray[i] < gindices[j]) buf[k++] = garray[i++];
180: else if (garray[i] > gindices[j]) buf[k++] = gindices[j++];
181: else {
182: buf[k++] = garray[i++];
183: j++;
184: }
185: }
186: /* Copy leftover in garray[] or gindices[] */
187: if (i < n) {
188: PetscCall(PetscArraycpy(buf + k, garray + i, n - i));
189: nuniq = k + n - i;
190: } else if (j < nuniq) {
191: PetscCall(PetscArraycpy(buf + k, gindices + j, nuniq - j));
192: nuniq = k + nuniq - j;
193: } else nuniq = k;
194: /* Swap gindices and buf to merge garray of the next matrix */
195: tmp = gindices;
196: gindices = buf;
197: buf = tmp;
198: }
199: PetscCall(PetscFree(buf));
201: /* Go through matrices third time to build a map from gindices[] to garray[] */
202: tot = 0;
203: for (cur = shell->head, j = 0; cur; cur = cur->next, j++) { /* j-th matrix */
204: PetscCall(MatMPIAIJGetSeqAIJ(cur->mat, NULL, &B, &garray));
205: PetscCall(MatGetLocalSize(B, NULL, &n));
206: PetscCall(VecCreateSeqWithArray(PETSC_COMM_SELF, 1, n, NULL, &shell->lvecs[j]));
207: /* This is an optimized PetscFindInt(garray[i],nuniq,gindices,&shell->location[tot+i]), using the fact that garray[] is also sorted */
208: lo = 0;
209: for (i = 0; i < n; i++) {
210: hi = nuniq;
211: while (hi - lo > 1) {
212: mid = lo + (hi - lo) / 2;
213: if (garray[i] < gindices[mid]) hi = mid;
214: else lo = mid;
215: }
216: shell->location[tot + i] = lo; /* gindices[lo] = garray[i] */
217: lo++; /* Since garray[i+1] > garray[i], we can safely advance lo */
218: }
219: tot += n;
220: }
222: /* Build merged Mvctx */
223: PetscCall(ISCreateGeneral(PETSC_COMM_SELF, nuniq, gindices, PETSC_OWN_POINTER, &ix));
224: PetscCall(ISCreateStride(PETSC_COMM_SELF, nuniq, 0, 1, &iy));
225: PetscCall(VecCreateMPIWithArray(PetscObjectComm((PetscObject)mat), 1, mat->cmap->n, mat->cmap->N, NULL, &xin));
226: PetscCall(VecCreateSeq(PETSC_COMM_SELF, nuniq, &shell->gvec));
227: PetscCall(VecScatterCreate(xin, ix, shell->gvec, iy, &shell->Mvctx));
228: PetscCall(VecDestroy(&xin));
229: PetscCall(ISDestroy(&ix));
230: PetscCall(ISDestroy(&iy));
231: }
233: skip_merge_mvctx:
234: PetscCall(VecSet(y, 0));
235: if (!((Mat_Shell *)mat->data)->left_work) PetscCall(VecDuplicate(y, &(((Mat_Shell *)mat->data)->left_work)));
236: y2 = ((Mat_Shell *)mat->data)->left_work;
238: if (shell->Mvctx) { /* Have a merged Mvctx */
239: /* Suppose we want to compute y = sMx, where s is the scaling factor and A, B are matrix M's diagonal/off-diagonal part. We could do
240: in y = s(Ax1 + Bx2) or y = sAx1 + sBx2. The former incurs less FLOPS than the latter, but the latter provides an opportunity to
241: overlap communication/computation since we can do sAx1 while communicating x2. Here, we use the former approach.
242: */
243: PetscCall(VecScatterBegin(shell->Mvctx, x, shell->gvec, INSERT_VALUES, SCATTER_FORWARD));
244: PetscCall(VecScatterEnd(shell->Mvctx, x, shell->gvec, INSERT_VALUES, SCATTER_FORWARD));
246: PetscCall(VecGetArrayRead(shell->gvec, &vals));
247: for (i = 0; i < shell->len; i++) shell->larray[i] = vals[shell->location[i]];
248: PetscCall(VecRestoreArrayRead(shell->gvec, &vals));
250: for (cur = shell->head, tot = i = 0; cur; cur = cur->next, i++) { /* i-th matrix */
251: PetscCall(MatMPIAIJGetSeqAIJ(cur->mat, &A, &B, NULL));
252: PetscUseTypeMethod(A, mult, x, y2);
253: PetscCall(MatGetLocalSize(B, NULL, &n));
254: PetscCall(VecPlaceArray(shell->lvecs[i], &shell->larray[tot]));
255: PetscUseTypeMethod(B, multadd, shell->lvecs[i], y2, y2);
256: PetscCall(VecResetArray(shell->lvecs[i]));
257: PetscCall(VecAXPY(y, shell->scalings ? shell->scalings[i] : 1.0, y2));
258: tot += n;
259: }
260: } else {
261: if (shell->scalings) {
262: for (cur = shell->head, i = 0; cur; cur = cur->next, i++) {
263: PetscCall(MatMult(cur->mat, x, y2));
264: PetscCall(VecAXPY(y, shell->scalings[i], y2));
265: }
266: } else {
267: for (cur = shell->head; cur; cur = cur->next) PetscCall(MatMultAdd(cur->mat, x, y, y));
268: }
269: }
270: PetscFunctionReturn(PETSC_SUCCESS);
271: }
273: static PetscErrorCode MatMultTranspose_Composite(Mat A, Vec x, Vec y)
274: {
275: Mat_Composite *shell;
276: Mat_CompositeLink next;
277: Vec y2 = NULL;
278: PetscInt i;
280: PetscFunctionBegin;
281: PetscCall(MatShellGetContext(A, &shell));
282: next = shell->head;
283: PetscCheck(next, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Must provide at least one matrix with MatCompositeAddMat()");
285: PetscCall(MatMultTranspose(next->mat, x, y));
286: if (shell->scalings) {
287: PetscCall(VecScale(y, shell->scalings[0]));
288: if (!((Mat_Shell *)A->data)->right_work) PetscCall(VecDuplicate(y, &(((Mat_Shell *)A->data)->right_work)));
289: y2 = ((Mat_Shell *)A->data)->right_work;
290: }
291: i = 1;
292: while ((next = next->next)) {
293: if (!shell->scalings) PetscCall(MatMultTransposeAdd(next->mat, x, y, y));
294: else {
295: PetscCall(MatMultTranspose(next->mat, x, y2));
296: PetscCall(VecAXPY(y, shell->scalings[i++], y2));
297: }
298: }
299: PetscFunctionReturn(PETSC_SUCCESS);
300: }
302: static PetscErrorCode MatGetDiagonal_Composite(Mat A, Vec v)
303: {
304: Mat_Composite *shell;
305: Mat_CompositeLink next;
306: PetscInt i;
308: PetscFunctionBegin;
309: PetscCall(MatShellGetContext(A, &shell));
310: next = shell->head;
311: PetscCheck(next, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Must provide at least one matrix with MatCompositeAddMat()");
312: PetscCall(MatGetDiagonal(next->mat, v));
313: if (shell->scalings) PetscCall(VecScale(v, shell->scalings[0]));
315: if (next->next && !shell->work) PetscCall(VecDuplicate(v, &shell->work));
316: i = 1;
317: while ((next = next->next)) {
318: PetscCall(MatGetDiagonal(next->mat, shell->work));
319: PetscCall(VecAXPY(v, shell->scalings ? shell->scalings[i++] : 1.0, shell->work));
320: }
321: PetscFunctionReturn(PETSC_SUCCESS);
322: }
324: static PetscErrorCode MatAssemblyEnd_Composite(Mat Y, MatAssemblyType t)
325: {
326: Mat_Composite *shell;
328: PetscFunctionBegin;
329: PetscCall(MatShellGetContext(Y, &shell));
330: if (shell->merge) PetscCall(MatCompositeMerge(Y));
331: else PetscCall(MatAssemblyEnd_Shell(Y, t));
332: PetscFunctionReturn(PETSC_SUCCESS);
333: }
335: static PetscErrorCode MatSetFromOptions_Composite(Mat A, PetscOptionItems PetscOptionsObject)
336: {
337: Mat_Composite *a;
339: PetscFunctionBegin;
340: PetscCall(MatShellGetContext(A, &a));
341: PetscOptionsHeadBegin(PetscOptionsObject, "MATCOMPOSITE options");
342: PetscCall(PetscOptionsBool("-mat_composite_merge", "Merge at MatAssemblyEnd", "MatCompositeMerge", a->merge, &a->merge, NULL));
343: PetscCall(PetscOptionsEnum("-mat_composite_merge_type", "Set composite merge direction", "MatCompositeSetMergeType", MatCompositeMergeTypes, (PetscEnum)a->mergetype, (PetscEnum *)&a->mergetype, NULL));
344: PetscCall(PetscOptionsBool("-mat_composite_merge_mvctx", "Merge MatMult() vecscat contexts", "MatCreateComposite", a->merge_mvctx, &a->merge_mvctx, NULL));
345: PetscOptionsHeadEnd();
346: PetscFunctionReturn(PETSC_SUCCESS);
347: }
349: /*@
350: MatCreateComposite - Creates a matrix as the sum or product of one or more matrices
352: Collective
354: Input Parameters:
355: + comm - MPI communicator
356: . nmat - number of matrices to put in
357: - mats - the matrices
359: Output Parameter:
360: . mat - the matrix
362: Options Database Keys:
363: + -mat_composite_merge - merge in `MatAssemblyEnd()`
364: . -mat_composite_merge_mvctx - merge Mvctx of component matrices to optimize communication in `MatMult()` for ADDITIVE matrices
365: - -mat_composite_merge_type - set merge direction
367: Level: advanced
369: Note:
370: Alternative construction
371: .vb
372: MatCreate(comm,&mat);
373: MatSetSizes(mat,m,n,M,N);
374: MatSetType(mat,MATCOMPOSITE);
375: MatCompositeAddMat(mat,mats[0]);
376: ....
377: MatCompositeAddMat(mat,mats[nmat-1]);
378: MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
379: MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);
380: .ve
382: For the multiplicative form the product is mat[nmat-1]*mat[nmat-2]*....*mat[0]
384: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MatMult()`, `MatCompositeAddMat()`, `MatCompositeGetMat()`, `MatCompositeMerge()`, `MatCompositeSetType()`,
385: `MATCOMPOSITE`, `MatCompositeType`
386: @*/
387: PetscErrorCode MatCreateComposite(MPI_Comm comm, PetscInt nmat, const Mat *mats, Mat *mat)
388: {
389: PetscFunctionBegin;
390: PetscCheck(nmat >= 1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Must pass in at least one matrix");
391: PetscAssertPointer(mat, 4);
392: PetscCall(MatCreate(comm, mat));
393: PetscCall(MatSetType(*mat, MATCOMPOSITE));
394: for (PetscInt i = 0; i < nmat; i++) PetscCall(MatCompositeAddMat(*mat, mats[i]));
395: PetscCall(MatAssemblyBegin(*mat, MAT_FINAL_ASSEMBLY));
396: PetscCall(MatAssemblyEnd(*mat, MAT_FINAL_ASSEMBLY));
397: PetscFunctionReturn(PETSC_SUCCESS);
398: }
400: static PetscErrorCode MatCompositeAddMat_Composite(Mat mat, Mat smat)
401: {
402: Mat_Composite *shell;
403: Mat_CompositeLink ilink, next;
404: VecType vtype_mat, vtype_smat;
405: PetscBool match;
407: PetscFunctionBegin;
408: PetscCall(MatShellGetContext(mat, &shell));
409: next = shell->head;
410: PetscCall(PetscNew(&ilink));
411: ilink->next = NULL;
412: PetscCall(PetscObjectReference((PetscObject)smat));
413: ilink->mat = smat;
415: if (!next) shell->head = ilink;
416: else {
417: while (next->next) next = next->next;
418: next->next = ilink;
419: ilink->prev = next;
420: }
421: shell->tail = ilink;
422: shell->nmat += 1;
424: /* If all of the partial matrices have the same default vector type, then the composite matrix should also have this default type.
425: Otherwise, the default type should be "standard". */
426: PetscCall(MatGetVecType(smat, &vtype_smat));
427: if (shell->nmat == 1) PetscCall(MatSetVecType(mat, vtype_smat));
428: else {
429: PetscCall(MatGetVecType(mat, &vtype_mat));
430: PetscCall(PetscStrcmp(vtype_smat, vtype_mat, &match));
431: if (!match) PetscCall(MatSetVecType(mat, VECSTANDARD));
432: }
434: /* Retain the old scalings (if any) and expand it with a 1.0 for the newly added matrix */
435: if (shell->scalings) {
436: PetscCall(PetscRealloc(sizeof(PetscScalar) * shell->nmat, &shell->scalings));
437: shell->scalings[shell->nmat - 1] = 1.0;
438: }
440: /* The composite matrix requires PetscLayouts for its rows and columns; we copy these from the constituent partial matrices. */
441: if (shell->nmat == 1) PetscCall(PetscLayoutReference(smat->cmap, &mat->cmap));
442: PetscCall(PetscLayoutReference(smat->rmap, &mat->rmap));
443: PetscFunctionReturn(PETSC_SUCCESS);
444: }
446: /*@
447: MatCompositeAddMat - Add another matrix to a composite matrix.
449: Collective
451: Input Parameters:
452: + mat - the composite matrix
453: - smat - the partial matrix
455: Level: advanced
457: .seealso: [](ch_matrices), `Mat`, `MatCreateComposite()`, `MatCompositeGetMat()`, `MATCOMPOSITE`
458: @*/
459: PetscErrorCode MatCompositeAddMat(Mat mat, Mat smat)
460: {
461: PetscFunctionBegin;
464: PetscUseMethod(mat, "MatCompositeAddMat_C", (Mat, Mat), (mat, smat));
465: PetscFunctionReturn(PETSC_SUCCESS);
466: }
468: static PetscErrorCode MatCompositeSetType_Composite(Mat mat, MatCompositeType type)
469: {
470: Mat_Composite *b;
472: PetscFunctionBegin;
473: PetscCall(MatShellGetContext(mat, &b));
474: b->type = type;
475: if (type == MAT_COMPOSITE_MULTIPLICATIVE) {
476: PetscCall(MatShellSetOperation(mat, MATOP_GET_DIAGONAL, NULL));
477: PetscCall(MatShellSetOperation(mat, MATOP_MULT, (PetscErrorCodeFn *)MatMult_Composite_Multiplicative));
478: PetscCall(MatShellSetOperation(mat, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_Composite_Multiplicative));
479: b->merge_mvctx = PETSC_FALSE;
480: } else {
481: PetscCall(MatShellSetOperation(mat, MATOP_GET_DIAGONAL, (PetscErrorCodeFn *)MatGetDiagonal_Composite));
482: PetscCall(MatShellSetOperation(mat, MATOP_MULT, (PetscErrorCodeFn *)MatMult_Composite));
483: PetscCall(MatShellSetOperation(mat, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_Composite));
484: }
485: PetscFunctionReturn(PETSC_SUCCESS);
486: }
488: /*@
489: MatCompositeSetType - Indicates if the matrix is defined as the sum of a set of matrices or the product.
491: Logically Collective
493: Input Parameters:
494: + mat - the composite matrix
495: - type - the `MatCompositeType` to use for the matrix
497: Level: advanced
499: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MatMult()`, `MatCompositeAddMat()`, `MatCreateComposite()`, `MatCompositeGetType()`, `MATCOMPOSITE`,
500: `MatCompositeType`
501: @*/
502: PetscErrorCode MatCompositeSetType(Mat mat, MatCompositeType type)
503: {
504: PetscFunctionBegin;
507: PetscUseMethod(mat, "MatCompositeSetType_C", (Mat, MatCompositeType), (mat, type));
508: PetscFunctionReturn(PETSC_SUCCESS);
509: }
511: static PetscErrorCode MatCompositeGetType_Composite(Mat mat, MatCompositeType *type)
512: {
513: Mat_Composite *shell;
515: PetscFunctionBegin;
516: PetscCall(MatShellGetContext(mat, &shell));
517: *type = shell->type;
518: PetscFunctionReturn(PETSC_SUCCESS);
519: }
521: /*@
522: MatCompositeGetType - Returns type of composite.
524: Not Collective
526: Input Parameter:
527: . mat - the composite matrix
529: Output Parameter:
530: . type - type of composite
532: Level: advanced
534: .seealso: [](ch_matrices), `Mat`, `MatCreateComposite()`, `MatCompositeSetType()`, `MATCOMPOSITE`, `MatCompositeType`
535: @*/
536: PetscErrorCode MatCompositeGetType(Mat mat, MatCompositeType *type)
537: {
538: PetscFunctionBegin;
540: PetscAssertPointer(type, 2);
541: PetscUseMethod(mat, "MatCompositeGetType_C", (Mat, MatCompositeType *), (mat, type));
542: PetscFunctionReturn(PETSC_SUCCESS);
543: }
545: static PetscErrorCode MatCompositeSetMatStructure_Composite(Mat mat, MatStructure str)
546: {
547: Mat_Composite *shell;
549: PetscFunctionBegin;
550: PetscCall(MatShellGetContext(mat, &shell));
551: shell->structure = str;
552: PetscFunctionReturn(PETSC_SUCCESS);
553: }
555: /*@
556: MatCompositeSetMatStructure - Indicates structure of matrices in the composite matrix.
558: Not Collective
560: Input Parameters:
561: + mat - the composite matrix
562: - str - either `SAME_NONZERO_PATTERN`, `DIFFERENT_NONZERO_PATTERN` (default) or `SUBSET_NONZERO_PATTERN`
564: Level: advanced
566: Note:
567: Information about the matrices structure is used in `MatCompositeMerge()` for additive composite matrix.
569: .seealso: [](ch_matrices), `Mat`, `MatAXPY()`, `MatCreateComposite()`, `MatCompositeMerge()`, `MatCompositeGetMatStructure()`, `MATCOMPOSITE`
570: @*/
571: PetscErrorCode MatCompositeSetMatStructure(Mat mat, MatStructure str)
572: {
573: PetscFunctionBegin;
575: PetscUseMethod(mat, "MatCompositeSetMatStructure_C", (Mat, MatStructure), (mat, str));
576: PetscFunctionReturn(PETSC_SUCCESS);
577: }
579: static PetscErrorCode MatCompositeGetMatStructure_Composite(Mat mat, MatStructure *str)
580: {
581: Mat_Composite *shell;
583: PetscFunctionBegin;
584: PetscCall(MatShellGetContext(mat, &shell));
585: *str = shell->structure;
586: PetscFunctionReturn(PETSC_SUCCESS);
587: }
589: /*@
590: MatCompositeGetMatStructure - Returns the structure of matrices in the composite matrix.
592: Not Collective
594: Input Parameter:
595: . mat - the composite matrix
597: Output Parameter:
598: . str - structure of the matrices
600: Level: advanced
602: .seealso: [](ch_matrices), `Mat`, `MatCreateComposite()`, `MatCompositeSetMatStructure()`, `MATCOMPOSITE`
603: @*/
604: PetscErrorCode MatCompositeGetMatStructure(Mat mat, MatStructure *str)
605: {
606: PetscFunctionBegin;
608: PetscAssertPointer(str, 2);
609: PetscUseMethod(mat, "MatCompositeGetMatStructure_C", (Mat, MatStructure *), (mat, str));
610: PetscFunctionReturn(PETSC_SUCCESS);
611: }
613: static PetscErrorCode MatCompositeSetMergeType_Composite(Mat mat, MatCompositeMergeType type)
614: {
615: Mat_Composite *shell;
617: PetscFunctionBegin;
618: PetscCall(MatShellGetContext(mat, &shell));
619: shell->mergetype = type;
620: PetscFunctionReturn(PETSC_SUCCESS);
621: }
623: /*@
624: MatCompositeSetMergeType - Sets order of `MatCompositeMerge()`.
626: Logically Collective
628: Input Parameters:
629: + mat - the composite matrix
630: - type - `MAT_COMPOSITE_MERGE RIGHT` (default) to start merge from right with the first added matrix (mat[0]),
631: `MAT_COMPOSITE_MERGE_LEFT` to start merge from left with the last added matrix (mat[nmat-1])
633: Level: advanced
635: Note:
636: The resulting matrix is the same regardless of the `MatCompositeMergeType`. Only the order of operation is changed.
637: If set to `MAT_COMPOSITE_MERGE_RIGHT` the order of the merge is mat[nmat-1]*(mat[nmat-2]*(...*(mat[1]*mat[0])))
638: otherwise the order is (((mat[nmat-1]*mat[nmat-2])*mat[nmat-3])*...)*mat[0].
640: .seealso: [](ch_matrices), `Mat`, `MatCreateComposite()`, `MatCompositeMerge()`, `MATCOMPOSITE`
641: @*/
642: PetscErrorCode MatCompositeSetMergeType(Mat mat, MatCompositeMergeType type)
643: {
644: PetscFunctionBegin;
647: PetscUseMethod(mat, "MatCompositeSetMergeType_C", (Mat, MatCompositeMergeType), (mat, type));
648: PetscFunctionReturn(PETSC_SUCCESS);
649: }
651: static PetscErrorCode MatCompositeMerge_Composite(Mat mat)
652: {
653: Mat_Composite *shell;
654: Mat_CompositeLink next, prev;
655: Mat tmat, newmat;
656: Vec left, right, dshift;
657: PetscScalar scale, shift;
658: PetscInt i;
660: PetscFunctionBegin;
661: PetscCall(MatShellGetContext(mat, &shell));
662: next = shell->head;
663: prev = shell->tail;
664: PetscCheck(next, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Must provide at least one matrix with MatCompositeAddMat()");
665: PetscCall(MatShellGetScalingShifts(mat, &shift, &scale, &dshift, &left, &right, (Mat *)MAT_SHELL_NOT_ALLOWED, (IS *)MAT_SHELL_NOT_ALLOWED, (IS *)MAT_SHELL_NOT_ALLOWED));
666: if (shell->type == MAT_COMPOSITE_ADDITIVE) {
667: if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) {
668: i = 0;
669: PetscCall(MatDuplicate(next->mat, MAT_COPY_VALUES, &tmat));
670: if (shell->scalings) PetscCall(MatScale(tmat, shell->scalings[i++]));
671: while ((next = next->next)) PetscCall(MatAXPY(tmat, shell->scalings ? shell->scalings[i++] : 1.0, next->mat, shell->structure));
672: } else {
673: i = shell->nmat - 1;
674: PetscCall(MatDuplicate(prev->mat, MAT_COPY_VALUES, &tmat));
675: if (shell->scalings) PetscCall(MatScale(tmat, shell->scalings[i--]));
676: while ((prev = prev->prev)) PetscCall(MatAXPY(tmat, shell->scalings ? shell->scalings[i--] : 1.0, prev->mat, shell->structure));
677: }
678: } else {
679: if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) {
680: PetscCall(MatDuplicate(next->mat, MAT_COPY_VALUES, &tmat));
681: while ((next = next->next)) {
682: PetscCall(MatMatMult(next->mat, tmat, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &newmat));
683: PetscCall(MatDestroy(&tmat));
684: tmat = newmat;
685: }
686: } else {
687: PetscCall(MatDuplicate(prev->mat, MAT_COPY_VALUES, &tmat));
688: while ((prev = prev->prev)) {
689: PetscCall(MatMatMult(tmat, prev->mat, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &newmat));
690: PetscCall(MatDestroy(&tmat));
691: tmat = newmat;
692: }
693: }
694: if (shell->scalings) {
695: for (i = 0; i < shell->nmat; i++) scale *= shell->scalings[i];
696: }
697: }
699: PetscCall(PetscObjectReference((PetscObject)left));
700: PetscCall(PetscObjectReference((PetscObject)right));
701: PetscCall(PetscObjectReference((PetscObject)dshift));
703: PetscCall(MatHeaderReplace(mat, &tmat));
705: PetscCall(MatDiagonalScale(mat, left, right));
706: PetscCall(MatScale(mat, scale));
707: PetscCall(MatShift(mat, shift));
708: PetscCall(VecDestroy(&left));
709: PetscCall(VecDestroy(&right));
710: if (dshift) {
711: PetscCall(MatDiagonalSet(mat, dshift, ADD_VALUES));
712: PetscCall(VecDestroy(&dshift));
713: }
714: PetscFunctionReturn(PETSC_SUCCESS);
715: }
717: /*@
718: MatCompositeMerge - Given a composite matrix, replaces it with a "regular" matrix
719: by summing or computing the product of all the matrices inside the composite matrix.
721: Collective
723: Input Parameter:
724: . mat - the composite matrix
726: Options Database Keys:
727: + -mat_composite_merge - merge in `MatAssemblyEnd()`
728: - -mat_composite_merge_type - set merge direction
730: Level: advanced
732: Note:
733: The `MatType` of the resulting matrix will be the same as the `MatType` of the FIRST matrix in the composite matrix.
735: .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MatMult()`, `MatCompositeAddMat()`, `MatCreateComposite()`, `MatCompositeSetMatStructure()`, `MatCompositeSetMergeType()`, `MATCOMPOSITE`
736: @*/
737: PetscErrorCode MatCompositeMerge(Mat mat)
738: {
739: PetscFunctionBegin;
741: PetscUseMethod(mat, "MatCompositeMerge_C", (Mat), (mat));
742: PetscFunctionReturn(PETSC_SUCCESS);
743: }
745: static PetscErrorCode MatCompositeGetNumberMat_Composite(Mat mat, PetscInt *nmat)
746: {
747: Mat_Composite *shell;
749: PetscFunctionBegin;
750: PetscCall(MatShellGetContext(mat, &shell));
751: *nmat = shell->nmat;
752: PetscFunctionReturn(PETSC_SUCCESS);
753: }
755: /*@
756: MatCompositeGetNumberMat - Returns the number of matrices in the composite matrix.
758: Not Collective
760: Input Parameter:
761: . mat - the composite matrix
763: Output Parameter:
764: . nmat - number of matrices in the composite matrix
766: Level: advanced
768: .seealso: [](ch_matrices), `Mat`, `MatCreateComposite()`, `MatCompositeGetMat()`, `MATCOMPOSITE`
769: @*/
770: PetscErrorCode MatCompositeGetNumberMat(Mat mat, PetscInt *nmat)
771: {
772: PetscFunctionBegin;
774: PetscAssertPointer(nmat, 2);
775: PetscUseMethod(mat, "MatCompositeGetNumberMat_C", (Mat, PetscInt *), (mat, nmat));
776: PetscFunctionReturn(PETSC_SUCCESS);
777: }
779: static PetscErrorCode MatCompositeGetMat_Composite(Mat mat, PetscInt i, Mat *Ai)
780: {
781: Mat_Composite *shell;
782: Mat_CompositeLink ilink;
783: PetscInt k;
785: PetscFunctionBegin;
786: PetscCall(MatShellGetContext(mat, &shell));
787: PetscCheck(i < shell->nmat, PetscObjectComm((PetscObject)mat), PETSC_ERR_ARG_OUTOFRANGE, "index out of range: %" PetscInt_FMT " >= %" PetscInt_FMT, i, shell->nmat);
788: ilink = shell->head;
789: for (k = 0; k < i; k++) ilink = ilink->next;
790: *Ai = ilink->mat;
791: PetscFunctionReturn(PETSC_SUCCESS);
792: }
794: /*@
795: MatCompositeGetMat - Returns the ith matrix from the composite matrix.
797: Logically Collective
799: Input Parameters:
800: + mat - the composite matrix
801: - i - the number of requested matrix
803: Output Parameter:
804: . Ai - ith matrix in composite
806: Level: advanced
808: .seealso: [](ch_matrices), `Mat`, `MatCreateComposite()`, `MatCompositeGetNumberMat()`, `MatCompositeAddMat()`, `MATCOMPOSITE`
809: @*/
810: PetscErrorCode MatCompositeGetMat(Mat mat, PetscInt i, Mat *Ai)
811: {
812: PetscFunctionBegin;
815: PetscAssertPointer(Ai, 3);
816: PetscUseMethod(mat, "MatCompositeGetMat_C", (Mat, PetscInt, Mat *), (mat, i, Ai));
817: PetscFunctionReturn(PETSC_SUCCESS);
818: }
820: static PetscErrorCode MatCompositeSetScalings_Composite(Mat mat, const PetscScalar *scalings)
821: {
822: Mat_Composite *shell;
823: PetscInt nmat;
825: PetscFunctionBegin;
826: PetscCall(MatShellGetContext(mat, &shell));
827: PetscCall(MatCompositeGetNumberMat(mat, &nmat));
828: if (!shell->scalings) PetscCall(PetscMalloc1(nmat, &shell->scalings));
829: PetscCall(PetscArraycpy(shell->scalings, scalings, nmat));
830: PetscFunctionReturn(PETSC_SUCCESS);
831: }
833: /*@
834: MatCompositeSetScalings - Sets separate scaling factors for component matrices.
836: Logically Collective
838: Input Parameters:
839: + mat - the composite matrix
840: - scalings - array of scaling factors with scalings[i] being factor of i-th matrix, for i in [0, nmat)
842: Level: advanced
844: .seealso: [](ch_matrices), `Mat`, `MatScale()`, `MatDiagonalScale()`, `MATCOMPOSITE`
845: @*/
846: PetscErrorCode MatCompositeSetScalings(Mat mat, const PetscScalar *scalings)
847: {
848: PetscFunctionBegin;
850: PetscAssertPointer(scalings, 2);
852: PetscUseMethod(mat, "MatCompositeSetScalings_C", (Mat, const PetscScalar *), (mat, scalings));
853: PetscFunctionReturn(PETSC_SUCCESS);
854: }
856: /*MC
857: MATCOMPOSITE - A matrix defined by the sum (or product) of one or more matrices.
858: The matrices need to have a correct size and parallel layout for the sum or product to be valid.
860: Level: advanced
862: Note:
863: To use the product of the matrices call `MatCompositeSetType`(mat,`MAT_COMPOSITE_MULTIPLICATIVE`);
865: Developer Notes:
866: This is implemented on top of `MATSHELL` to get support for scaling and shifting without requiring duplicate code
868: Users can not call `MatShellSetOperation()` operations on this class, there is some error checking for that incorrect usage
870: .seealso: [](ch_matrices), `Mat`, `MatCreateComposite()`, `MatCompositeSetScalings()`, `MatCompositeAddMat()`, `MatSetType()`, `MatCompositeSetType()`, `MatCompositeGetType()`,
871: `MatCompositeSetMatStructure()`, `MatCompositeGetMatStructure()`, `MatCompositeMerge()`, `MatCompositeSetMergeType()`, `MatCompositeGetNumberMat()`, `MatCompositeGetMat()`
872: M*/
874: PETSC_EXTERN PetscErrorCode MatCreate_Composite(Mat A)
875: {
876: Mat_Composite *b;
878: PetscFunctionBegin;
879: PetscCall(PetscNew(&b));
881: b->type = MAT_COMPOSITE_ADDITIVE;
882: b->nmat = 0;
883: b->merge = PETSC_FALSE;
884: b->mergetype = MAT_COMPOSITE_MERGE_RIGHT;
885: b->structure = DIFFERENT_NONZERO_PATTERN;
886: b->merge_mvctx = PETSC_TRUE;
888: PetscCall(MatSetType(A, MATSHELL));
889: PetscCall(MatShellSetContext(A, b));
890: PetscCall(MatShellSetOperation(A, MATOP_DESTROY, (PetscErrorCodeFn *)MatDestroy_Composite));
891: PetscCall(MatShellSetOperation(A, MATOP_MULT, (PetscErrorCodeFn *)MatMult_Composite));
892: PetscCall(MatShellSetOperation(A, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_Composite));
893: PetscCall(MatShellSetOperation(A, MATOP_GET_DIAGONAL, (PetscErrorCodeFn *)MatGetDiagonal_Composite));
894: PetscCall(MatShellSetOperation(A, MATOP_ASSEMBLY_END, (PetscErrorCodeFn *)MatAssemblyEnd_Composite));
895: PetscCall(MatShellSetOperation(A, MATOP_SET_FROM_OPTIONS, (PetscErrorCodeFn *)MatSetFromOptions_Composite));
896: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeAddMat_C", MatCompositeAddMat_Composite));
897: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeSetType_C", MatCompositeSetType_Composite));
898: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeGetType_C", MatCompositeGetType_Composite));
899: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeSetMergeType_C", MatCompositeSetMergeType_Composite));
900: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeSetMatStructure_C", MatCompositeSetMatStructure_Composite));
901: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeGetMatStructure_C", MatCompositeGetMatStructure_Composite));
902: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeMerge_C", MatCompositeMerge_Composite));
903: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeGetNumberMat_C", MatCompositeGetNumberMat_Composite));
904: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeGetMat_C", MatCompositeGetMat_Composite));
905: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatCompositeSetScalings_C", MatCompositeSetScalings_Composite));
906: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatShellSetContext_C", MatShellSetContext_Immutable));
907: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatShellSetContextDestroy_C", MatShellSetContextDestroy_Immutable));
908: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatShellSetManageScalingShifts_C", MatShellSetManageScalingShifts_Immutable));
909: PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATCOMPOSITE));
910: PetscFunctionReturn(PETSC_SUCCESS);
911: }