xref: /petsc/src/mat/tutorials/ex8.c (revision 5f80ce2ab25dff0f4601e710601cbbcecf323266) 
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*5f80ce2aSJacob Faibussowitsch   CHKERRQ(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*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatMPIAIJGetSeqAIJ(A,&AA,&AB,NULL));
22*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatScaleUserImpl(AA,aa));
23*5f80ce2aSJacob Faibussowitsch   CHKERRQ(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*5f80ce2aSJacob Faibussowitsch   CHKERRMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
36c4762a1bSJed Brown   if (size == 1) { /* SeqAIJ Matrix */
37*5f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ));
38c4762a1bSJed Brown   } else { /* MPIAIJ Matrix */
39c4762a1bSJed Brown     Mat AA,AB;
40*5f80ce2aSJacob Faibussowitsch     CHKERRQ(MatMPIAIJGetSeqAIJ(mat,&AA,&AB,NULL));
41*5f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_MPIAIJ));
42*5f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscObjectComposeFunction((PetscObject)AA,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ));
43*5f80ce2aSJacob Faibussowitsch     CHKERRQ(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 {
55*5f80ce2aSJacob Faibussowitsch   PetscErrorCode (*f)(Mat,PetscScalar);
56c4762a1bSJed Brown 
57c4762a1bSJed Brown   PetscFunctionBegin;
58*5f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscObjectQueryFunction((PetscObject)mat,"MatScaleUserImpl_C",&f));
59*5f80ce2aSJacob Faibussowitsch   if (f) CHKERRQ((*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   PetscErrorCode ierr;
70c4762a1bSJed Brown   PetscScalar    v,none = -1.0;
71c4762a1bSJed Brown   PetscMPIInt    rank,size;
72c4762a1bSJed Brown 
73c4762a1bSJed Brown   ierr = PetscInitialize(&argc,&args,(char*)0,help);if (ierr) return ierr;
74*5f80ce2aSJacob Faibussowitsch   CHKERRMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&rank));
75*5f80ce2aSJacob Faibussowitsch   CHKERRMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
76c4762a1bSJed Brown   n    = 2*size;
77c4762a1bSJed Brown 
78c4762a1bSJed Brown   /* create the matrix */
79*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatCreate(PETSC_COMM_WORLD,&mat));
80*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatSetSizes(mat,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n));
81*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatSetType(mat,MATAIJ));
82*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatSetUp(mat));
83c4762a1bSJed Brown 
84c4762a1bSJed Brown   /* register user defined MatScaleUser() operation for both SeqAIJ
85c4762a1bSJed Brown      and MPIAIJ types */
86*5f80ce2aSJacob Faibussowitsch   CHKERRQ(RegisterMatScaleUserImpl(mat));
87c4762a1bSJed Brown 
88c4762a1bSJed Brown   /* assemble the matrix */
89c4762a1bSJed Brown   for (i=0; i<m; i++) {
90c4762a1bSJed Brown     for (j=2*rank; j<2*rank+2; j++) {
91c4762a1bSJed Brown       v = -1.0;  Ii = j + n*i;
92*5f80ce2aSJacob Faibussowitsch       if (i>0)   {J = Ii - n; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
93*5f80ce2aSJacob Faibussowitsch       if (i<m-1) {J = Ii + n; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
94*5f80ce2aSJacob Faibussowitsch       if (j>0)   {J = Ii - 1; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
95*5f80ce2aSJacob Faibussowitsch       if (j<n-1) {J = Ii + 1; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));}
96*5f80ce2aSJacob Faibussowitsch       v = 4.0; CHKERRQ(MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES));
97c4762a1bSJed Brown     }
98c4762a1bSJed Brown   }
99*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY));
100*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY));
101c4762a1bSJed Brown 
102c4762a1bSJed Brown   /* check the matrix before and after scaling by -1.0 */
103*5f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Matrix _before_ MatScaleUserImpl() operation\n"));
104*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatView(mat,PETSC_VIEWER_STDOUT_WORLD));
105*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatScaleUserImpl(mat,none));
106*5f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Matrix _after_ MatScaleUserImpl() operation\n"));
107*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatView(mat,PETSC_VIEWER_STDOUT_WORLD));
108c4762a1bSJed Brown 
109*5f80ce2aSJacob Faibussowitsch   CHKERRQ(MatDestroy(&mat));
110c4762a1bSJed Brown   ierr = PetscFinalize();
111c4762a1bSJed Brown   return ierr;
112c4762a1bSJed Brown }
113c4762a1bSJed Brown 
114c4762a1bSJed Brown /*TEST
115c4762a1bSJed Brown 
116c4762a1bSJed Brown    test:
117c4762a1bSJed Brown 
118c4762a1bSJed Brown    test:
119c4762a1bSJed Brown       suffix: 2
120c4762a1bSJed Brown       nsize: 2
121c4762a1bSJed Brown 
122c4762a1bSJed Brown TEST*/
123