Date: 2013-07-02T14:44:28+00:00

Actual source code: ex8.c
petsc-3.4.2 2013-07-02

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