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