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