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