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: {
  9:   MatScale(inA, alpha);
 10:   return 0;
 11: }

 13: extern PetscErrorCode MatScaleUserImpl(Mat, PetscScalar);

 15: static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A, PetscScalar aa)
 16: {
 17:   Mat AA, AB;

 19:   MatMPIAIJGetSeqAIJ(A, &AA, &AB, NULL);
 20:   MatScaleUserImpl(AA, aa);
 21:   MatScaleUserImpl(AB, aa);
 22:   return 0;
 23: }

 25: /* This routine registers MatScaleUserImpl_SeqAIJ() and
 26:    MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl()
 27:    functionality for SeqAIJ and MPIAIJ matrix-types */
 28: PetscErrorCode RegisterMatScaleUserImpl(Mat mat)
 29: {
 30:   PetscMPIInt size;

 32:   MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size);
 33:   if (size == 1) { /* SeqAIJ Matrix */
 34:     PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ);
 35:   } else { /* MPIAIJ Matrix */
 36:     Mat AA, AB;
 37:     MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL);
 38:     PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_MPIAIJ);
 39:     PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ);
 40:     PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ);
 41:   }
 42:   return 0;
 43: }

 45: /* This routine deregisters MatScaleUserImpl_SeqAIJ() and
 46:    MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl()
 47:    functionality for SeqAIJ and MPIAIJ matrix-types */
 48: PetscErrorCode DeRegisterMatScaleUserImpl(Mat mat)
 49: {
 50:   PetscMPIInt size;

 52:   MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size);
 53:   if (size == 1) { /* SeqAIJ Matrix */
 54:     PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL);
 55:   } else { /* MPIAIJ Matrix */
 56:     Mat AA, AB;
 57:     MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL);
 58:     PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL);
 59:     PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", NULL);
 60:     PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", NULL);
 61:   }
 62:   return 0;
 63: }

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

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

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

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

 89:   PetscInitialize(&argc, &args, (char *)0, help);
 90:   MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
 91:   MPI_Comm_size(PETSC_COMM_WORLD, &size);
 92:   n = 2 * size;

 94:   /* create the matrix */
 95:   MatCreate(PETSC_COMM_WORLD, &mat);
 96:   MatSetSizes(mat, PETSC_DECIDE, PETSC_DECIDE, m * n, m * n);
 97:   MatSetType(mat, MATAIJ);
 98:   MatSetUp(mat);

100:   /* register user defined MatScaleUser() operation for both SeqAIJ
101:      and MPIAIJ types */
102:   RegisterMatScaleUserImpl(mat);

104:   /* assemble the matrix */
105:   for (i = 0; i < m; i++) {
106:     for (j = 2 * rank; j < 2 * rank + 2; j++) {
107:       v  = -1.0;
108:       Ii = j + n * i;
109:       if (i > 0) {
110:         J = Ii - n;
111:         MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES);
112:       }
113:       if (i < m - 1) {
114:         J = Ii + n;
115:         MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES);
116:       }
117:       if (j > 0) {
118:         J = Ii - 1;
119:         MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES);
120:       }
121:       if (j < n - 1) {
122:         J = Ii + 1;
123:         MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES);
124:       }
125:       v = 4.0;
126:       MatSetValues(mat, 1, &Ii, 1, &Ii, &v, INSERT_VALUES);
127:     }
128:   }
129:   MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY);
130:   MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY);

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

139:   /* deregister user defined MatScaleUser() operation for both SeqAIJ
140:      and MPIAIJ types */
141:   DeRegisterMatScaleUserImpl(mat);
142:   MatDestroy(&mat);
143:   PetscFinalize();
144:   return 0;
145: }

147: /*TEST

149:    test:

151:    test:
152:       suffix: 2
153:       nsize: 2

155: TEST*/