2: static char help[] = "Shows how to add a new MatOperation to AIJ MatType\n\n";

  4: #include <petscmat.h>
  5: #include <petscblaslapack.h>

  7: static PetscErrorCode MatScaleUserImpl_SeqAIJ(Mat inA,PetscScalar alpha)
  8: {
  9:   MatScale(inA,alpha);
 10:   return 0;
 11: }

 13: extern PetscErrorCode MatScaleUserImpl(Mat,PetscScalar);

 15: static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A,PetscScalar aa)
 16: {
 17:   Mat            AA,AB;

 19:   MatMPIAIJGetSeqAIJ(A,&AA,&AB,NULL);
 20:   MatScaleUserImpl(AA,aa);
 21:   MatScaleUserImpl(AB,aa);
 22:   return 0;
 23: }

 25: /* This routine registers MatScaleUserImpl_SeqAIJ() and
 26:    MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl()
 27:    functionality for SeqAIJ and MPIAIJ matrix-types */
 28: PetscErrorCode RegisterMatScaleUserImpl(Mat mat)
 29: {
 30:   PetscMPIInt    size;

 32:   MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size);
 33:   if (size == 1) { /* SeqAIJ Matrix */
 34:     PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 35:   } else { /* MPIAIJ Matrix */
 36:     Mat AA,AB;
 37:     MatMPIAIJGetSeqAIJ(mat,&AA,&AB,NULL);
 38:     PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_MPIAIJ);
 39:     PetscObjectComposeFunction((PetscObject)AA,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 40:     PetscObjectComposeFunction((PetscObject)AB,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 41:   }
 42:   return 0;
 43: }

 45: /* this routines queries the already registered MatScaleUserImp_XXX
 46:    implementations for the given matrix, and calls the correct
 47:    routine. i.e if MatType is SeqAIJ, MatScaleUserImpl_SeqAIJ() gets
 48:    called, and if MatType is MPIAIJ, MatScaleUserImpl_MPIAIJ() gets
 49:    called */
 50: PetscErrorCode MatScaleUserImpl(Mat mat,PetscScalar a)
 51: {
 52:   PetscErrorCode (*f)(Mat,PetscScalar);

 54:   PetscObjectQueryFunction((PetscObject)mat,"MatScaleUserImpl_C",&f);
 55:   if (f) (*f)(mat,a);
 56:   return 0;
 57: }

 59: /* Main user code that uses MatScaleUserImpl() */

 61: int main(int argc,char **args)
 62: {
 63:   Mat            mat;
 64:   PetscInt       i,j,m = 2,n,Ii,J;
 65:   PetscScalar    v,none = -1.0;
 66:   PetscMPIInt    rank,size;

 68:   PetscInitialize(&argc,&args,(char*)0,help);
 69:   MPI_Comm_rank(PETSC_COMM_WORLD,&rank);
 70:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
 71:   n    = 2*size;

 73:   /* create the matrix */
 74:   MatCreate(PETSC_COMM_WORLD,&mat);
 75:   MatSetSizes(mat,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n);
 76:   MatSetType(mat,MATAIJ);
 77:   MatSetUp(mat);

 79:   /* register user defined MatScaleUser() operation for both SeqAIJ
 80:      and MPIAIJ types */
 81:   RegisterMatScaleUserImpl(mat);

 83:   /* assemble the matrix */
 84:   for (i=0; i<m; i++) {
 85:     for (j=2*rank; j<2*rank+2; j++) {
 86:       v = -1.0;  Ii = j + n*i;
 87:       if (i>0)   {J = Ii - n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
 88:       if (i<m-1) {J = Ii + n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
 89:       if (j>0)   {J = Ii - 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
 90:       if (j<n-1) {J = Ii + 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
 91:       v = 4.0; MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES);
 92:     }
 93:   }
 94:   MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
 95:   MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);

 97:   /* check the matrix before and after scaling by -1.0 */
 98:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _before_ MatScaleUserImpl() operation\n");
 99:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);
100:   MatScaleUserImpl(mat,none);
101:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _after_ MatScaleUserImpl() operation\n");
102:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);

104:   MatDestroy(&mat);
105:   PetscFinalize();
106:   return 0;
107: }

109: /*TEST

111:    test:

113:    test:
114:       suffix: 2
115:       nsize: 2

117: TEST*/