xref: /petsc/src/mat/impls/composite/mcomposite.c (revision b6cfcaf5f2b3474f3cd87e956126e5b03ee5b3e1)
1 
2 #include <petsc/private/matimpl.h>        /*I "petscmat.h" I*/
3 
4 typedef struct _Mat_CompositeLink *Mat_CompositeLink;
5 struct _Mat_CompositeLink {
6   Mat               mat;
7   Vec               work;
8   Mat_CompositeLink next,prev;
9 };
10 
11 typedef struct {
12   MatCompositeType  type;
13   Mat_CompositeLink head,tail;
14   Vec               work;
15   PetscScalar       scale;        /* scale factor supplied with MatScale() */
16   Vec               left,right;   /* left and right diagonal scaling provided with MatDiagonalScale() */
17   Vec               leftwork,rightwork;
18   PetscInt          nmat;
19 } Mat_Composite;
20 
21 PetscErrorCode MatDestroy_Composite(Mat mat)
22 {
23   PetscErrorCode    ierr;
24   Mat_Composite     *shell = (Mat_Composite*)mat->data;
25   Mat_CompositeLink next   = shell->head,oldnext;
26 
27   PetscFunctionBegin;
28   while (next) {
29     ierr = MatDestroy(&next->mat);CHKERRQ(ierr);
30     if (next->work && (!next->next || next->work != next->next->work)) {
31       ierr = VecDestroy(&next->work);CHKERRQ(ierr);
32     }
33     oldnext = next;
34     next    = next->next;
35     ierr    = PetscFree(oldnext);CHKERRQ(ierr);
36   }
37   ierr = VecDestroy(&shell->work);CHKERRQ(ierr);
38   ierr = VecDestroy(&shell->left);CHKERRQ(ierr);
39   ierr = VecDestroy(&shell->right);CHKERRQ(ierr);
40   ierr = VecDestroy(&shell->leftwork);CHKERRQ(ierr);
41   ierr = VecDestroy(&shell->rightwork);CHKERRQ(ierr);
42   ierr = PetscFree(mat->data);CHKERRQ(ierr);
43   PetscFunctionReturn(0);
44 }
45 
46 PetscErrorCode MatMult_Composite_Multiplicative(Mat A,Vec x,Vec y)
47 {
48   Mat_Composite     *shell = (Mat_Composite*)A->data;
49   Mat_CompositeLink next   = shell->head;
50   PetscErrorCode    ierr;
51   Vec               in,out;
52 
53   PetscFunctionBegin;
54   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
55   in = x;
56   if (shell->right) {
57     if (!shell->rightwork) {
58       ierr = VecDuplicate(shell->right,&shell->rightwork);CHKERRQ(ierr);
59     }
60     ierr = VecPointwiseMult(shell->rightwork,shell->right,in);CHKERRQ(ierr);
61     in   = shell->rightwork;
62   }
63   while (next->next) {
64     if (!next->work) { /* should reuse previous work if the same size */
65       ierr = MatCreateVecs(next->mat,NULL,&next->work);CHKERRQ(ierr);
66     }
67     out  = next->work;
68     ierr = MatMult(next->mat,in,out);CHKERRQ(ierr);
69     in   = out;
70     next = next->next;
71   }
72   ierr = MatMult(next->mat,in,y);CHKERRQ(ierr);
73   if (shell->left) {
74     ierr = VecPointwiseMult(y,shell->left,y);CHKERRQ(ierr);
75   }
76   ierr = VecScale(y,shell->scale);CHKERRQ(ierr);
77   PetscFunctionReturn(0);
78 }
79 
80 PetscErrorCode MatMultTranspose_Composite_Multiplicative(Mat A,Vec x,Vec y)
81 {
82   Mat_Composite     *shell = (Mat_Composite*)A->data;
83   Mat_CompositeLink tail   = shell->tail;
84   PetscErrorCode    ierr;
85   Vec               in,out;
86 
87   PetscFunctionBegin;
88   if (!tail) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
89   in = x;
90   if (shell->left) {
91     if (!shell->leftwork) {
92       ierr = VecDuplicate(shell->left,&shell->leftwork);CHKERRQ(ierr);
93     }
94     ierr = VecPointwiseMult(shell->leftwork,shell->left,in);CHKERRQ(ierr);
95     in   = shell->leftwork;
96   }
97   while (tail->prev) {
98     if (!tail->prev->work) { /* should reuse previous work if the same size */
99       ierr = MatCreateVecs(tail->mat,NULL,&tail->prev->work);CHKERRQ(ierr);
100     }
101     out  = tail->prev->work;
102     ierr = MatMultTranspose(tail->mat,in,out);CHKERRQ(ierr);
103     in   = out;
104     tail = tail->prev;
105   }
106   ierr = MatMultTranspose(tail->mat,in,y);CHKERRQ(ierr);
107   if (shell->right) {
108     ierr = VecPointwiseMult(y,shell->right,y);CHKERRQ(ierr);
109   }
110   ierr = VecScale(y,shell->scale);CHKERRQ(ierr);
111   PetscFunctionReturn(0);
112 }
113 
114 PetscErrorCode MatMult_Composite(Mat A,Vec x,Vec y)
115 {
116   Mat_Composite     *shell = (Mat_Composite*)A->data;
117   Mat_CompositeLink next   = shell->head;
118   PetscErrorCode    ierr;
119   Vec               in;
120 
121   PetscFunctionBegin;
122   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
123   in = x;
124   if (shell->right) {
125     if (!shell->rightwork) {
126       ierr = VecDuplicate(shell->right,&shell->rightwork);CHKERRQ(ierr);
127     }
128     ierr = VecPointwiseMult(shell->rightwork,shell->right,in);CHKERRQ(ierr);
129     in   = shell->rightwork;
130   }
131   ierr = MatMult(next->mat,in,y);CHKERRQ(ierr);
132   while ((next = next->next)) {
133     ierr = MatMultAdd(next->mat,in,y,y);CHKERRQ(ierr);
134   }
135   if (shell->left) {
136     ierr = VecPointwiseMult(y,shell->left,y);CHKERRQ(ierr);
137   }
138   ierr = VecScale(y,shell->scale);CHKERRQ(ierr);
139   PetscFunctionReturn(0);
140 }
141 
142 PetscErrorCode MatMultTranspose_Composite(Mat A,Vec x,Vec y)
143 {
144   Mat_Composite     *shell = (Mat_Composite*)A->data;
145   Mat_CompositeLink next   = shell->head;
146   PetscErrorCode    ierr;
147   Vec               in;
148 
149   PetscFunctionBegin;
150   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
151   in = x;
152   if (shell->left) {
153     if (!shell->leftwork) {
154       ierr = VecDuplicate(shell->left,&shell->leftwork);CHKERRQ(ierr);
155     }
156     ierr = VecPointwiseMult(shell->leftwork,shell->left,in);CHKERRQ(ierr);
157     in   = shell->leftwork;
158   }
159   ierr = MatMultTranspose(next->mat,in,y);CHKERRQ(ierr);
160   while ((next = next->next)) {
161     ierr = MatMultTransposeAdd(next->mat,in,y,y);CHKERRQ(ierr);
162   }
163   if (shell->right) {
164     ierr = VecPointwiseMult(y,shell->right,y);CHKERRQ(ierr);
165   }
166   ierr = VecScale(y,shell->scale);CHKERRQ(ierr);
167   PetscFunctionReturn(0);
168 }
169 
170 PetscErrorCode MatGetDiagonal_Composite(Mat A,Vec v)
171 {
172   Mat_Composite     *shell = (Mat_Composite*)A->data;
173   Mat_CompositeLink next   = shell->head;
174   PetscErrorCode    ierr;
175 
176   PetscFunctionBegin;
177   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
178   if (shell->right || shell->left) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Cannot get diagonal if left or right scaling");
179 
180   ierr = MatGetDiagonal(next->mat,v);CHKERRQ(ierr);
181   if (next->next && !shell->work) {
182     ierr = VecDuplicate(v,&shell->work);CHKERRQ(ierr);
183   }
184   while ((next = next->next)) {
185     ierr = MatGetDiagonal(next->mat,shell->work);CHKERRQ(ierr);
186     ierr = VecAXPY(v,1.0,shell->work);CHKERRQ(ierr);
187   }
188   ierr = VecScale(v,shell->scale);CHKERRQ(ierr);
189   PetscFunctionReturn(0);
190 }
191 
192 PetscErrorCode MatAssemblyEnd_Composite(Mat Y,MatAssemblyType t)
193 {
194   PetscErrorCode ierr;
195   PetscBool      flg = PETSC_FALSE;
196 
197   PetscFunctionBegin;
198   ierr = PetscOptionsGetBool(((PetscObject)Y)->options,((PetscObject)Y)->prefix,"-mat_composite_merge",&flg,NULL);CHKERRQ(ierr);
199   if (flg) {
200     ierr = MatCompositeMerge(Y);CHKERRQ(ierr);
201   }
202   PetscFunctionReturn(0);
203 }
204 
205 PetscErrorCode MatScale_Composite(Mat inA,PetscScalar alpha)
206 {
207   Mat_Composite *a = (Mat_Composite*)inA->data;
208 
209   PetscFunctionBegin;
210   a->scale *= alpha;
211   PetscFunctionReturn(0);
212 }
213 
214 PetscErrorCode MatDiagonalScale_Composite(Mat inA,Vec left,Vec right)
215 {
216   Mat_Composite  *a = (Mat_Composite*)inA->data;
217   PetscErrorCode ierr;
218 
219   PetscFunctionBegin;
220   if (left) {
221     if (!a->left) {
222       ierr = VecDuplicate(left,&a->left);CHKERRQ(ierr);
223       ierr = VecCopy(left,a->left);CHKERRQ(ierr);
224     } else {
225       ierr = VecPointwiseMult(a->left,left,a->left);CHKERRQ(ierr);
226     }
227   }
228   if (right) {
229     if (!a->right) {
230       ierr = VecDuplicate(right,&a->right);CHKERRQ(ierr);
231       ierr = VecCopy(right,a->right);CHKERRQ(ierr);
232     } else {
233       ierr = VecPointwiseMult(a->right,right,a->right);CHKERRQ(ierr);
234     }
235   }
236   PetscFunctionReturn(0);
237 }
238 
239 static struct _MatOps MatOps_Values = {0,
240                                        0,
241                                        0,
242                                        MatMult_Composite,
243                                        0,
244                                 /*  5*/ MatMultTranspose_Composite,
245                                        0,
246                                        0,
247                                        0,
248                                        0,
249                                 /* 10*/ 0,
250                                        0,
251                                        0,
252                                        0,
253                                        0,
254                                 /* 15*/ 0,
255                                        0,
256                                        MatGetDiagonal_Composite,
257                                        MatDiagonalScale_Composite,
258                                        0,
259                                 /* 20*/ 0,
260                                        MatAssemblyEnd_Composite,
261                                        0,
262                                        0,
263                                /* 24*/ 0,
264                                        0,
265                                        0,
266                                        0,
267                                        0,
268                                /* 29*/ 0,
269                                        0,
270                                        0,
271                                        0,
272                                        0,
273                                /* 34*/ 0,
274                                        0,
275                                        0,
276                                        0,
277                                        0,
278                                /* 39*/ 0,
279                                        0,
280                                        0,
281                                        0,
282                                        0,
283                                /* 44*/ 0,
284                                        MatScale_Composite,
285                                        MatShift_Basic,
286                                        0,
287                                        0,
288                                /* 49*/ 0,
289                                        0,
290                                        0,
291                                        0,
292                                        0,
293                                /* 54*/ 0,
294                                        0,
295                                        0,
296                                        0,
297                                        0,
298                                /* 59*/ 0,
299                                        MatDestroy_Composite,
300                                        0,
301                                        0,
302                                        0,
303                                /* 64*/ 0,
304                                        0,
305                                        0,
306                                        0,
307                                        0,
308                                /* 69*/ 0,
309                                        0,
310                                        0,
311                                        0,
312                                        0,
313                                /* 74*/ 0,
314                                        0,
315                                        0,
316                                        0,
317                                        0,
318                                /* 79*/ 0,
319                                        0,
320                                        0,
321                                        0,
322                                        0,
323                                /* 84*/ 0,
324                                        0,
325                                        0,
326                                        0,
327                                        0,
328                                /* 89*/ 0,
329                                        0,
330                                        0,
331                                        0,
332                                        0,
333                                /* 94*/ 0,
334                                        0,
335                                        0,
336                                        0,
337                                        0,
338                                 /*99*/ 0,
339                                        0,
340                                        0,
341                                        0,
342                                        0,
343                                /*104*/ 0,
344                                        0,
345                                        0,
346                                        0,
347                                        0,
348                                /*109*/ 0,
349                                        0,
350                                        0,
351                                        0,
352                                        0,
353                                /*114*/ 0,
354                                        0,
355                                        0,
356                                        0,
357                                        0,
358                                /*119*/ 0,
359                                        0,
360                                        0,
361                                        0,
362                                        0,
363                                /*124*/ 0,
364                                        0,
365                                        0,
366                                        0,
367                                        0,
368                                /*129*/ 0,
369                                        0,
370                                        0,
371                                        0,
372                                        0,
373                                /*134*/ 0,
374                                        0,
375                                        0,
376                                        0,
377                                        0,
378                                /*139*/ 0,
379                                        0,
380                                        0
381 };
382 
383 /*MC
384    MATCOMPOSITE - A matrix defined by the sum (or product) of one or more matrices.
385     The matrices need to have a correct size and parallel layout for the sum or product to be valid.
386 
387    Notes:
388     to use the product of the matrices call MatCompositeSetType(mat,MAT_COMPOSITE_MULTIPLICATIVE);
389 
390   Level: advanced
391 
392 .seealso: MatCreateComposite(), MatCompositeAddMat(), MatSetType(), MatCompositeMerge(), MatCompositeSetType(), MatCompositeType
393 M*/
394 
395 PETSC_EXTERN PetscErrorCode MatCreate_Composite(Mat A)
396 {
397   Mat_Composite  *b;
398   PetscErrorCode ierr;
399 
400   PetscFunctionBegin;
401   ierr    = PetscNewLog(A,&b);CHKERRQ(ierr);
402   A->data = (void*)b;
403   ierr    = PetscMemcpy(A->ops,&MatOps_Values,sizeof(struct _MatOps));CHKERRQ(ierr);
404 
405   ierr = PetscLayoutSetUp(A->rmap);CHKERRQ(ierr);
406   ierr = PetscLayoutSetUp(A->cmap);CHKERRQ(ierr);
407 
408   A->assembled    = PETSC_TRUE;
409   A->preallocated = PETSC_TRUE;
410   b->type         = MAT_COMPOSITE_ADDITIVE;
411   b->scale        = 1.0;
412   b->nmat         = 0;
413   ierr            = PetscObjectChangeTypeName((PetscObject)A,MATCOMPOSITE);CHKERRQ(ierr);
414   PetscFunctionReturn(0);
415 }
416 
417 /*@
418    MatCreateComposite - Creates a matrix as the sum of zero or more matrices
419 
420   Collective on MPI_Comm
421 
422    Input Parameters:
423 +  comm - MPI communicator
424 .  nmat - number of matrices to put in
425 -  mats - the matrices
426 
427    Output Parameter:
428 .  mat - the matrix
429 
430    Level: advanced
431 
432    Notes:
433      Alternative construction
434 $       MatCreate(comm,&mat);
435 $       MatSetSizes(mat,m,n,M,N);
436 $       MatSetType(mat,MATCOMPOSITE);
437 $       MatCompositeAddMat(mat,mats[0]);
438 $       ....
439 $       MatCompositeAddMat(mat,mats[nmat-1]);
440 $       MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
441 $       MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);
442 
443      For the multiplicative form the product is mat[nmat-1]*mat[nmat-2]*....*mat[0]
444 
445 .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCompositeMerge(), MatCompositeSetType(), MatCompositeType
446 
447 @*/
448 PetscErrorCode  MatCreateComposite(MPI_Comm comm,PetscInt nmat,const Mat *mats,Mat *mat)
449 {
450   PetscErrorCode ierr;
451   PetscInt       m,n,M,N,i;
452 
453   PetscFunctionBegin;
454   if (nmat < 1) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Must pass in at least one matrix");
455   PetscValidPointer(mat,3);
456 
457   ierr = MatGetLocalSize(mats[0],PETSC_IGNORE,&n);CHKERRQ(ierr);
458   ierr = MatGetLocalSize(mats[nmat-1],&m,PETSC_IGNORE);CHKERRQ(ierr);
459   ierr = MatGetSize(mats[0],PETSC_IGNORE,&N);CHKERRQ(ierr);
460   ierr = MatGetSize(mats[nmat-1],&M,PETSC_IGNORE);CHKERRQ(ierr);
461   ierr = MatCreate(comm,mat);CHKERRQ(ierr);
462   ierr = MatSetSizes(*mat,m,n,M,N);CHKERRQ(ierr);
463   ierr = MatSetType(*mat,MATCOMPOSITE);CHKERRQ(ierr);
464   for (i=0; i<nmat; i++) {
465     ierr = MatCompositeAddMat(*mat,mats[i]);CHKERRQ(ierr);
466   }
467   ierr = MatAssemblyBegin(*mat,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
468   ierr = MatAssemblyEnd(*mat,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
469   PetscFunctionReturn(0);
470 }
471 
472 /*@
473     MatCompositeAddMat - add another matrix to a composite matrix
474 
475    Collective on Mat
476 
477     Input Parameters:
478 +   mat - the composite matrix
479 -   smat - the partial matrix
480 
481    Level: advanced
482 
483 .seealso: MatCreateComposite()
484 @*/
485 PetscErrorCode  MatCompositeAddMat(Mat mat,Mat smat)
486 {
487   Mat_Composite     *shell;
488   PetscErrorCode    ierr;
489   Mat_CompositeLink ilink,next;
490 
491   PetscFunctionBegin;
492   PetscValidHeaderSpecific(mat,MAT_CLASSID,1);
493   PetscValidHeaderSpecific(smat,MAT_CLASSID,2);
494   ierr        = PetscNewLog(mat,&ilink);CHKERRQ(ierr);
495   ilink->next = 0;
496   ierr        = PetscObjectReference((PetscObject)smat);CHKERRQ(ierr);
497   ilink->mat  = smat;
498 
499   shell = (Mat_Composite*)mat->data;
500   next  = shell->head;
501   if (!next) shell->head = ilink;
502   else {
503     while (next->next) {
504       next = next->next;
505     }
506     next->next  = ilink;
507     ilink->prev = next;
508   }
509   shell->tail =  ilink;
510   shell->nmat += 1;
511   PetscFunctionReturn(0);
512 }
513 
514 /*@
515    MatCompositeSetType - Indicates if the matrix is defined as the sum of a set of matrices or the product
516 
517   Collective on MPI_Comm
518 
519    Input Parameters:
520 .  mat - the composite matrix
521 
522 
523    Level: advanced
524 
525    Notes:
526       The MatType of the resulting matrix will be the same as the MatType of the FIRST
527     matrix in the composite matrix.
528 
529 .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCreateComposite(), MATCOMPOSITE
530 
531 @*/
532 PetscErrorCode  MatCompositeSetType(Mat mat,MatCompositeType type)
533 {
534   Mat_Composite  *b = (Mat_Composite*)mat->data;
535   PetscBool      flg;
536   PetscErrorCode ierr;
537 
538   PetscFunctionBegin;
539   ierr = PetscObjectTypeCompare((PetscObject)mat,MATCOMPOSITE,&flg);CHKERRQ(ierr);
540   if (!flg) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Can only use with composite matrix");
541   if (type == MAT_COMPOSITE_MULTIPLICATIVE) {
542     mat->ops->getdiagonal   = 0;
543     mat->ops->mult          = MatMult_Composite_Multiplicative;
544     mat->ops->multtranspose = MatMultTranspose_Composite_Multiplicative;
545     b->type                 = MAT_COMPOSITE_MULTIPLICATIVE;
546   } else {
547     mat->ops->getdiagonal   = MatGetDiagonal_Composite;
548     mat->ops->mult          = MatMult_Composite;
549     mat->ops->multtranspose = MatMultTranspose_Composite;
550     b->type                 = MAT_COMPOSITE_ADDITIVE;
551   }
552   PetscFunctionReturn(0);
553 }
554 
555 
556 /*@
557    MatCompositeMerge - Given a composite matrix, replaces it with a "regular" matrix
558      by summing all the matrices inside the composite matrix.
559 
560   Collective on MPI_Comm
561 
562    Input Parameters:
563 .  mat - the composite matrix
564 
565 
566    Options Database:
567 .  -mat_composite_merge  (you must call MatAssemblyBegin()/MatAssemblyEnd() to have this checked)
568 
569    Level: advanced
570 
571    Notes:
572       The MatType of the resulting matrix will be the same as the MatType of the FIRST
573     matrix in the composite matrix.
574 
575 .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCreateComposite(), MATCOMPOSITE
576 
577 @*/
578 PetscErrorCode  MatCompositeMerge(Mat mat)
579 {
580   Mat_Composite     *shell = (Mat_Composite*)mat->data;
581   Mat_CompositeLink next   = shell->head, prev = shell->tail;
582   PetscErrorCode    ierr;
583   Mat               tmat,newmat;
584   Vec               left,right;
585   PetscScalar       scale;
586 
587   PetscFunctionBegin;
588   if (!next) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
589 
590   PetscFunctionBegin;
591   if (shell->type == MAT_COMPOSITE_ADDITIVE) {
592     ierr = MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);CHKERRQ(ierr);
593     while ((next = next->next)) {
594       ierr = MatAXPY(tmat,1.0,next->mat,DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr);
595     }
596   } else {
597     ierr = MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);CHKERRQ(ierr);
598     while ((next = next->next)) {
599       ierr = MatMatMult(next->mat,tmat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat);CHKERRQ(ierr);
600       ierr = MatDestroy(&tmat);CHKERRQ(ierr);
601       tmat = newmat;
602     }
603   }
604 
605   scale = shell->scale;
606   if ((left = shell->left)) {ierr = PetscObjectReference((PetscObject)left);CHKERRQ(ierr);}
607   if ((right = shell->right)) {ierr = PetscObjectReference((PetscObject)right);CHKERRQ(ierr);}
608 
609   ierr = MatHeaderReplace(mat,&tmat);CHKERRQ(ierr);
610 
611   ierr = MatDiagonalScale(mat,left,right);CHKERRQ(ierr);
612   ierr = MatScale(mat,scale);CHKERRQ(ierr);
613   ierr = VecDestroy(&left);CHKERRQ(ierr);
614   ierr = VecDestroy(&right);CHKERRQ(ierr);
615   PetscFunctionReturn(0);
616 }
617 
618 /*@
619    MatCompositeGetNMat - Returns the number of matrices in composite.
620 
621    Not Collective
622 
623    Input Parameter:
624 .  A - the composite matrix
625 
626    Output Parameter:
627 .  size - the local size
628 .  nmat - number of matrices in composite
629 
630    Level: advanced
631 
632 .seealso: MatCreateComposite(), MatCompositeGetMat()
633 
634 @*/
635 PetscErrorCode MatCompositeGetNMat(Mat A,PetscInt *nmat)
636 {
637   Mat_Composite  *shell;
638 
639   PetscFunctionBegin;
640   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
641   PetscValidPointer(nmat,2);
642   shell = (Mat_Composite*)A->data;
643   *nmat = shell->nmat;
644   PetscFunctionReturn(0);
645 }
646 
647 /*@
648    MatCompositeGetMat - Returns the ith matrix from composite.
649 
650    Logically Collective on Mat
651 
652    Input Parameter:
653 +  A - the composite matrix
654 -  i - the number of requested matrix
655 
656    Output Parameter:
657 .  Ai - ith matrix in composite
658 
659    Level: advanced
660 
661 .seealso: MatCreateComposite(), MatCompositeGetNMat()
662 
663 @*/
664 PetscErrorCode MatCompositeGetMat(Mat A,PetscInt i,Mat *Ai)
665 {
666   Mat_Composite     *shell;
667   Mat_CompositeLink ilink;
668   PetscInt          k;
669 
670   PetscFunctionBegin;
671   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
672   PetscValidLogicalCollectiveInt(A,i,2);
673   PetscValidPointer(Ai,3);
674   shell = (Mat_Composite*)A->data;
675   if (i >= shell->nmat) SETERRQ2(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_OUTOFRANGE,"index out of range: %d >= %d",i,shell->nmat);
676   ilink = shell->head;
677   for (k=0; k<i; k++) {
678     if (ilink) {
679       ilink = ilink->next;
680     } else {
681       break;
682     }
683   }
684   *Ai = ilink->mat;
685   PetscFunctionReturn(0);
686 }
687 
688