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