Actual source code: ex8.c


  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:   PetscFunctionBegin;
 10:   PetscCall(MatScale(inA, alpha));
 11:   PetscFunctionReturn(PETSC_SUCCESS);
 12: }

 14: extern PetscErrorCode MatScaleUserImpl(Mat, PetscScalar);

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

 20:   PetscFunctionBegin;
 21:   PetscCall(MatMPIAIJGetSeqAIJ(A, &AA, &AB, NULL));
 22:   PetscCall(MatScaleUserImpl(AA, aa));
 23:   PetscCall(MatScaleUserImpl(AB, aa));
 24:   PetscFunctionReturn(PETSC_SUCCESS);
 25: }

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

 34:   PetscFunctionBegin;
 35:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
 36:   if (size == 1) { /* SeqAIJ Matrix */
 37:     PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ));
 38:   } else { /* MPIAIJ Matrix */
 39:     Mat AA, AB;
 40:     PetscCall(MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL));
 41:     PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_MPIAIJ));
 42:     PetscCall(PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ));
 43:     PetscCall(PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ));
 44:   }
 45:   PetscFunctionReturn(PETSC_SUCCESS);
 46: }

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

 55:   PetscFunctionBegin;
 56:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
 57:   if (size == 1) { /* SeqAIJ Matrix */
 58:     PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL));
 59:   } else { /* MPIAIJ Matrix */
 60:     Mat AA, AB;
 61:     PetscCall(MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL));
 62:     PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL));
 63:     PetscCall(PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", NULL));
 64:     PetscCall(PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", NULL));
 65:   }
 66:   PetscFunctionReturn(PETSC_SUCCESS);
 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 (*f)(Mat, PetscScalar);

 78:   PetscFunctionBegin;
 79:   PetscCall(PetscObjectQueryFunction((PetscObject)mat, "MatScaleUserImpl_C", &f));
 80:   if (f) PetscCall((*f)(mat, a));
 81:   PetscFunctionReturn(PETSC_SUCCESS);
 82: }

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

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

 93:   PetscFunctionBeginUser;
 94:   PetscCall(PetscInitialize(&argc, &args, (char *)0, help));
 95:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
 96:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 97:   n = 2 * size;

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

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

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

137:   /* check the matrix before and after scaling by -1.0 */
138:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Matrix _before_ MatScaleUserImpl() operation\n"));
139:   PetscCall(MatView(mat, PETSC_VIEWER_STDOUT_WORLD));
140:   PetscCall(MatScaleUserImpl(mat, none));
141:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Matrix _after_ MatScaleUserImpl() operation\n"));
142:   PetscCall(MatView(mat, PETSC_VIEWER_STDOUT_WORLD));

144:   /* deregister user defined MatScaleUser() operation for both SeqAIJ
145:      and MPIAIJ types */
146:   PetscCall(DeRegisterMatScaleUserImpl(mat));
147:   PetscCall(MatDestroy(&mat));
148:   PetscCall(PetscFinalize());
149:   return 0;
150: }

152: /*TEST

154:    test:

156:    test:
157:       suffix: 2
158:       nsize: 2

160: TEST*/