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