Date: 2022-03-04T17:37:54+00:00

Actual source code: mcomposite.c

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

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

      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;

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

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

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

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

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

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

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

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

    471:   Collective

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

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

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

    486:    Level: advanced

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

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

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

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

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

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

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

    535:   PetscNewLog(mat,&ilink);
    536:   ilink->next = NULL;
    537:   PetscObjectReference((PetscObject)smat);
    538:   ilink->mat  = smat;

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

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

    559: /*@
    560:     MatCompositeAddMat - Add another matrix to a composite matrix.

    562:    Collective on Mat

    564:     Input Parameters:
    565: +   mat - the composite matrix
    566: -   smat - the partial matrix

    568:    Level: advanced

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

    579:   PetscUseMethod(mat,"MatCompositeAddMat_C",(Mat,Mat),(mat,smat));
    580:   return(0);
    581: }

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

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

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

    605:    Logically Collective on Mat

    607:    Input Parameters:
    608: .  mat - the composite matrix

    610:    Level: advanced

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

    614: @*/
    615: PetscErrorCode MatCompositeSetType(Mat mat,MatCompositeType type)
    616: {

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

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

    631:   *type = b->type;
    632:   return(0);
    633: }

    635: /*@
    636:    MatCompositeGetType - Returns type of composite.

    638:    Not Collective

    640:    Input Parameter:
    641: .  mat - the composite matrix

    643:    Output Parameter:
    644: .  type - type of composite

    646:    Level: advanced

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

    650: @*/
    651: PetscErrorCode MatCompositeGetType(Mat mat,MatCompositeType *type)
    652: {

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

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

    667:   b->structure = str;
    668:   return(0);
    669: }

    671: /*@
    672:    MatCompositeSetMatStructure - Indicates structure of matrices in the composite matrix.

    674:    Not Collective

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

    680:    Level: advanced

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

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

    687: @*/
    688: PetscErrorCode MatCompositeSetMatStructure(Mat mat,MatStructure str)
    689: {

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

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

    703:   *str = b->structure;
    704:   return(0);
    705: }

    707: /*@
    708:    MatCompositeGetMatStructure - Returns the structure of matrices in the composite matrix.

    710:    Not Collective

    712:    Input Parameter:
    713: .  mat - the composite matrix

    715:    Output Parameter:
    716: .  str - structure of the matrices

    718:    Level: advanced

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

    722: @*/
    723: PetscErrorCode MatCompositeGetMatStructure(Mat mat,MatStructure *str)
    724: {

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

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

    739:   shell->mergetype = type;
    740:   return(0);
    741: }

    743: /*@
    744:    MatCompositeSetMergeType - Sets order of MatCompositeMerge().

    746:    Logically Collective on Mat

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

    753:    Level: advanced

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

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

    762: @*/
    763: PetscErrorCode MatCompositeSetMergeType(Mat mat,MatCompositeMergeType type)
    764: {

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

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

    785:   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
    786:   scale = shell->scale;
    787:   if (shell->type == MAT_COMPOSITE_ADDITIVE) {
    788:     if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) {
    789:       i = 0;
    790:       MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);
    791:       if (shell->scalings) {MatScale(tmat,shell->scalings[i++]);}
    792:       while ((next = next->next)) {
    793:         MatAXPY(tmat,(shell->scalings ? shell->scalings[i++] : 1.0),next->mat,shell->structure);
    794:       }
    795:     } else {
    796:       i = shell->nmat-1;
    797:       MatDuplicate(prev->mat,MAT_COPY_VALUES,&tmat);
    798:       if (shell->scalings) {MatScale(tmat,shell->scalings[i--]);}
    799:       while ((prev = prev->prev)) {
    800:         MatAXPY(tmat,(shell->scalings ? shell->scalings[i--] : 1.0),prev->mat,shell->structure);
    801:       }
    802:     }
    803:   } else {
    804:     if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) {
    805:       MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);
    806:       while ((next = next->next)) {
    807:         MatMatMult(next->mat,tmat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat);
    808:         MatDestroy(&tmat);
    809:         tmat = newmat;
    810:       }
    811:     } else {
    812:       MatDuplicate(prev->mat,MAT_COPY_VALUES,&tmat);
    813:       while ((prev = prev->prev)) {
    814:         MatMatMult(tmat,prev->mat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat);
    815:         MatDestroy(&tmat);
    816:         tmat = newmat;
    817:       }
    818:     }
    819:     if (shell->scalings) {for (i=0; i<shell->nmat; i++) scale *= shell->scalings[i];}
    820:   }

    822:   if ((left = shell->left)) {PetscObjectReference((PetscObject)left);}
    823:   if ((right = shell->right)) {PetscObjectReference((PetscObject)right);}

    825:   MatHeaderReplace(mat,&tmat);

    827:   MatDiagonalScale(mat,left,right);
    828:   MatScale(mat,scale);
    829:   VecDestroy(&left);
    830:   VecDestroy(&right);
    831:   return(0);
    832: }

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

    838:   Collective

    840:    Input Parameter:
    841: .  mat - the composite matrix

    843:    Options Database Keys:
    844: +  -mat_composite_merge - merge in MatAssemblyEnd()
    845: -  -mat_composite_merge_type - set merge direction

    847:    Level: advanced

    849:    Notes:
    850:       The MatType of the resulting matrix will be the same as the MatType of the FIRST
    851:     matrix in the composite matrix.

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

    855: @*/
    856: PetscErrorCode MatCompositeMerge(Mat mat)
    857: {

    862:   PetscUseMethod(mat,"MatCompositeMerge_C",(Mat),(mat));
    863:   return(0);
    864: }

    866: static PetscErrorCode MatCompositeGetNumberMat_Composite(Mat mat,PetscInt *nmat)
    867: {
    868:   Mat_Composite     *shell = (Mat_Composite*)mat->data;

    871:   *nmat = shell->nmat;
    872:   return(0);
    873: }

    875: /*@
    876:    MatCompositeGetNumberMat - Returns the number of matrices in the composite matrix.

    878:    Not Collective

    880:    Input Parameter:
    881: .  mat - the composite matrix

    883:    Output Parameter:
    884: .  nmat - number of matrices in the composite matrix

    886:    Level: advanced

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

    890: @*/
    891: PetscErrorCode MatCompositeGetNumberMat(Mat mat,PetscInt *nmat)
    892: {

    898:   PetscUseMethod(mat,"MatCompositeGetNumberMat_C",(Mat,PetscInt*),(mat,nmat));
    899:   return(0);
    900: }

    902: static PetscErrorCode MatCompositeGetMat_Composite(Mat mat,PetscInt i,Mat *Ai)
    903: {
    904:   Mat_Composite     *shell = (Mat_Composite*)mat->data;
    905:   Mat_CompositeLink ilink;
    906:   PetscInt          k;

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

    918: /*@
    919:    MatCompositeGetMat - Returns the ith matrix from the composite matrix.

    921:    Logically Collective on Mat

    923:    Input Parameters:
    924: +  mat - the composite matrix
    925: -  i - the number of requested matrix

    927:    Output Parameter:
    928: .  Ai - ith matrix in composite

    930:    Level: advanced

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

    934: @*/
    935: PetscErrorCode MatCompositeGetMat(Mat mat,PetscInt i,Mat *Ai)
    936: {

    943:   PetscUseMethod(mat,"MatCompositeGetMat_C",(Mat,PetscInt,Mat*),(mat,i,Ai));
    944:   return(0);
    945: }

    947: PetscErrorCode MatCompositeSetScalings_Composite(Mat mat,const PetscScalar *scalings)
    948: {
    950:   Mat_Composite  *shell = (Mat_Composite*)mat->data;
    951:   PetscInt       nmat;

    954:   MatCompositeGetNumberMat(mat,&nmat);
    955:   if (!shell->scalings) {PetscMalloc1(nmat,&shell->scalings);}
    956:   PetscArraycpy(shell->scalings,scalings,nmat);
    957:   return(0);
    958: }

    960: /*@
    961:    MatCompositeSetScalings - Sets separate scaling factors for component matrices.

    963:    Logically Collective on Mat

    965:    Input Parameters:
    966: +  mat      - the composite matrix
    967: -  scalings - array of scaling factors with scalings[i] being factor of i-th matrix, for i in [0, nmat)

    969:    Level: advanced

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

    973: @*/
    974: PetscErrorCode MatCompositeSetScalings(Mat mat,const PetscScalar *scalings)
    975: {

    982:   PetscUseMethod(mat,"MatCompositeSetScalings_C",(Mat,const PetscScalar*),(mat,scalings));
    983:   return(0);
    984: }

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

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

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

    1137:   Level: advanced

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

    1142: PETSC_EXTERN PetscErrorCode MatCreate_Composite(Mat A)
    1143: {
    1144:   Mat_Composite  *b;

    1148:   PetscNewLog(A,&b);
    1149:   A->data = (void*)b;
    1150:   PetscMemcpy(A->ops,&MatOps_Values,sizeof(struct _MatOps));

    1152:   PetscLayoutSetUp(A->rmap);
    1153:   PetscLayoutSetUp(A->cmap);

    1155:   A->assembled    = PETSC_TRUE;
    1156:   A->preallocated = PETSC_TRUE;
    1157:   b->type         = MAT_COMPOSITE_ADDITIVE;
    1158:   b->scale        = 1.0;
    1159:   b->nmat         = 0;
    1160:   b->merge        = PETSC_FALSE;
    1161:   b->mergetype    = MAT_COMPOSITE_MERGE_RIGHT;
    1162:   b->structure    = DIFFERENT_NONZERO_PATTERN;
    1163:   b->merge_mvctx  = PETSC_TRUE;

    1165:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeAddMat_C",MatCompositeAddMat_Composite);
    1166:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetType_C",MatCompositeSetType_Composite);
    1167:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetType_C",MatCompositeGetType_Composite);
    1168:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetMergeType_C",MatCompositeSetMergeType_Composite);
    1169:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetMatStructure_C",MatCompositeSetMatStructure_Composite);
    1170:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetMatStructure_C",MatCompositeGetMatStructure_Composite);
    1171:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeMerge_C",MatCompositeMerge_Composite);
    1172:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetNumberMat_C",MatCompositeGetNumberMat_Composite);
    1173:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetMat_C",MatCompositeGetMat_Composite);
    1174:   PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetScalings_C",MatCompositeSetScalings_Composite);

    1176:   PetscObjectChangeTypeName((PetscObject)A,MATCOMPOSITE);
    1177:   return(0);
    1178: }