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: /* This routine implments MatScaleUserImpl() functionality for MatType
  8:    SeqAIJ. MatScale_SeqAIJ() code duplicated here */
  9: #include <../src/mat/impls/aij/seq/aij.h>
 10: static PetscErrorCode MatScaleUserImpl_SeqAIJ(Mat inA,PetscScalar alpha)
 11: {
 12:   Mat_SeqAIJ     *a     = (Mat_SeqAIJ*)inA->data;
 13:   PetscScalar    oalpha = alpha;
 14:   PetscBLASInt   one    = 1,bnz;

 18:   PetscBLASIntCast(a->nz,&bnz);
 19:   BLASscal_(&bnz,&oalpha,a->a,&one);
 20:   return(0);
 21: }

 23: /* This routine implments MatScaleUserImpl() functionality for MatType
 24:    SeqAIJ. MatScale_MPIAIJ() code duplicated here */
 25: extern PetscErrorCode MatScaleUserImpl(Mat,PetscScalar);
 26: #include <../src/mat/impls/aij/mpi/mpiaij.h>
 27: static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A,PetscScalar aa)
 28: {
 29:   Mat_MPIAIJ     *a = (Mat_MPIAIJ*)A->data;

 33:   /* we can call MatScaleUserImpl_SeqAIJ() directly here instead of
 34:      going through MatScaleUserImpl() wrapper */
 35:   MatScaleUserImpl(a->A,aa);
 36:   MatScaleUserImpl(a->B,aa);
 37:   return(0);
 38: }

 40: /* This routine registers MatScaleUserImpl_SeqAIJ() and
 41:    MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl()
 42:    functionality for SeqAIJ and MPIAIJ matrix-types */
 43: PetscErrorCode RegisterMatScaleUserImpl(Mat mat)
 44: {
 46:   PetscMPIInt    size;

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

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

 53:   } else { /* MPIAIJ Matrix */
 54:     Mat_MPIAIJ *a = (Mat_MPIAIJ*)mat->data;
 55:     PetscObjectComposeFunction((PetscObject)mat,"MatScaleUserImpl_C",MatScaleUserImpl_MPIAIJ);
 56:     PetscObjectComposeFunction((PetscObject)(a->A),"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 57:     PetscObjectComposeFunction((PetscObject)(a->B),"MatScaleUserImpl_C",MatScaleUserImpl_SeqAIJ);
 58:   }
 59:   return(0);
 60: }

 62: /* this routines queries the already registered MatScaleUserImp_XXX
 63:    implementations for the given matrix, and calls the correct
 64:    routine. i.e if MatType is SeqAIJ, MatScaleUserImpl_SeqAIJ() gets
 65:    called, and if MatType is MPIAIJ, MatScaleUserImpl_MPIAIJ() gets
 66:    called */
 67: PetscErrorCode MatScaleUserImpl(Mat mat,PetscScalar a)
 68: {
 69:   PetscErrorCode ierr,(*f)(Mat,PetscScalar);

 72:   PetscObjectQueryFunction((PetscObject)mat,"MatScaleUserImpl_C",&f);
 73:   if (f) {
 74:     (*f)(mat,a);
 75:   }
 76:   return(0);
 77: }

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

 83: int main(int argc,char **args)
 84: {
 85:   Mat            mat;
 86:   PetscInt       i,j,m = 2,n,Ii,J;
 88:   PetscScalar    v,none = -1.0;
 89:   PetscMPIInt    rank,size;


 92:   PetscInitialize(&argc,&args,(char*)0,help);

 94:   MPI_Comm_rank(PETSC_COMM_WORLD,&rank);
 95:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
 96:   n    = 2*size;

 98:   /* create the matrix */
 99:   MatCreate(PETSC_COMM_WORLD,&mat);
100:   MatSetSizes(mat,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n);
101:   MatSetType(mat,MATAIJ);
102:   MatSetUp(mat);

104:   /* register user defined MatScaleUser() operation for both SeqAIJ
105:      and MPIAIJ types */
106:   RegisterMatScaleUserImpl(mat);

108:   /* assemble the matrix */
109:   for (i=0; i<m; i++) {
110:     for (j=2*rank; j<2*rank+2; j++) {
111:       v = -1.0;  Ii = j + n*i;
112:       if (i>0)   {J = Ii - n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
113:       if (i<m-1) {J = Ii + n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
114:       if (j>0)   {J = Ii - 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
115:       if (j<n-1) {J = Ii + 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
116:       v = 4.0; MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES);
117:     }
118:   }
119:   MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
120:   MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);

122:   /* check the matrix before and after scaling by -1.0 */
123:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _before_ MatScaleUserImpl() operation\n");
124:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);
125:   MatScaleUserImpl(mat,none);
126:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _after_ MatScaleUserImpl() operation\n");
127:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);

129:   MatDestroy(&mat);
130:   PetscFinalize();
131:   return 0;
132: }