Actual source code: ex8.c

      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*/