Actual source code: ex8.c

petsc-master 2016-09-30
Report Typos and Errors
  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);if (ierr) return ierr;
 93:   MPI_Comm_rank(PETSC_COMM_WORLD,&rank);
 94:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
 95:   n    = 2*size;

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

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

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

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

128:   MatDestroy(&mat);
129:   PetscFinalize();
130:   return ierr;
131: }