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: {

 12:   MatScale(inA,alpha);
 13:   return(0);
 14: }

 16: extern PetscErrorCode MatScaleUserImpl(Mat,PetscScalar);

 18: static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A,PetscScalar aa)
 19: {
 21:   Mat            AA,AB;

 24:   MatMPIAIJGetSeqAIJ(A,&AA,&AB,NULL);
 25:   MatScaleUserImpl(AA,aa);
 26:   MatScaleUserImpl(AB,aa);
 27:   return(0);
 28: }

 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: {
 36:   PetscMPIInt    size;

 39:   MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size);
 40:   if (size == 1) { /* SeqAIJ Matrix */
 41:     PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 42:   } else { /* MPIAIJ Matrix */
 43:     Mat AA,AB;
 44:     MatMPIAIJGetSeqAIJ(mat,&AA,&AB,NULL);
 45:     PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_MPIAIJ);
 46:     PetscObjectComposeFunction((PetscObject)AA,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 47:     PetscObjectComposeFunction((PetscObject)AB,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 48:   }
 49:   return(0);
 50: }

 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);

 62:   PetscObjectQueryFunction((PetscObject)mat,"MatScaleUserImpl_C",&f);
 63:   if (f) {
 64:     (*f)(mat,a);
 65:   }
 66:   return(0);
 67: }

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

 71: int main(int argc,char **args)
 72: {
 73:   Mat            mat;
 74:   PetscInt       i,j,m = 2,n,Ii,J;
 76:   PetscScalar    v,none = -1.0;
 77:   PetscMPIInt    rank,size;

 79:   PetscInitialize(&argc,&args,(char*)0,help);if (ierr) return ierr;
 80:   MPI_Comm_rank(PETSC_COMM_WORLD,&rank);
 81:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
 82:   n    = 2*size;

 84:   /* create the matrix */
 85:   MatCreate(PETSC_COMM_WORLD,&mat);
 86:   MatSetSizes(mat,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n);
 87:   MatSetType(mat,MATAIJ);
 88:   MatSetUp(mat);

 90:   /* register user defined MatScaleUser() operation for both SeqAIJ
 91:      and MPIAIJ types */
 92:   RegisterMatScaleUserImpl(mat);

 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; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
 99:       if (i<m-1) {J = Ii + n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
100:       if (j>0)   {J = Ii - 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
101:       if (j<n-1) {J = Ii + 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
102:       v = 4.0; MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES);
103:     }
104:   }
105:   MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
106:   MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);

108:   /* check the matrix before and after scaling by -1.0 */
109:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _before_ MatScaleUserImpl() operation\n");
110:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);
111:   MatScaleUserImpl(mat,none);
112:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _after_ MatScaleUserImpl() operation\n");
113:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);

115:   MatDestroy(&mat);
116:   PetscFinalize();
117:   return ierr;
118: }

120: /*TEST

122:    test:

124:    test:
125:       suffix: 2
126:       nsize: 2

128: TEST*/