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: 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 = PetscBLASIntCast(a->nz);

 17:   BLASscal_(&bnz,&oalpha,a->a,&one);
 18:   return(0);
 19: }

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

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

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

 46:   MPI_Comm_size(((PetscObject)mat)->comm, &size);
 47: 
 48:   if (size == 1) { /* SeqAIJ Matrix */
 49:     PetscObjectComposeFunctionDynamic((PetscObject)mat,"MatScaleUserImpl_C",
 50:                                              "MatScaleUserImpl_SeqAIJ",
 51:                                              MatScaleUserImpl_SeqAIJ);

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

 66:   return(0);
 67: }

 69: /* this routines queries the already registered MatScaleUserImp_XXX
 70:    implementations for the given matrix, and calls the correct
 71:    routine. i.e if MatType is SeqAIJ, MatScaleUserImpl_SeqAIJ() gets
 72:    called, and if MatType is MPIAIJ, MatScaleUserImpl_MPIAIJ() gets
 73:    called */
 74: PetscErrorCode MatScaleUserImpl(Mat mat,PetscScalar a)
 75: {
 76:   PetscErrorCode ierr,(*f)(Mat,PetscScalar);

 79:   PetscObjectQueryFunction((PetscObject)mat,"MatScaleUserImpl_C",(void (**)(void))&f);
 80:   if (f) {
 81:     (*f)(mat,a);
 82:   }
 83:   return(0);
 84: }

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

 90: int main(int argc,char **args)
 91: {
 92:   Mat            mat;
 93:   PetscInt       i,j,m = 2,n,Ii,J;
 95:   PetscScalar    v,none = -1.0;
 96:   PetscMPIInt    rank,size;
 97: 

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

101:   MPI_Comm_rank(PETSC_COMM_WORLD,&rank);
102:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
103:   n = 2*size;

105:   /* create the matrix */
106:   MatCreate(PETSC_COMM_WORLD,&mat);
107:   MatSetSizes(mat,PETSC_DECIDE,PETSC_DECIDE,m*n,m*n);
108:   MatSetType(mat,MATAIJ);
109:   MatSetUp(mat);

111:   /* register user defined MatScaleUser() operation for both SeqAIJ
112:      and MPIAIJ types */
113:   RegisterMatScaleUserImpl(mat);

115:   /* assemble the matrix */
116:   for (i=0; i<m; i++) {
117:     for (j=2*rank; j<2*rank+2; j++) {
118:       v = -1.0;  Ii = j + n*i;
119:       if (i>0)   {J = Ii - n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
120:       if (i<m-1) {J = Ii + n; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
121:       if (j>0)   {J = Ii - 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
122:       if (j<n-1) {J = Ii + 1; MatSetValues(mat,1,&Ii,1,&J,&v,INSERT_VALUES);}
123:       v = 4.0; MatSetValues(mat,1,&Ii,1,&Ii,&v,INSERT_VALUES);
124:     }
125:   }
126:   MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY);
127:   MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY);

129:   /* check the matrix before and after scaling by -1.0 */
130:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _before_ MatScaleUserImpl() operation\n");
131:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);
132:   MatScaleUserImpl(mat,none);
133:   PetscPrintf(PETSC_COMM_WORLD,"Matrix _after_ MatScaleUserImpl() operation\n");
134:   MatView(mat,PETSC_VIEWER_STDOUT_WORLD);

136:   MatDestroy(&mat);
137:   PetscFinalize();
138:   return 0;
139: }