xref: /petsc/src/mat/tutorials/ex8.c (revision 9566063d113dddea24716c546802770db7481bc0) 
1c4762a1bSJed Brown 
2c4762a1bSJed Brown static char help[] = "Shows how to add a new MatOperation to AIJ MatType\n\n";
3c4762a1bSJed Brown 
4c4762a1bSJed Brown #include <petscmat.h>
5c4762a1bSJed Brown #include <petscblaslapack.h>
6c4762a1bSJed Brown 
7c4762a1bSJed Brown static PetscErrorCode MatScaleUserImpl_SeqAIJ(Mat inA,PetscScalar alpha)
8c4762a1bSJed Brown {
9c4762a1bSJed Brown   PetscFunctionBegin;
10*9566063dSJacob Faibussowitsch   PetscCall(MatScale(inA,alpha));
11c4762a1bSJed Brown   PetscFunctionReturn(0);
12c4762a1bSJed Brown }
13c4762a1bSJed Brown 
14c4762a1bSJed Brown extern PetscErrorCode MatScaleUserImpl(Mat,PetscScalar);
15c4762a1bSJed Brown 
16c4762a1bSJed Brown static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A,PetscScalar aa)
17c4762a1bSJed Brown {
18c4762a1bSJed Brown   Mat            AA,AB;
19c4762a1bSJed Brown 
20c4762a1bSJed Brown   PetscFunctionBegin;
21*9566063dSJacob Faibussowitsch   PetscCall(MatMPIAIJGetSeqAIJ(A,&AA,&AB,NULL));
22*9566063dSJacob Faibussowitsch   PetscCall(MatScaleUserImpl(AA,aa));
23*9566063dSJacob Faibussowitsch   PetscCall(MatScaleUserImpl(AB,aa));
24c4762a1bSJed Brown   PetscFunctionReturn(0);
25c4762a1bSJed Brown }
26c4762a1bSJed Brown 
27c4762a1bSJed Brown /* This routine registers MatScaleUserImpl_SeqAIJ() and
28c4762a1bSJed Brown    MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl()
29c4762a1bSJed Brown    functionality for SeqAIJ and MPIAIJ matrix-types */
30c4762a1bSJed Brown PetscErrorCode RegisterMatScaleUserImpl(Mat mat)
31c4762a1bSJed Brown {
32c4762a1bSJed Brown   PetscMPIInt    size;
33c4762a1bSJed Brown 
34362febeeSStefano Zampini   PetscFunctionBegin;
35*9566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
36c4762a1bSJed Brown   if (size == 1) { /* SeqAIJ Matrix */
37*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ));
38c4762a1bSJed Brown   } else { /* MPIAIJ Matrix */
39c4762a1bSJed Brown     Mat AA,AB;
40*9566063dSJacob Faibussowitsch     PetscCall(MatMPIAIJGetSeqAIJ(mat,&AA,&AB,NULL));
41*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_MPIAIJ));
42*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)AA,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ));
43*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)AB,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ));
44c4762a1bSJed Brown   }
45c4762a1bSJed Brown   PetscFunctionReturn(0);
46c4762a1bSJed Brown }
47c4762a1bSJed Brown 
48c4762a1bSJed Brown /* this routines queries the already registered MatScaleUserImp_XXX
49c4762a1bSJed Brown    implementations for the given matrix, and calls the correct
50c4762a1bSJed Brown    routine. i.e if MatType is SeqAIJ, MatScaleUserImpl_SeqAIJ() gets
51c4762a1bSJed Brown    called, and if MatType is MPIAIJ, MatScaleUserImpl_MPIAIJ() gets
52c4762a1bSJed Brown    called */
53c4762a1bSJed Brown PetscErrorCode MatScaleUserImpl(Mat mat,PetscScalar a)
54c4762a1bSJed Brown {
555f80ce2aSJacob Faibussowitsch   PetscErrorCode (*f)(Mat,PetscScalar);
56c4762a1bSJed Brown 
57c4762a1bSJed Brown   PetscFunctionBegin;
58*9566063dSJacob Faibussowitsch   PetscCall(PetscObjectQueryFunction((PetscObject)mat,"MatScaleUserImpl_C",&f));
59*9566063dSJacob Faibussowitsch   if (f) PetscCall((*f)(mat,a));
60c4762a1bSJed Brown   PetscFunctionReturn(0);
61c4762a1bSJed Brown }
62c4762a1bSJed Brown 
63c4762a1bSJed Brown /* Main user code that uses MatScaleUserImpl() */
64c4762a1bSJed Brown 
65c4762a1bSJed Brown int main(int argc,char **args)
66c4762a1bSJed Brown {
67c4762a1bSJed Brown   Mat            mat;
68c4762a1bSJed Brown   PetscInt       i,j,m = 2,n,Ii,J;
69c4762a1bSJed Brown   PetscScalar    v,none = -1.0;
70c4762a1bSJed Brown   PetscMPIInt    rank,size;
71c4762a1bSJed Brown 
72*9566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc,&args,(char*)0,help));
73*9566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&rank));
74*9566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
75c4762a1bSJed Brown   n    = 2*size;
76c4762a1bSJed Brown 
77c4762a1bSJed Brown   /* create the matrix */
78*9566063dSJacob Faibussowitsch   PetscCall(MatCreate(PETSC_COMM_WORLD,&mat));
79*9566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(mat,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n));
80*9566063dSJacob Faibussowitsch   PetscCall(MatSetType(mat,MATAIJ));
81*9566063dSJacob Faibussowitsch   PetscCall(MatSetUp(mat));
82c4762a1bSJed Brown 
83c4762a1bSJed Brown   /* register user defined MatScaleUser() operation for both SeqAIJ
84c4762a1bSJed Brown      and MPIAIJ types */
85*9566063dSJacob Faibussowitsch   PetscCall(RegisterMatScaleUserImpl(mat));
86c4762a1bSJed Brown 
87c4762a1bSJed Brown   /* assemble the matrix */
88c4762a1bSJed Brown   for (i=0; i<m; i++) {
89c4762a1bSJed Brown     for (j=2*rank; j<2*rank+2; j++) {
90c4762a1bSJed Brown       v = -1.0;  Ii = j + n*i;
91*9566063dSJacob Faibussowitsch       if (i>0)   {J = Ii - n; PetscCall(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
92*9566063dSJacob Faibussowitsch       if (i<m-1) {J = Ii + n; PetscCall(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
93*9566063dSJacob Faibussowitsch       if (j>0)   {J = Ii - 1; PetscCall(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
94*9566063dSJacob Faibussowitsch       if (j<n-1) {J = Ii + 1; PetscCall(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
95*9566063dSJacob Faibussowitsch       v = 4.0; PetscCall(MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES));
96c4762a1bSJed Brown     }
97c4762a1bSJed Brown   }
98*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY));
99*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY));
100c4762a1bSJed Brown 
101c4762a1bSJed Brown   /* check the matrix before and after scaling by -1.0 */
102*9566063dSJacob Faibussowitsch   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"Matrix _before_ MatScaleUserImpl() operation\n"));
103*9566063dSJacob Faibussowitsch   PetscCall(MatView(mat,PETSC_VIEWER_STDOUT_WORLD));
104*9566063dSJacob Faibussowitsch   PetscCall(MatScaleUserImpl(mat,none));
105*9566063dSJacob Faibussowitsch   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"Matrix _after_ MatScaleUserImpl() operation\n"));
106*9566063dSJacob Faibussowitsch   PetscCall(MatView(mat,PETSC_VIEWER_STDOUT_WORLD));
107c4762a1bSJed Brown 
108*9566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&mat));
109*9566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
110b122ec5aSJacob Faibussowitsch   return 0;
111c4762a1bSJed Brown }
112c4762a1bSJed Brown 
113c4762a1bSJed Brown /*TEST
114c4762a1bSJed Brown 
115c4762a1bSJed Brown    test:
116c4762a1bSJed Brown 
117c4762a1bSJed Brown    test:
118c4762a1bSJed Brown       suffix: 2
119c4762a1bSJed Brown       nsize: 2
120c4762a1bSJed Brown 
121c4762a1bSJed Brown TEST*/
122