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