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;

 38:   MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size);

 40:   if (size == 1) { /* SeqAIJ Matrix */
 41:     PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);

 43:   } else { /* MPIAIJ Matrix */
 44:     Mat AA,AB;
 45:     MatMPIAIJGetSeqAIJ(mat,&AA,&AB,NULL);
 46:     PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_MPIAIJ);
 47:     PetscObjectComposeFunction((PetscObject)AA,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 48:     PetscObjectComposeFunction((PetscObject)AB,"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 49:   }
 50:   return(0);
 51: }

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

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

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

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


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

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

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

 96:   /* assemble the matrix */
 97:   for (i=0; i<m; i++) {
 98:     for (j=2*rank; j<2*rank+2; j++) {
 99:       v = -1.0;  Ii = j + n*i;
100:       if (i>0)   {J = Ii - n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
101:       if (i<m-1) {J = Ii + n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
102:       if (j>0)   {J = Ii - 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
103:       if (j<n-1) {J = Ii + 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
104:       v = 4.0; MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES);
105:     }
106:   }
107:   MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
108:   MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);

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

117:   MatDestroy(&mat);
118:   PetscFinalize();
119:   return ierr;
120: }