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