1c4762a1bSJed Brown 2c4762a1bSJed Brown static char help[] = "Shows how to add a new MatOperation to AIJ MatType\n\n"; 3c4762a1bSJed Brown 4c4762a1bSJed Brown #include <petscmat.h> 5c4762a1bSJed Brown #include <petscblaslapack.h> 6c4762a1bSJed Brown 7c4762a1bSJed Brown static PetscErrorCode MatScaleUserImpl_SeqAIJ(Mat inA,PetscScalar alpha) 8c4762a1bSJed Brown { 9c4762a1bSJed Brown PetscFunctionBegin; 105f80ce2aSJacob Faibussowitsch CHKERRQ(MatScale(inA,alpha)); 11c4762a1bSJed Brown PetscFunctionReturn(0); 12c4762a1bSJed Brown } 13c4762a1bSJed Brown 14c4762a1bSJed Brown extern PetscErrorCode MatScaleUserImpl(Mat,PetscScalar); 15c4762a1bSJed Brown 16c4762a1bSJed Brown static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A,PetscScalar aa) 17c4762a1bSJed Brown { 18c4762a1bSJed Brown Mat AA,AB; 19c4762a1bSJed Brown 20c4762a1bSJed Brown PetscFunctionBegin; 215f80ce2aSJacob Faibussowitsch CHKERRQ(MatMPIAIJGetSeqAIJ(A,&AA,&AB,NULL)); 225f80ce2aSJacob Faibussowitsch CHKERRQ(MatScaleUserImpl(AA,aa)); 235f80ce2aSJacob Faibussowitsch CHKERRQ(MatScaleUserImpl(AB,aa)); 24c4762a1bSJed Brown PetscFunctionReturn(0); 25c4762a1bSJed Brown } 26c4762a1bSJed Brown 27c4762a1bSJed Brown /* This routine registers MatScaleUserImpl_SeqAIJ() and 28c4762a1bSJed Brown MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl() 29c4762a1bSJed Brown functionality for SeqAIJ and MPIAIJ matrix-types */ 30c4762a1bSJed Brown PetscErrorCode RegisterMatScaleUserImpl(Mat mat) 31c4762a1bSJed Brown { 32c4762a1bSJed Brown PetscMPIInt size; 33c4762a1bSJed Brown 34362febeeSStefano Zampini PetscFunctionBegin; 355f80ce2aSJacob Faibussowitsch CHKERRMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size)); 36c4762a1bSJed Brown if (size == 1) { /* SeqAIJ Matrix */ 375f80ce2aSJacob Faibussowitsch CHKERRQ(PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ)); 38c4762a1bSJed Brown } else { /* MPIAIJ Matrix */ 39c4762a1bSJed Brown Mat AA,AB; 405f80ce2aSJacob Faibussowitsch CHKERRQ(MatMPIAIJGetSeqAIJ(mat,&AA,&AB,NULL)); 415f80ce2aSJacob Faibussowitsch CHKERRQ(PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_MPIAIJ)); 425f80ce2aSJacob Faibussowitsch CHKERRQ(PetscObjectComposeFunction((PetscObject)AA,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ)); 435f80ce2aSJacob Faibussowitsch CHKERRQ(PetscObjectComposeFunction((PetscObject)AB,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ)); 44c4762a1bSJed Brown } 45c4762a1bSJed Brown PetscFunctionReturn(0); 46c4762a1bSJed Brown } 47c4762a1bSJed Brown 48c4762a1bSJed Brown /* this routines queries the already registered MatScaleUserImp_XXX 49c4762a1bSJed Brown implementations for the given matrix, and calls the correct 50c4762a1bSJed Brown routine. i.e if MatType is SeqAIJ, MatScaleUserImpl_SeqAIJ() gets 51c4762a1bSJed Brown called, and if MatType is MPIAIJ, MatScaleUserImpl_MPIAIJ() gets 52c4762a1bSJed Brown called */ 53c4762a1bSJed Brown PetscErrorCode MatScaleUserImpl(Mat mat,PetscScalar a) 54c4762a1bSJed Brown { 555f80ce2aSJacob Faibussowitsch PetscErrorCode (*f)(Mat,PetscScalar); 56c4762a1bSJed Brown 57c4762a1bSJed Brown PetscFunctionBegin; 585f80ce2aSJacob Faibussowitsch CHKERRQ(PetscObjectQueryFunction((PetscObject)mat,"MatScaleUserImpl_C",&f)); 595f80ce2aSJacob Faibussowitsch if (f) CHKERRQ((*f)(mat,a)); 60c4762a1bSJed Brown PetscFunctionReturn(0); 61c4762a1bSJed Brown } 62c4762a1bSJed Brown 63c4762a1bSJed Brown /* Main user code that uses MatScaleUserImpl() */ 64c4762a1bSJed Brown 65c4762a1bSJed Brown int main(int argc,char **args) 66c4762a1bSJed Brown { 67c4762a1bSJed Brown Mat mat; 68c4762a1bSJed Brown PetscInt i,j,m = 2,n,Ii,J; 69c4762a1bSJed Brown PetscScalar v,none = -1.0; 70c4762a1bSJed Brown PetscMPIInt rank,size; 71c4762a1bSJed Brown 72*b122ec5aSJacob Faibussowitsch CHKERRQ(PetscInitialize(&argc,&args,(char*)0,help)); 735f80ce2aSJacob Faibussowitsch CHKERRMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&rank)); 745f80ce2aSJacob Faibussowitsch CHKERRMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size)); 75c4762a1bSJed Brown n = 2*size; 76c4762a1bSJed Brown 77c4762a1bSJed Brown /* create the matrix */ 785f80ce2aSJacob Faibussowitsch CHKERRQ(MatCreate(PETSC_COMM_WORLD,&mat)); 795f80ce2aSJacob Faibussowitsch CHKERRQ(MatSetSizes(mat,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n)); 805f80ce2aSJacob Faibussowitsch CHKERRQ(MatSetType(mat,MATAIJ)); 815f80ce2aSJacob Faibussowitsch CHKERRQ(MatSetUp(mat)); 82c4762a1bSJed Brown 83c4762a1bSJed Brown /* register user defined MatScaleUser() operation for both SeqAIJ 84c4762a1bSJed Brown and MPIAIJ types */ 855f80ce2aSJacob Faibussowitsch CHKERRQ(RegisterMatScaleUserImpl(mat)); 86c4762a1bSJed Brown 87c4762a1bSJed Brown /* assemble the matrix */ 88c4762a1bSJed Brown for (i=0; i<m; i++) { 89c4762a1bSJed Brown for (j=2*rank; j<2*rank+2; j++) { 90c4762a1bSJed Brown v = -1.0; Ii = j + n*i; 915f80ce2aSJacob Faibussowitsch if (i>0) {J = Ii - n; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));} 925f80ce2aSJacob Faibussowitsch if (i<m-1) {J = Ii + n; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));} 935f80ce2aSJacob Faibussowitsch if (j>0) {J = Ii - 1; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));} 945f80ce2aSJacob Faibussowitsch if (j<n-1) {J = Ii + 1; CHKERRQ(MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES));} 955f80ce2aSJacob Faibussowitsch v = 4.0; CHKERRQ(MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES)); 96c4762a1bSJed Brown } 97c4762a1bSJed Brown } 985f80ce2aSJacob Faibussowitsch CHKERRQ(MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY)); 995f80ce2aSJacob Faibussowitsch CHKERRQ(MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY)); 100c4762a1bSJed Brown 101c4762a1bSJed Brown /* check the matrix before and after scaling by -1.0 */ 1025f80ce2aSJacob Faibussowitsch CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Matrix _before_ MatScaleUserImpl() operation\n")); 1035f80ce2aSJacob Faibussowitsch CHKERRQ(MatView(mat,PETSC_VIEWER_STDOUT_WORLD)); 1045f80ce2aSJacob Faibussowitsch CHKERRQ(MatScaleUserImpl(mat,none)); 1055f80ce2aSJacob Faibussowitsch CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Matrix _after_ MatScaleUserImpl() operation\n")); 1065f80ce2aSJacob Faibussowitsch CHKERRQ(MatView(mat,PETSC_VIEWER_STDOUT_WORLD)); 107c4762a1bSJed Brown 1085f80ce2aSJacob Faibussowitsch CHKERRQ(MatDestroy(&mat)); 109*b122ec5aSJacob Faibussowitsch CHKERRQ(PetscFinalize()); 110*b122ec5aSJacob Faibussowitsch return 0; 111c4762a1bSJed Brown } 112c4762a1bSJed Brown 113c4762a1bSJed Brown /*TEST 114c4762a1bSJed Brown 115c4762a1bSJed Brown test: 116c4762a1bSJed Brown 117c4762a1bSJed Brown test: 118c4762a1bSJed Brown suffix: 2 119c4762a1bSJed Brown nsize: 2 120c4762a1bSJed Brown 121c4762a1bSJed Brown TEST*/ 122