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