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