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