xref: /petsc/src/mat/impls/composite/mcomposite.c (revision c4e433427e906e3f2dbd70553bb6d0901adfc1f1)
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_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_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_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_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_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
191   if (shell->right || shell->left) SETERRQ(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   PetscTruth     flg;
211 
212   PetscFunctionBegin;
213   ierr = PetscOptionsHasName(((PetscObject)Y)->prefix,"-mat_composite_merge",&flg);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        0,
283 /*25*/ 0,
284        0,
285        0,
286        0,
287        0,
288 /*30*/ 0,
289        0,
290        0,
291        0,
292        0,
293 /*35*/ 0,
294        0,
295        0,
296        0,
297        0,
298 /*40*/ 0,
299        0,
300        0,
301        0,
302        0,
303 /*45*/ 0,
304        MatScale_Composite,
305        0,
306        0,
307        0,
308 /*50*/ 0,
309        0,
310        0,
311        0,
312        0,
313 /*55*/ 0,
314        0,
315        0,
316        0,
317        0,
318 /*60*/ 0,
319        MatDestroy_Composite,
320        0,
321        0,
322        0,
323 /*65*/ 0,
324        0,
325        0,
326        0,
327        0,
328 /*70*/ 0,
329        0,
330        0,
331        0,
332        0,
333 /*75*/ 0,
334        0,
335        0,
336        0,
337        0,
338 /*80*/ 0,
339        0,
340        0,
341        0,
342        0,
343 /*85*/ 0,
344        0,
345        0,
346        0,
347        0,
348 /*90*/ 0,
349        0,
350        0,
351        0,
352        0,
353 /*95*/ 0,
354        0,
355        0,
356        0};
357 
358 /*MC
359    MATCOMPOSITE - A matrix defined by the sum (or product) of one or more matrices (all matrices are of same size and parallel layout).
360 
361    Notes: to use the product of the matrices call MatCompositeSetType(mat,MAT_COMPOSITE_MULTIPLICATIVE);
362 
363   Level: advanced
364 
365 .seealso: MatCreateComposite(), MatCompositeAddMat(), MatSetType(), MatCompositeMerge(), MatCompositeSetType(), MatCompositeType
366 M*/
367 
368 EXTERN_C_BEGIN
369 #undef __FUNCT__
370 #define __FUNCT__ "MatCreate_Composite"
371 PetscErrorCode PETSCMAT_DLLEXPORT MatCreate_Composite(Mat A)
372 {
373   Mat_Composite  *b;
374   PetscErrorCode ierr;
375 
376   PetscFunctionBegin;
377   ierr = PetscNewLog(A,Mat_Composite,&b);CHKERRQ(ierr);
378   A->data = (void*)b;
379   ierr = PetscMemcpy(A->ops,&MatOps_Values,sizeof(struct _MatOps));CHKERRQ(ierr);
380 
381   ierr = PetscMapSetBlockSize(A->rmap,1);CHKERRQ(ierr);
382   ierr = PetscMapSetBlockSize(A->cmap,1);CHKERRQ(ierr);
383   ierr = PetscMapSetUp(A->rmap);CHKERRQ(ierr);
384   ierr = PetscMapSetUp(A->cmap);CHKERRQ(ierr);
385 
386   A->assembled     = PETSC_TRUE;
387   A->preallocated  = PETSC_FALSE;
388   b->type          = MAT_COMPOSITE_ADDITIVE;
389   b->scale         = 1.0;
390   ierr = PetscObjectChangeTypeName((PetscObject)A,MATCOMPOSITE);CHKERRQ(ierr);
391   PetscFunctionReturn(0);
392 }
393 EXTERN_C_END
394 
395 #undef __FUNCT__
396 #define __FUNCT__ "MatCreateComposite"
397 /*@C
398    MatCreateComposite - Creates a matrix as the sum of zero or more matrices
399 
400   Collective on MPI_Comm
401 
402    Input Parameters:
403 +  comm - MPI communicator
404 .  nmat - number of matrices to put in
405 -  mats - the matrices
406 
407    Output Parameter:
408 .  A - the matrix
409 
410    Level: advanced
411 
412    Notes:
413      Alternative construction
414 $       MatCreate(comm,&mat);
415 $       MatSetSizes(mat,m,n,M,N);
416 $       MatSetType(mat,MATCOMPOSITE);
417 $       MatCompositeAddMat(mat,mats[0]);
418 $       ....
419 $       MatCompositeAddMat(mat,mats[nmat-1]);
420 $       MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
421 $       MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);
422 
423      For the multiplicative form the product is mat[nmat-1]*mat[nmat-2]*....*mat[0]
424 
425 .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCompositeMerge(), MatCompositeSetType(), MatCompositeType
426 
427 @*/
428 PetscErrorCode PETSCMAT_DLLEXPORT MatCreateComposite(MPI_Comm comm,PetscInt nmat,const Mat *mats,Mat *mat)
429 {
430   PetscErrorCode ierr;
431   PetscInt       m,n,M,N,i;
432 
433   PetscFunctionBegin;
434   if (nmat < 1) SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"Must pass in at least one matrix");
435   PetscValidPointer(mat,3);
436 
437   ierr = MatGetLocalSize(mats[0],&m,&n);CHKERRQ(ierr);
438   ierr = MatGetSize(mats[0],&M,&N);CHKERRQ(ierr);
439   ierr = MatCreate(comm,mat);CHKERRQ(ierr);
440   ierr = MatSetSizes(*mat,m,n,M,N);CHKERRQ(ierr);
441   ierr = MatSetType(*mat,MATCOMPOSITE);CHKERRQ(ierr);
442   for (i=0; i<nmat; i++) {
443     ierr = MatCompositeAddMat(*mat,mats[i]);CHKERRQ(ierr);
444   }
445   ierr = MatAssemblyBegin(*mat,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
446   ierr = MatAssemblyEnd(*mat,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
447   PetscFunctionReturn(0);
448 }
449 
450 #undef __FUNCT__
451 #define __FUNCT__ "MatCompositeAddMat"
452 /*@
453     MatCompositeAddMat - add another matrix to a composite matrix
454 
455    Collective on Mat
456 
457     Input Parameters:
458 +   mat - the composite matrix
459 -   smat - the partial matrix
460 
461    Level: advanced
462 
463 .seealso: MatCreateComposite()
464 @*/
465 PetscErrorCode PETSCMAT_DLLEXPORT MatCompositeAddMat(Mat mat,Mat smat)
466 {
467   Mat_Composite     *shell;
468   PetscErrorCode    ierr;
469   Mat_CompositeLink ilink,next;
470 
471   PetscFunctionBegin;
472   PetscValidHeaderSpecific(mat,MAT_COOKIE,1);
473   PetscValidHeaderSpecific(smat,MAT_COOKIE,2);
474   ierr        = PetscNewLog(mat,struct _Mat_CompositeLink,&ilink);CHKERRQ(ierr);
475   ilink->next = 0;
476   ierr        = PetscObjectReference((PetscObject)smat);CHKERRQ(ierr);
477   ilink->mat  = smat;
478 
479   shell = (Mat_Composite*)mat->data;
480   next = shell->head;
481   if (!next) {
482     shell->head  = ilink;
483   } else {
484     while (next->next) {
485       next = next->next;
486     }
487     next->next      = ilink;
488     ilink->prev     = next;
489   }
490   shell->tail = ilink;
491   PetscFunctionReturn(0);
492 }
493 
494 #undef __FUNCT__
495 #define __FUNCT__ "MatCompositeSetType"
496 /*@C
497    MatCompositeSetType - Indicates if the matrix is defined as the sum of a set of matrices or the product
498 
499   Collective on MPI_Comm
500 
501    Input Parameters:
502 .  mat - the composite matrix
503 
504 
505    Level: advanced
506 
507    Notes:
508       The MatType of the resulting matrix will be the same as the MatType of the FIRST
509     matrix in the composite matrix.
510 
511 .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCreateComposite(), MATCOMPOSITE
512 
513 @*/
514 PetscErrorCode PETSCMAT_DLLEXPORT MatCompositeSetType(Mat mat,MatCompositeType type)
515 {
516   Mat_Composite  *b = (Mat_Composite*)mat->data;
517   PetscTruth     flg;
518   PetscErrorCode ierr;
519 
520   PetscFunctionBegin;
521   ierr = PetscTypeCompare((PetscObject)mat,MATCOMPOSITE,&flg);CHKERRQ(ierr);
522   if (!flg) SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Can only use with composite matrix");
523   if (type == MAT_COMPOSITE_MULTIPLICATIVE) {
524     mat->ops->getdiagonal   = 0;
525     mat->ops->mult          = MatMult_Composite_Multiplicative;
526     mat->ops->multtranspose = MatMultTranspose_Composite_Multiplicative;
527     b->type                 = MAT_COMPOSITE_MULTIPLICATIVE;
528   } else {
529     mat->ops->getdiagonal   = MatGetDiagonal_Composite;
530     mat->ops->mult          = MatMult_Composite;
531     mat->ops->multtranspose = MatMultTranspose_Composite;
532     b->type                 = MAT_COMPOSITE_ADDITIVE;
533   }
534   PetscFunctionReturn(0);
535 }
536 
537 
538 #undef __FUNCT__
539 #define __FUNCT__ "MatCompositeMerge"
540 /*@C
541    MatCompositeMerge - Given a composite matrix, replaces it with a "regular" matrix
542      by summing all the matrices inside the composite matrix.
543 
544   Collective on MPI_Comm
545 
546    Input Parameters:
547 .  mat - the composite matrix
548 
549 
550    Options Database:
551 .  -mat_composite_merge  (you must call MatAssemblyBegin()/MatAssemblyEnd() to have this checked)
552 
553    Level: advanced
554 
555    Notes:
556       The MatType of the resulting matrix will be the same as the MatType of the FIRST
557     matrix in the composite matrix.
558 
559 .seealso: MatDestroy(), MatMult(), MatCompositeAddMat(), MatCreateComposite(), MATCOMPOSITE
560 
561 @*/
562 PetscErrorCode PETSCMAT_DLLEXPORT MatCompositeMerge(Mat mat)
563 {
564   Mat_Composite     *shell = (Mat_Composite*)mat->data;
565   Mat_CompositeLink next = shell->head, prev = shell->tail;
566   PetscErrorCode    ierr;
567   Mat               tmat,newmat;
568 
569   PetscFunctionBegin;
570   if (!next) SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()");
571 
572   PetscFunctionBegin;
573   if (shell->type == MAT_COMPOSITE_ADDITIVE) {
574     ierr = MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);CHKERRQ(ierr);
575     while ((next = next->next)) {
576       ierr = MatAXPY(tmat,1.0,next->mat,DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr);
577     }
578   } else {
579     ierr = MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat);CHKERRQ(ierr);
580     while ((prev = prev->prev)) {
581       ierr = MatMatMult(tmat,prev->mat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat);CHKERRQ(ierr);
582       ierr = MatDestroy(tmat);CHKERRQ(ierr);
583       tmat = newmat;
584     }
585   }
586   ierr = MatHeaderReplace(mat,tmat);CHKERRQ(ierr);
587   ierr = MatDiagonalScale(mat,shell->left,shell->right);CHKERRQ(ierr);
588   ierr = MatScale(mat,shell->scale);CHKERRQ(ierr);
589   PetscFunctionReturn(0);
590 }
591