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