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