1: static char help[] = "Shows how to add a new MatOperation to AIJ MatType\n\n";

  3: #include <petscmat.h>
  4: #include <petscblaslapack.h>

  6: static PetscErrorCode MatScaleUserImpl_SeqAIJ(Mat inA, PetscScalar alpha)
  7: {
  8:   PetscFunctionBegin;
  9:   PetscCall(MatScale(inA, alpha));
 10:   PetscFunctionReturn(PETSC_SUCCESS);
 11: }

 13: extern PetscErrorCode MatScaleUserImpl(Mat, PetscScalar);

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

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

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

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

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

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

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

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

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

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

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

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

104:   /* register user defined MatScaleUser() operation for both SeqAIJ
105:      and MPIAIJ types */
106:   PetscCall(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;
112:       Ii = j + n * i;
113:       if (i > 0) {
114:         J = Ii - n;
115:         PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
116:       }
117:       if (i < m - 1) {
118:         J = Ii + n;
119:         PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
120:       }
121:       if (j > 0) {
122:         J = Ii - 1;
123:         PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
124:       }
125:       if (j < n - 1) {
126:         J = Ii + 1;
127:         PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
128:       }
129:       v = 4.0;
130:       PetscCall(MatSetValues(mat, 1, &Ii, 1, &Ii, &v, INSERT_VALUES));
131:     }
132:   }
133:   PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY));
134:   PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY));

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

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

151: /*TEST

153:    test:

155:    test:
156:       suffix: 2
157:       nsize: 2

159: TEST*/