Actual source code: mcomposite.c
petsc-3.13.6 2020-09-29
Report Typos and Errors

      2:  #include <petsc/private/matimpl.h>

      4: const char *const MatCompositeMergeTypes[] = {"left","right","MatCompositeMergeType","MAT_COMPOSITE_",0};

      6: typedef struct _Mat_CompositeLink *Mat_CompositeLink;
      7: struct _Mat_CompositeLink {
      8:   Mat               mat;
      9:   Vec               work;
     10:   Mat_CompositeLink next,prev;
     11: };

     13: typedef struct {
     14:   MatCompositeType      type;
     15:   Mat_CompositeLink     head,tail;
     16:   Vec                   work;
     17:   PetscScalar           scale;        /* scale factor supplied with MatScale() */
     18:   Vec                   left,right;   /* left and right diagonal scaling provided with MatDiagonalScale() */
     19:   Vec                   leftwork,rightwork,leftwork2,rightwork2; /* Two pairs of working vectors */
     20:   PetscInt              nmat;
     21:   PetscBool             merge;
     22:   MatCompositeMergeType mergetype;
     23:   MatStructure          structure;

     25:   PetscScalar           *scalings;
     26:   PetscBool             merge_mvctx;  /* Whether need to merge mvctx of component matrices */
     27:   Vec                   *lvecs;       /* [nmat] Basically, they are Mvctx->lvec of each component matrix */
     28:   PetscScalar           *larray;      /* [len] Data arrays of lvecs[] are stored consecutively in larray */
     29:   PetscInt              len;          /* Length of larray[] */
     30:   Vec                   gvec;         /* Union of lvecs[] without duplicated entries */
     31:   PetscInt              *location;    /* A map that maps entries in garray[] to larray[] */
     32:   VecScatter            Mvctx;
     33: } Mat_Composite;

     35: PetscErrorCode MatDestroy_Composite(Mat mat)
     36: {
     37:   PetscErrorCode    ierr;
     38:   Mat_Composite     *shell = (Mat_Composite*)mat->data;
     39:   Mat_CompositeLink next   = shell->head,oldnext;
     40:   PetscInt          i;

     43:   while (next) {
     44:     MatDestroy(&next->mat);
     45:     if (next->work && (!next->next || next->work != next->next->work)) {
     46:       VecDestroy(&next->work);
     47:     }
     48:     oldnext = next;
     49:     next    = next->next;
     50:     PetscFree(oldnext);
     51:   }
     52:   VecDestroy(&shell->work);
     53:   VecDestroy(&shell->left);
     54:   VecDestroy(&shell->right);
     55:   VecDestroy(&shell->leftwork);
     56:   VecDestroy(&shell->rightwork);
     57:   VecDestroy(&shell->leftwork2);
     58:   VecDestroy(&shell->rightwork2);

     60:   if (shell->Mvctx) {
     61:     for (i=0; i<shell->nmat; i++) {VecDestroy(&shell->lvecs[i]);}
     62:     PetscFree3(shell->location,shell->larray,shell->lvecs);
     63:     PetscFree(shell->larray);
     64:     VecDestroy(&shell->gvec);
     65:     VecScatterDestroy(&shell->Mvctx);
     66:   }

     68:   PetscFree(shell->scalings);
     69:   PetscFree(mat->data);
     70:   return(0);
     71: }

     73: PetscErrorCode MatMult_Composite_Multiplicative(Mat A,Vec x,Vec y)
     74: {
     75:   Mat_Composite     *shell = (Mat_Composite*)A->data;
     76:   Mat_CompositeLink next   = shell->head;
     77:   PetscErrorCode    ierr;
     78:   Vec               in,out;
     79:   PetscScalar       scale;
     80:   PetscInt          i;

     83:   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
     84:   in = x;
     85:   if (shell->right) {
     86:     if (!shell->rightwork) {
     87:       VecDuplicate(shell->right,&shell->rightwork);
     88:     }
     89:     VecPointwiseMult(shell->rightwork,shell->right,in);
     90:     in   = shell->rightwork;
     91:   }
     92:   while (next->next) {
     93:     if (!next->work) { /* should reuse previous work if the same size */
     94:       MatCreateVecs(next->mat,NULL,&next->work);
     95:     }
     96:     out  = next->work;
     97:     MatMult(next->mat,in,out);
     98:     in   = out;
     99:     next = next->next;
    100:   }
    101:   MatMult(next->mat,in,y);
    102:   if (shell->left) {
    103:     VecPointwiseMult(y,shell->left,y);
    104:   }
    105:   scale = shell->scale;
    106:   if (shell->scalings) {for (i=0; i<shell->nmat; i++) scale *= shell->scalings[i];}
    107:   VecScale(y,scale);
    108:   return(0);
    109: }

    111: PetscErrorCode MatMultTranspose_Composite_Multiplicative(Mat A,Vec x,Vec y)
    112: {
    113:   Mat_Composite     *shell = (Mat_Composite*)A->data;
    114:   Mat_CompositeLink tail   = shell->tail;
    115:   PetscErrorCode    ierr;
    116:   Vec               in,out;
    117:   PetscScalar       scale;
    118:   PetscInt          i;

    121:   if (!tail) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
    122:   in = x;
    123:   if (shell->left) {
    124:     if (!shell->leftwork) {
    125:       VecDuplicate(shell->left,&shell->leftwork);
    126:     }
    127:     VecPointwiseMult(shell->leftwork,shell->left,in);
    128:     in   = shell->leftwork;
    129:   }
    130:   while (tail->prev) {
    131:     if (!tail->prev->work) { /* should reuse previous work if the same size */
    132:       MatCreateVecs(tail->mat,NULL,&tail->prev->work);
    133:     }
    134:     out  = tail->prev->work;
    135:     MatMultTranspose(tail->mat,in,out);
    136:     in   = out;
    137:     tail = tail->prev;
    138:   }
    139:   MatMultTranspose(tail->mat,in,y);
    140:   if (shell->right) {
    141:     VecPointwiseMult(y,shell->right,y);
    142:   }

    144:   scale = shell->scale;
    145:   if (shell->scalings) {for (i=0; i<shell->nmat; i++) scale *= shell->scalings[i];}
    146:   VecScale(y,scale);
    147:   return(0);
    148: }

    150: PetscErrorCode MatMult_Composite(Mat mat,Vec x,Vec y)
    151: {
    152:   PetscErrorCode    ierr;
    153:   Mat_Composite     *shell = (Mat_Composite*)mat->data;
    154:   Mat_CompositeLink cur = shell->head;
    155:   Vec               in,y2,xin;
    156:   Mat               A,B;
    157:   PetscInt          i,j,k,n,nuniq,lo,hi,mid,*gindices,*buf,*tmp,tot;
    158:   const PetscScalar *vals;
    159:   const PetscInt    *garray;
    160:   IS                ix,iy;
    161:   PetscBool         match;

    164:   if (!cur) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
    165:   in = x;
    166:   if (shell->right) {
    167:     if (!shell->rightwork) {
    168:       VecDuplicate(shell->right,&shell->rightwork);
    169:     }
    170:     VecPointwiseMult(shell->rightwork,shell->right,in);
    171:     in   = shell->rightwork;
    172:   }

    174:   /* Try to merge Mvctx when instructed but not yet done. We did not do it in MatAssemblyEnd() since at that time
    175:      we did not know whether mat is ADDITIVE or MULTIPLICATIVE. Only now we are assured mat is ADDITIVE and
    176:      it is legal to merge Mvctx, because all component matrices have the same size.
    177:    */
    178:   if (shell->merge_mvctx && !shell->Mvctx) {
    179:     /* Currently only implemented for MATMPIAIJ */
    180:     for (cur=shell->head; cur; cur=cur->next) {
    181:       PetscObjectTypeCompare((PetscObject)cur->mat,MATMPIAIJ,&match);
    182:       if (!match) {
    183:         shell->merge_mvctx = PETSC_FALSE;
    184:         goto skip_merge_mvctx;
    185:       }
    186:     }

    188:     /* Go through matrices first time to count total number of nonzero off-diag columns (may have dups) */
    189:     tot = 0;
    190:     for (cur=shell->head; cur; cur=cur->next) {
    191:       MatMPIAIJGetSeqAIJ(cur->mat,NULL,&B,NULL);
    192:       MatGetLocalSize(B,NULL,&n);
    193:       tot += n;
    194:     }
    195:     PetscMalloc3(tot,&shell->location,tot,&shell->larray,shell->nmat,&shell->lvecs);
    196:     shell->len = tot;

    198:     /* Go through matrices second time to sort off-diag columns and remove dups */
    199:     PetscMalloc1(tot,&gindices); /* No Malloc2() since we will give one to petsc and free the other */
    200:     PetscMalloc1(tot,&buf);
    201:     nuniq = 0; /* Number of unique nonzero columns */
    202:     for (cur=shell->head; cur; cur=cur->next) {
    203:       MatMPIAIJGetSeqAIJ(cur->mat,NULL,&B,&garray);
    204:       MatGetLocalSize(B,NULL,&n);
    205:       /* Merge pre-sorted garray[0,n) and gindices[0,nuniq) to buf[] */
    206:       i = j = k = 0;
    207:       while (i < n && j < nuniq) {
    208:         if (garray[i] < gindices[j]) buf[k++] = garray[i++];
    209:         else if (garray[i] > gindices[j]) buf[k++] = gindices[j++];
    210:         else {buf[k++] = garray[i++]; j++;}
    211:       }
    212:       /* Copy leftover in garray[] or gindices[] */
    213:       if (i < n) {
    214:         PetscArraycpy(buf+k,garray+i,n-i);
    215:         nuniq = k + n-i;
    216:       } else if (j < nuniq) {
    217:         PetscArraycpy(buf+k,gindices+j,nuniq-j);
    218:         nuniq = k + nuniq-j;
    219:       } else nuniq = k;
    220:       /* Swap gindices and buf to merge garray of the next matrix */
    221:       tmp      = gindices;
    222:       gindices = buf;
    223:       buf      = tmp;
    224:     }
    225:     PetscFree(buf);

    227:     /* Go through matrices third time to build a map from gindices[] to garray[] */
    228:     tot = 0;
    229:     for (cur=shell->head,j=0; cur; cur=cur->next,j++) { /* j-th matrix */
    230:       MatMPIAIJGetSeqAIJ(cur->mat,NULL,&B,&garray);
    231:       MatGetLocalSize(B,NULL,&n);
    232:       VecCreateSeqWithArray(PETSC_COMM_SELF,1,n,NULL,&shell->lvecs[j]);
    233:       /* This is an optimized PetscFindInt(garray[i],nuniq,gindices,&shell->location[tot+i]), using the fact that garray[] is also sorted */
    234:       lo   = 0;
    235:       for (i=0; i<n; i++) {
    236:         hi = nuniq;
    237:         while (hi - lo > 1) {
    238:           mid = lo + (hi - lo)/2;
    239:           if (garray[i] < gindices[mid]) hi = mid;
    240:           else lo = mid;
    241:         }
    242:         shell->location[tot+i] = lo; /* gindices[lo] = garray[i] */
    243:         lo++; /* Since garray[i+1] > garray[i], we can safely advance lo */
    244:       }
    245:       tot += n;
    246:     }

    248:     /* Build merged Mvctx */
    249:     ISCreateGeneral(PETSC_COMM_SELF,nuniq,gindices,PETSC_OWN_POINTER,&ix);
    250:     ISCreateStride(PETSC_COMM_SELF,nuniq,0,1,&iy);
    251:     VecCreateMPIWithArray(PetscObjectComm((PetscObject)mat),1,mat->cmap->n,mat->cmap->N,NULL,&xin);
    252:     VecCreateSeq(PETSC_COMM_SELF,nuniq,&shell->gvec);
    253:     VecScatterCreate(xin,ix,shell->gvec,iy,&shell->Mvctx);
    254:     VecDestroy(&xin);
    255:     ISDestroy(&ix);
    256:     ISDestroy(&iy);
    257:   }

    259: skip_merge_mvctx:
    260:   VecSet(y,0);
    261:   if (!shell->leftwork2) {VecDuplicate(y,&shell->leftwork2);}
    262:   y2 = shell->leftwork2;

    264:   if (shell->Mvctx) { /* Have a merged Mvctx */
    265:     /* 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
    266:        in y = s(Ax1 + Bx2) or y = sAx1 + sBx2. The former incurs less FLOPS than the latter, but the latter provides an oppertunity to
    267:        overlap communication/computation since we can do sAx1 while communicating x2. Here, we use the former approach.
    268:      */
    269:     VecScatterBegin(shell->Mvctx,in,shell->gvec,INSERT_VALUES,SCATTER_FORWARD);
    270:     VecScatterEnd(shell->Mvctx,in,shell->gvec,INSERT_VALUES,SCATTER_FORWARD);

    272:     VecGetArrayRead(shell->gvec,&vals);
    273:     for (i=0; i<shell->len; i++) shell->larray[i] = vals[shell->location[i]];
    274:     VecRestoreArrayRead(shell->gvec,&vals);

    276:     for (cur=shell->head,tot=i=0; cur; cur=cur->next,i++) { /* i-th matrix */
    277:       MatMPIAIJGetSeqAIJ(cur->mat,&A,&B,NULL);
    278:       (*A->ops->mult)(A,in,y2);
    279:       MatGetLocalSize(B,NULL,&n);
    280:       VecPlaceArray(shell->lvecs[i],&shell->larray[tot]);
    281:       (*B->ops->multadd)(B,shell->lvecs[i],y2,y2);
    282:       VecResetArray(shell->lvecs[i]);
    283:       VecAXPY(y,(shell->scalings ? shell->scalings[i] : 1.0),y2);
    284:       tot += n;
    285:     }
    286:   } else {
    287:     if (shell->scalings) {
    288:       for (cur=shell->head,i=0; cur; cur=cur->next,i++) {
    289:         MatMult(cur->mat,in,y2);
    290:         VecAXPY(y,shell->scalings[i],y2);
    291:       }
    292:     } else {
    293:       for (cur=shell->head; cur; cur=cur->next) {MatMultAdd(cur->mat,in,y,y);}
    294:     }
    295:   }

    297:   if (shell->left) {VecPointwiseMult(y,shell->left,y);}
    298:   VecScale(y,shell->scale);
    299:   return(0);
    300: }

    302: PetscErrorCode MatMultTranspose_Composite(Mat A,Vec x,Vec y)
    303: {
    304:   Mat_Composite     *shell = (Mat_Composite*)A->data;
    305:   Mat_CompositeLink next   = shell->head;
    306:   PetscErrorCode    ierr;
    307:   Vec               in,y2 = NULL;
    308:   PetscInt          i;

    311:   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
    312:   in = x;
    313:   if (shell->left) {
    314:     if (!shell->leftwork) {
    315:       VecDuplicate(shell->left,&shell->leftwork);
    316:     }
    317:     VecPointwiseMult(shell->leftwork,shell->left,in);
    318:     in   = shell->leftwork;
    319:   }

    321:   MatMultTranspose(next->mat,in,y);
    322:   if (shell->scalings) {
    323:     VecScale(y,shell->scalings[0]);
    324:     if (!shell->rightwork2) {VecDuplicate(y,&shell->rightwork2);}
    325:     y2 = shell->rightwork2;
    326:   }
    327:   i = 1;
    328:   while ((next = next->next)) {
    329:     if (!shell->scalings) {MatMultTransposeAdd(next->mat,in,y,y);}
    330:     else {
    331:       MatMultTranspose(next->mat,in,y2);
    332:       VecAXPY(y,shell->scalings[i++],y2);
    333:     }
    334:   }
    335:   if (shell->right) {
    336:     VecPointwiseMult(y,shell->right,y);
    337:   }
    338:   VecScale(y,shell->scale);
    339:   return(0);
    340: }

    342: PetscErrorCode MatMultAdd_Composite(Mat A,Vec x,Vec y,Vec z)
    343: {
    344:   Mat_Composite     *shell = (Mat_Composite*)A->data;
    345:   PetscErrorCode    ierr;

    348:   if (y != z) {
    349:     MatMult(A,x,z);
    350:     VecAXPY(z,1.0,y);
    351:   } else {
    352:     if (!shell->leftwork) {
    353:       VecDuplicate(z,&shell->leftwork);
    354:     }
    355:     MatMult(A,x,shell->leftwork);
    356:     VecCopy(y,z);
    357:     VecAXPY(z,1.0,shell->leftwork);
    358:   }
    359:   return(0);
    360: }

    362: PetscErrorCode MatMultTransposeAdd_Composite(Mat A,Vec x,Vec y, Vec z)
    363: {
    364:   Mat_Composite     *shell = (Mat_Composite*)A->data;
    365:   PetscErrorCode    ierr;

    368:   if (y != z) {
    369:     MatMultTranspose(A,x,z);
    370:     VecAXPY(z,1.0,y);
    371:   } else {
    372:     if (!shell->rightwork) {
    373:       VecDuplicate(z,&shell->rightwork);
    374:     }
    375:     MatMultTranspose(A,x,shell->rightwork);
    376:     VecCopy(y,z);
    377:     VecAXPY(z,1.0,shell->rightwork);
    378:   }
    379:   return(0);
    380: }

    382: PetscErrorCode MatGetDiagonal_Composite(Mat A,Vec v)
    383: {
    384:   Mat_Composite     *shell = (Mat_Composite*)A->data;
    385:   Mat_CompositeLink next   = shell->head;
    386:   PetscErrorCode    ierr;
    387:   PetscInt          i;

    390:   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
    391:   if (shell->right || shell->left) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Cannot get diagonal if left or right scaling");

    393:   MatGetDiagonal(next->mat,v);
    394:   if (shell->scalings) {VecScale(v,shell->scalings[0]);}

    396:   if (next->next && !shell->work) {
    397:     VecDuplicate(v,&shell->work);
    398:   }
    399:   i = 1;
    400:   while ((next = next->next)) {
    401:     MatGetDiagonal(next->mat,shell->work);
    402:     VecAXPY(v,(shell->scalings ? shell->scalings[i++] : 1.0),shell->work);
    403:   }
    404:   VecScale(v,shell->scale);
    405:   return(0);
    406: }

    408: PetscErrorCode MatAssemblyEnd_Composite(Mat Y,MatAssemblyType t)
    409: {
    410:   Mat_Composite     *shell = (Mat_Composite*)Y->data;
    411:   PetscErrorCode    ierr;

    414:   if (shell->merge) {
    415:     MatCompositeMerge(Y);
    416:   }

    418:   return(0);
    419: }

    421: PetscErrorCode MatScale_Composite(Mat inA,PetscScalar alpha)
    422: {
    423:   Mat_Composite *a = (Mat_Composite*)inA->data;

    426:   a->scale *= alpha;
    427:   return(0);
    428: }

    430: PetscErrorCode MatDiagonalScale_Composite(Mat inA,Vec left,Vec right)
    431: {
    432:   Mat_Composite  *a = (Mat_Composite*)inA->data;

    436:   if (left) {
    437:     if (!a->left) {
    438:       VecDuplicate(left,&a->left);
    439:       VecCopy(left,a->left);
    440:     } else {
    441:       VecPointwiseMult(a->left,left,a->left);
    442:     }
    443:   }
    444:   if (right) {
    445:     if (!a->right) {
    446:       VecDuplicate(right,&a->right);
    447:       VecCopy(right,a->right);
    448:     } else {
    449:       VecPointwiseMult(a->right,right,a->right);
    450:     }
    451:   }
    452:   return(0);
    453: }

    455: PetscErrorCode MatSetFromOptions_Composite(PetscOptionItems *PetscOptionsObject,Mat A)
    456: {
    457:   Mat_Composite  *a = (Mat_Composite*)A->data;

    461:   PetscOptionsHead(PetscOptionsObject,"MATCOMPOSITE options");
    462:   PetscOptionsBool("-mat_composite_merge","Merge at MatAssemblyEnd","MatCompositeMerge",a->merge,&a->merge,NULL);
    463:   PetscOptionsEnum("-mat_composite_merge_type","Set composite merge direction","MatCompositeSetMergeType",MatCompositeMergeTypes,(PetscEnum)a->mergetype,(PetscEnum*)&a->mergetype,NULL);
    464:   PetscOptionsBool("-mat_composite_merge_mvctx","Merge MatMult() vecscat contexts","MatCreateComposite",a->merge_mvctx,&a->merge_mvctx,NULL);
    465:   PetscOptionsTail();
    466:   return(0);
    467: }

    469: /*@
    470:    MatCreateComposite - Creates a matrix as the sum or product of one or more matrices

    472:   Collective

    474:    Input Parameters:
    475: +  comm - MPI communicator
    476: .  nmat - number of matrices to put in
    477: -  mats - the matrices

    479:    Output Parameter:
    480: .  mat - the matrix

    482:    Options Database Keys:
    483: +  -mat_composite_merge         - merge in MatAssemblyEnd()
    484: .  -mat_composite_merge_mvctx   - merge Mvctx of component matrices to optimize communication in MatMult() for ADDITIVE matrices
    485: -  -mat_composite_merge_type    - set merge direction

    487:    Level: advanced

    489:    Notes:
    490:      Alternative construction
    491: $       MatCreate(comm,&mat);
    492: $       MatSetSizes(mat,m,n,M,N);
    493: $       MatSetType(mat,MATCOMPOSITE);
    494: $       MatCompositeAddMat(mat,mats[0]);
    495: $       ....
    496: $       MatCompositeAddMat(mat,mats[nmat-1]);
    497: $       MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
    498: $       MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);

    500:      For the multiplicative form the product is mat[nmat-1]*mat[nmat-2]*....*mat[0]

    502: .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCompositeGetMat(), MatCompositeMerge(), MatCompositeSetType(), MATCOMPOSITE

    504: @*/
    505: PetscErrorCode MatCreateComposite(MPI_Comm comm,PetscInt nmat,const Mat *mats,Mat *mat)
    506: {
    508:   PetscInt       m,n,M,N,i;

    511:   if (nmat < 1) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Must pass in at least one matrix");

    514:   MatGetLocalSize(mats[0],PETSC_IGNORE,&n);
    515:   MatGetLocalSize(mats[nmat-1],&m,PETSC_IGNORE);
    516:   MatGetSize(mats[0],PETSC_IGNORE,&N);
    517:   MatGetSize(mats[nmat-1],&M,PETSC_IGNORE);
    518:   MatCreate(comm,mat);
    519:   MatSetSizes(*mat,m,n,M,N);
    520:   MatSetType(*mat,MATCOMPOSITE);
    521:   for (i=0; i<nmat; i++) {
    522:     MatCompositeAddMat(*mat,mats[i]);
    523:   }
    524:   MatAssemblyBegin(*mat,MAT_FINAL_ASSEMBLY);
    525:   MatAssemblyEnd(*mat,MAT_FINAL_ASSEMBLY);
    526:   return(0);
    527: }

    530: static PetscErrorCode MatCompositeAddMat_Composite(Mat mat,Mat smat)
    531: {
    532:   Mat_Composite     *shell = (Mat_Composite*)mat->data;
    533:   Mat_CompositeLink ilink,next = shell->head;
    534:   PetscErrorCode    ierr;

    537:   PetscNewLog(mat,&ilink);
    538:   ilink->next = 0;
    539:   PetscObjectReference((PetscObject)smat);
    540:   ilink->mat  = smat;

    542:   if (!next) shell->head = ilink;
    543:   else {
    544:     while (next->next) {
    545:       next = next->next;
    546:     }
    547:     next->next  = ilink;
    548:     ilink->prev = next;
    549:   }
    550:   shell->tail =  ilink;
    551:   shell->nmat += 1;

    553:   /* Retain the old scalings (if any) and expand it with a 1.0 for the newly added matrix */
    554:   if (shell->scalings) {
    555:     PetscRealloc(sizeof(PetscScalar)*shell->nmat,&shell->scalings);
    556:     shell->scalings[shell->nmat-1] = 1.0;
    557:   }
    558:   return(0);
    559: }

    561: /*@
    562:     MatCompositeAddMat - Add another matrix to a composite matrix.

    564:    Collective on Mat

    566:     Input Parameters:
    567: +   mat - the composite matrix
    568: -   smat - the partial matrix

    570:    Level: advanced

    572: .seealso: MatCreateComposite(), MatCompositeGetMat(), MATCOMPOSITE
    573: @*/
    574: PetscErrorCode MatCompositeAddMat(Mat mat,Mat smat)
    575: {
    576:   PetscErrorCode    ierr;

    581:   PetscUseMethod(mat,"MatCompositeAddMat_C",(Mat,Mat),(mat,smat));
    582:   return(0);
    583: }

    585: static PetscErrorCode MatCompositeSetType_Composite(Mat mat,MatCompositeType type)
    586: {
    587:   Mat_Composite  *b = (Mat_Composite*)mat->data;

    590:   b->type = type;
    591:   if (type == MAT_COMPOSITE_MULTIPLICATIVE) {
    592:     mat->ops->getdiagonal   = 0;
    593:     mat->ops->mult          = MatMult_Composite_Multiplicative;
    594:     mat->ops->multtranspose = MatMultTranspose_Composite_Multiplicative;
    595:     b->merge_mvctx          = PETSC_FALSE;
    596:   } else {
    597:     mat->ops->getdiagonal   = MatGetDiagonal_Composite;
    598:     mat->ops->mult          = MatMult_Composite;
    599:     mat->ops->multtranspose = MatMultTranspose_Composite;
    600:   }
    601:   return(0);
    602: }

    604: /*@
    605:    MatCompositeSetType - Indicates if the matrix is defined as the sum of a set of matrices or the product.

    607:    Logically Collective on Mat

    609:    Input Parameters:
    610: .  mat - the composite matrix

    612:    Level: advanced

    614: .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCreateComposite(), MatCompositeGetType(), MATCOMPOSITE

    616: @*/
    617: PetscErrorCode MatCompositeSetType(Mat mat,MatCompositeType type)
    618: {

    624:   PetscUseMethod(mat,"MatCompositeSetType_C",(Mat,MatCompositeType),(mat,type));
    625:   return(0);
    626: }

    628: static PetscErrorCode MatCompositeGetType_Composite(Mat mat,MatCompositeType *type)
    629: {
    630:   Mat_Composite  *b = (Mat_Composite*)mat->data;

    633:   *type = b->type;
    634:   return(0);
    635: }

    637: /*@
    638:    MatCompositeGetType - Returns type of composite.

    640:    Not Collective

    642:    Input Parameter:
    643: .  mat - the composite matrix

    645:    Output Parameter:
    646: .  type - type of composite

    648:    Level: advanced

    650: .seealso: MatCreateComposite(), MatCompositeSetType(), MATCOMPOSITE

    652: @*/
    653: PetscErrorCode MatCompositeGetType(Mat mat,MatCompositeType *type)
    654: {

    660:   PetscUseMethod(mat,"MatCompositeGetType_C",(Mat,MatCompositeType*),(mat,type));
    661:   return(0);
    662: }

    664: static PetscErrorCode MatCompositeSetMatStructure_Composite(Mat mat,MatStructure str)
    665: {
    666:   Mat_Composite  *b = (Mat_Composite*)mat->data;

    669:   b->structure = str;
    670:   return(0);
    671: }

    673: /*@
    674:    MatCompositeSetMatStructure - Indicates structure of matrices in the composite matrix.

    676:    Not Collective

    678:    Input Parameters:
    679: +  mat - the composite matrix
    680: -  str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN (default) or SUBSET_NONZERO_PATTERN

    682:    Level: advanced

    684:    Notes:
    685:     Information about the matrices structure is used in MatCompositeMerge() for additive composite matrix.

    687: .seealso: MatAXPY(), MatCreateComposite(), MatCompositeMerge() MatCompositeGetMatStructure(), MATCOMPOSITE

    689: @*/
    690: PetscErrorCode MatCompositeSetMatStructure(Mat mat,MatStructure str)
    691: {

    696:   PetscUseMethod(mat,"MatCompositeSetMatStructure_C",(Mat,MatStructure),(mat,str));
    697:   return(0);
    698: }

    700: static PetscErrorCode MatCompositeGetMatStructure_Composite(Mat mat,MatStructure *str)
    701: {
    702:   Mat_Composite  *b = (Mat_Composite*)mat->data;

    705:   *str = b->structure;
    706:   return(0);
    707: }

    709: /*@
    710:    MatCompositeGetMatStructure - Returns the structure of matrices in the composite matrix.

    712:    Not Collective

    714:    Input Parameter:
    715: .  mat - the composite matrix

    717:    Output Parameter:
    718: .  str - structure of the matrices

    720:    Level: advanced

    722: .seealso: MatCreateComposite(), MatCompositeSetMatStructure(), MATCOMPOSITE

    724: @*/
    725: PetscErrorCode MatCompositeGetMatStructure(Mat mat,MatStructure *str)
    726: {

    732:   PetscUseMethod(mat,"MatCompositeGetMatStructure_C",(Mat,MatStructure*),(mat,str));
    733:   return(0);
    734: }

    736: static PetscErrorCode MatCompositeSetMergeType_Composite(Mat mat,MatCompositeMergeType type)
    737: {
    738:   Mat_Composite     *shell = (Mat_Composite*)mat->data;

    741:   shell->mergetype = type;
    742:   return(0);
    743: }

    745: /*@
    746:    MatCompositeSetMergeType - Sets order of MatCompositeMerge().

    748:    Logically Collective on Mat

    750:    Input Parameter:
    751: +  mat - the composite matrix
    752: -  type - MAT_COMPOSITE_MERGE RIGHT (default) to start merge from right with the first added matrix (mat[0]),
    753:           MAT_COMPOSITE_MERGE_LEFT to start merge from left with the last added matrix (mat[nmat-1])

    755:    Level: advanced

    757:    Notes:
    758:     The resulting matrix is the same regardles of the MergeType. Only the order of operation is changed.
    759:     If set to MAT_COMPOSITE_MERGE_RIGHT the order of the merge is mat[nmat-1]*(mat[nmat-2]*(...*(mat[1]*mat[0])))
    760:     otherwise the order is (((mat[nmat-1]*mat[nmat-2])*mat[nmat-3])*...)*mat[0].

    762: .seealso: MatCreateComposite(), MatCompositeMerge(), MATCOMPOSITE

    764: @*/
    765: PetscErrorCode MatCompositeSetMergeType(Mat mat,MatCompositeMergeType type)
    766: {

    772:   PetscUseMethod(mat,"MatCompositeSetMergeType_C",(Mat,MatCompositeMergeType),(mat,type));
    773:   return(0);
    774: }

    776: static PetscErrorCode MatCompositeMerge_Composite(Mat mat)
    777: {
    778:   Mat_Composite     *shell = (Mat_Composite*)mat->data;
    779:   Mat_CompositeLink next   = shell->head, prev = shell->tail;
    780:   PetscErrorCode    ierr;
    781:   Mat               tmat,newmat;
    782:   Vec               left,right;
    783:   PetscScalar       scale;
    784:   PetscInt          i;

    787:   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");

    790:   scale = shell->scale;
    791:   if (shell->type == MAT_COMPOSITE_ADDITIVE) {
    792:     if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) {
    793:       i = 0;
    794:       MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);
    795:       if (shell->scalings) {MatScale(tmat,shell->scalings[i++]);}
    796:       while ((next = next->next)) {
    797:         MatAXPY(tmat,(shell->scalings ? shell->scalings[i++] : 1.0),next->mat,shell->structure);
    798:       }
    799:     } else {
    800:       i = shell->nmat-1;
    801:       MatDuplicate(prev->mat,MAT_COPY_VALUES,&tmat);
    802:       if (shell->scalings) {MatScale(tmat,shell->scalings[i--]);}
    803:       while ((prev = prev->prev)) {
    804:         MatAXPY(tmat,(shell->scalings ? shell->scalings[i--] : 1.0),prev->mat,shell->structure);
    805:       }
    806:     }
    807:   } else {
    808:     if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) {
    809:       MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);
    810:       while ((next = next->next)) {
    811:         MatMatMult(next->mat,tmat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat);
    812:         MatDestroy(&tmat);
    813:         tmat = newmat;
    814:       }
    815:     } else {
    816:       MatDuplicate(prev->mat,MAT_COPY_VALUES,&tmat);
    817:       while ((prev = prev->prev)) {
    818:         MatMatMult(tmat,prev->mat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat);
    819:         MatDestroy(&tmat);
    820:         tmat = newmat;
    821:       }
    822:     }
    823:     if (shell->scalings) {for (i=0; i<shell->nmat; i++) scale *= shell->scalings[i];}
    824:   }

    826:   if ((left = shell->left)) {PetscObjectReference((PetscObject)left);}
    827:   if ((right = shell->right)) {PetscObjectReference((PetscObject)right);}

    829:   MatHeaderReplace(mat,&tmat);

    831:   MatDiagonalScale(mat,left,right);
    832:   MatScale(mat,scale);
    833:   VecDestroy(&left);
    834:   VecDestroy(&right);
    835:   return(0);
    836: }

    838: /*@
    839:    MatCompositeMerge - Given a composite matrix, replaces it with a "regular" matrix
    840:      by summing or computing the product of all the matrices inside the composite matrix.

    842:   Collective

    844:    Input Parameters:
    845: .  mat - the composite matrix

    848:    Options Database Keys:
    849: +  -mat_composite_merge - merge in MatAssemblyEnd()
    850: -  -mat_composite_merge_type - set merge direction

    852:    Level: advanced

    854:    Notes:
    855:       The MatType of the resulting matrix will be the same as the MatType of the FIRST
    856:     matrix in the composite matrix.

    858: .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCreateComposite(), MatCompositeSetMatStructure(), MatCompositeSetMergeType(), MATCOMPOSITE

    860: @*/
    861: PetscErrorCode MatCompositeMerge(Mat mat)
    862: {

    867:   PetscUseMethod(mat,"MatCompositeMerge_C",(Mat),(mat));
    868:   return(0);
    869: }

    871: static PetscErrorCode MatCompositeGetNumberMat_Composite(Mat mat,PetscInt *nmat)
    872: {
    873:   Mat_Composite     *shell = (Mat_Composite*)mat->data;

    876:   *nmat = shell->nmat;
    877:   return(0);
    878: }

    880: /*@
    881:    MatCompositeGetNumberMat - Returns the number of matrices in the composite matrix.

    883:    Not Collective

    885:    Input Parameter:
    886: .  mat - the composite matrix

    888:    Output Parameter:
    889: .  nmat - number of matrices in the composite matrix

    891:    Level: advanced

    893: .seealso: MatCreateComposite(), MatCompositeGetMat(), MATCOMPOSITE

    895: @*/
    896: PetscErrorCode MatCompositeGetNumberMat(Mat mat,PetscInt *nmat)
    897: {

    903:   PetscUseMethod(mat,"MatCompositeGetNumberMat_C",(Mat,PetscInt*),(mat,nmat));
    904:   return(0);
    905: }

    907: static PetscErrorCode MatCompositeGetMat_Composite(Mat mat,PetscInt i,Mat *Ai)
    908: {
    909:   Mat_Composite     *shell = (Mat_Composite*)mat->data;
    910:   Mat_CompositeLink ilink;
    911:   PetscInt          k;

    914:   if (i >= shell->nmat) SETERRQ2(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_OUTOFRANGE,"index out of range: %d >= %d",i,shell->nmat);
    915:   ilink = shell->head;
    916:   for (k=0; k<i; k++) {
    917:     ilink = ilink->next;
    918:   }
    919:   *Ai = ilink->mat;
    920:   return(0);
    921: }

    923: /*@
    924:    MatCompositeGetMat - Returns the ith matrix from the composite matrix.

    926:    Logically Collective on Mat

    928:    Input Parameter:
    929: +  mat - the composite matrix
    930: -  i - the number of requested matrix

    932:    Output Parameter:
    933: .  Ai - ith matrix in composite

    935:    Level: advanced

    937: .seealso: MatCreateComposite(), MatCompositeGetNumberMat(), MatCompositeAddMat(), MATCOMPOSITE

    939: @*/
    940: PetscErrorCode MatCompositeGetMat(Mat mat,PetscInt i,Mat *Ai)
    941: {

    948:   PetscUseMethod(mat,"MatCompositeGetMat_C",(Mat,PetscInt,Mat*),(mat,i,Ai));
    949:   return(0);
    950: }

    952: PetscErrorCode MatCompositeSetScalings_Composite(Mat mat,const PetscScalar *scalings)
    953: {
    955:   Mat_Composite  *shell = (Mat_Composite*)mat->data;
    956:   PetscInt       nmat;

    959:   MatCompositeGetNumberMat(mat,&nmat);
    960:   if (!shell->scalings) {PetscMalloc1(nmat,&shell->scalings);}
    961:   PetscArraycpy(shell->scalings,scalings,nmat);
    962:   return(0);
    963: }

    965: /*@
    966:    MatCompositeSetScalings - Sets separate scaling factors for component matrices.

    968:    Logically Collective on Mat

    970:    Input Parameter:
    971: +  mat      - the composite matrix
    972: -  scalings - array of scaling factors with scalings[i] being factor of i-th matrix, for i in [0, nmat)

    974:    Level: advanced

    976: .seealso: MatScale(), MatDiagonalScale(), MATCOMPOSITE

    978: @*/
    979: PetscErrorCode MatCompositeSetScalings(Mat mat,const PetscScalar *scalings)
    980: {

    987:   PetscUseMethod(mat,"MatCompositeSetScalings_C",(Mat,const PetscScalar*),(mat,scalings));
    988:   return(0);
    989: }

    991: static struct _MatOps MatOps_Values = {0,
    992:                                        0,
    993:                                        0,
    994:                                        MatMult_Composite,
    995:                                        MatMultAdd_Composite,
    996:                                 /*  5*/ MatMultTranspose_Composite,
    997:                                        MatMultTransposeAdd_Composite,
    998:                                        0,
    999:                                        0,
    1000:                                        0,
    1001:                                 /* 10*/ 0,
    1002:                                        0,
    1003:                                        0,
    1004:                                        0,
    1005:                                        0,
    1006:                                 /* 15*/ 0,
    1007:                                        0,
    1008:                                        MatGetDiagonal_Composite,
    1009:                                        MatDiagonalScale_Composite,
    1010:                                        0,
    1011:                                 /* 20*/ 0,
    1012:                                        MatAssemblyEnd_Composite,
    1013:                                        0,
    1014:                                        0,
    1015:                                /* 24*/ 0,
    1016:                                        0,
    1017:                                        0,
    1018:                                        0,
    1019:                                        0,
    1020:                                /* 29*/ 0,
    1021:                                        0,
    1022:                                        0,
    1023:                                        0,
    1024:                                        0,
    1025:                                /* 34*/ 0,
    1026:                                        0,
    1027:                                        0,
    1028:                                        0,
    1029:                                        0,
    1030:                                /* 39*/ 0,
    1031:                                        0,
    1032:                                        0,
    1033:                                        0,
    1034:                                        0,
    1035:                                /* 44*/ 0,
    1036:                                        MatScale_Composite,
    1037:                                        MatShift_Basic,
    1038:                                        0,
    1039:                                        0,
    1040:                                /* 49*/ 0,
    1041:                                        0,
    1042:                                        0,
    1043:                                        0,
    1044:                                        0,
    1045:                                /* 54*/ 0,
    1046:                                        0,
    1047:                                        0,
    1048:                                        0,
    1049:                                        0,
    1050:                                /* 59*/ 0,
    1051:                                        MatDestroy_Composite,
    1052:                                        0,
    1053:                                        0,
    1054:                                        0,
    1055:                                /* 64*/ 0,
    1056:                                        0,
    1057:                                        0,
    1058:                                        0,
    1059:                                        0,
    1060:                                /* 69*/ 0,
    1061:                                        0,
    1062:                                        0,
    1063:                                        0,
    1064:                                        0,
    1065:                                /* 74*/ 0,
    1066:                                        0,
    1067:                                        MatSetFromOptions_Composite,
    1068:                                        0,
    1069:                                        0,
    1070:                                /* 79*/ 0,
    1071:                                        0,
    1072:                                        0,
    1073:                                        0,
    1074:                                        0,
    1075:                                /* 84*/ 0,
    1076:                                        0,
    1077:                                        0,
    1078:                                        0,
    1079:                                        0,
    1080:                                /* 89*/ 0,
    1081:                                        0,
    1082:                                        0,
    1083:                                        0,
    1084:                                        0,
    1085:                                /* 94*/ 0,
    1086:                                        0,
    1087:                                        0,
    1088:                                        0,
    1089:                                        0,
    1090:                                 /*99*/ 0,
    1091:                                        0,
    1092:                                        0,
    1093:                                        0,
    1094:                                        0,
    1095:                                /*104*/ 0,
    1096:                                        0,
    1097:                                        0,
    1098:                                        0,
    1099:                                        0,
    1100:                                /*109*/ 0,
    1101:                                        0,
    1102:                                        0,
    1103:                                        0,
    1104:                                        0,
    1105:                                /*114*/ 0,
    1106:                                        0,
    1107:                                        0,
    1108:                                        0,
    1109:                                        0,
    1110:                                /*119*/ 0,
    1111:                                        0,
    1112:                                        0,
    1113:                                        0,
    1114:                                        0,
    1115:                                /*124*/ 0,
    1116:                                        0,
    1117:                                        0,
    1118:                                        0,
    1119:                                        0,
    1120:                                /*129*/ 0,
    1121:                                        0,
    1122:                                        0,
    1123:                                        0,
    1124:                                        0,
    1125:                                /*134*/ 0,
    1126:                                        0,
    1127:                                        0,
    1128:                                        0,
    1129:                                        0,
    1130:                                /*139*/ 0,
    1131:                                        0,
    1132:                                        0
    1133: };

    1135: /*MC
    1136:    MATCOMPOSITE - A matrix defined by the sum (or product) of one or more matrices.
    1137:     The matrices need to have a correct size and parallel layout for the sum or product to be valid.

    1139:    Notes:
    1140:     to use the product of the matrices call MatCompositeSetType(mat,MAT_COMPOSITE_MULTIPLICATIVE);

    1142:   Level: advanced

    1144: .seealso: MatCreateComposite(), MatCompositeSetScalings(), MatCompositeAddMat(), MatSetType(), MatCompositeSetType(), MatCompositeGetType(), MatCompositeSetMatStructure(), MatCompositeGetMatStructure(), MatCompositeMerge(), MatCompositeSetMergeType(), MatCompositeGetNumberMat(), MatCompositeGetMat()
    1145: M*/

    1147: PETSC_EXTERN PetscErrorCode MatCreate_Composite(Mat A)
    1148: {
    1149:   Mat_Composite  *b;

    1153:   PetscNewLog(A,&b);
    1154:   A->data = (void*)b;
    1155:   PetscMemcpy(A->ops,&MatOps_Values,sizeof(struct _MatOps));

    1157:   PetscLayoutSetUp(A->rmap);
    1158:   PetscLayoutSetUp(A->cmap);

    1160:   A->assembled    = PETSC_TRUE;
    1161:   A->preallocated = PETSC_TRUE;
    1162:   b->type         = MAT_COMPOSITE_ADDITIVE;
    1163:   b->scale        = 1.0;
    1164:   b->nmat         = 0;
    1165:   b->merge        = PETSC_FALSE;
    1166:   b->mergetype    = MAT_COMPOSITE_MERGE_RIGHT;
    1167:   b->structure    = DIFFERENT_NONZERO_PATTERN;
    1168:   b->merge_mvctx  = PETSC_TRUE;

    1172:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeAddMat_C",MatCompositeAddMat_Composite);
    1173:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetType_C",MatCompositeSetType_Composite);
    1174:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetType_C",MatCompositeGetType_Composite);
    1175:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetMergeType_C",MatCompositeSetMergeType_Composite);
    1176:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetMatStructure_C",MatCompositeSetMatStructure_Composite);
    1177:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetMatStructure_C",MatCompositeGetMatStructure_Composite);
    1178:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeMerge_C",MatCompositeMerge_Composite);
    1179:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetNumberMat_C",MatCompositeGetNumberMat_Composite);
    1180:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetMat_C",MatCompositeGetMat_Composite);
    1181:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetScalings_C",MatCompositeSetScalings_Composite);

    1183:   PetscObjectChangeTypeName((PetscObject)A,MATCOMPOSITE);
    1184:   return(0);
    1185: }