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