1: static char help[] = "Example of inverting a block diagonal matrix.\n"
  2:                      "\n";

  4: #include <petscmat.h>

  6: int main(int argc, char **args)
  7: {
  8:   Mat          A, A_inv;
  9:   PetscMPIInt  rank, size;
 10:   PetscInt     M, m, bs, rstart, rend, j, x, y, pass, row;
 11:   PetscInt    *dnnz;
 12:   PetscScalar *v;
 13:   Vec          X, Y;
 14:   PetscReal    norm;

 16:   PetscFunctionBeginUser;
 17:   PetscCall(PetscInitialize(&argc, &args, NULL, help));
 18:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 19:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));

 21:   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "ex184", "Mat");
 22:   M = 8;
 23:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-mat_size", &M, NULL));
 24:   bs = 3;
 25:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-mat_block_size", &bs, NULL));
 26:   PetscOptionsEnd();

 28:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 29:   PetscCall(MatSetFromOptions(A));
 30:   PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, M * bs, M * bs));
 31:   PetscCall(MatSetBlockSize(A, bs));
 32:   PetscCall(MatSetUp(A)); /* called so that MatGetLocalSize() will work */
 33:   PetscCall(MatGetLocalSize(A, &m, NULL));
 34:   PetscCall(PetscMalloc1(m / bs, &dnnz));
 35:   for (j = 0; j < m / bs; j++) dnnz[j] = 1;
 36:   PetscCall(MatXAIJSetPreallocation(A, bs, dnnz, NULL, NULL, NULL));
 37:   PetscCall(PetscFree(dnnz));

 39:   PetscCall(PetscMalloc1(bs * bs, &v));
 40:   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
 41:   /* pass 2 re-assembles new values into the nonzero pattern of pass 1 and pass 3 changes values
 42:      through MatZeroRows(); after each the check must run against an inverse that
 43:      MatInvertBlockDiagonal() recomputed rather than took from its cache */
 44:   for (pass = 1; pass <= 3; pass++) {
 45:     if (pass < 3) {
 46:       for (j = rstart / bs; j < rend / bs; j++) {
 47:         for (x = 0; x < bs; x++) {
 48:           for (y = 0; y < bs; y++) {
 49:             if (x == y) v[y + bs * x] = 2 * bs * pass;
 50:             else v[y + bs * x] = (-1 * (x < y) - 2 * (x > y)) * pass;
 51:           }
 52:         }
 53:         PetscCall(MatSetValuesBlocked(A, 1, &j, 1, &j, v, INSERT_VALUES));
 54:       }
 55:       PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
 56:       PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
 57:     } else {
 58:       row = rstart;
 59:       PetscCall(MatZeroRows(A, 1, &row, 2 * bs, NULL, NULL));
 60:     }

 62:     /* check that A  = inv(inv(A)) */
 63:     PetscCall(MatCreate(PETSC_COMM_WORLD, &A_inv));
 64:     PetscCall(MatSetFromOptions(A_inv));
 65:     PetscCall(MatInvertBlockDiagonalMat(A, A_inv));

 67:     /* Test A_inv * A on a random vector */
 68:     PetscCall(MatCreateVecs(A, &X, &Y));
 69:     PetscCall(VecSetRandom(X, NULL));
 70:     PetscCall(MatMult(A, X, Y));
 71:     PetscCall(VecScale(X, -1));
 72:     PetscCall(MatMultAdd(A_inv, Y, X, X));
 73:     PetscCall(VecNorm(X, NORM_MAX, &norm));
 74:     if (norm > PETSC_SMALL) {
 75:       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Norm of error exceeds tolerance on pass %" PetscInt_FMT ".\nInverse of block diagonal A\n", pass));
 76:       PetscCall(MatView(A_inv, PETSC_VIEWER_STDOUT_WORLD));
 77:     }

 79:     PetscCall(MatDestroy(&A_inv));
 80:     PetscCall(VecDestroy(&X));
 81:     PetscCall(VecDestroy(&Y));
 82:   }
 83:   PetscCall(PetscFree(v));
 84:   PetscCall(MatDestroy(&A));

 86:   PetscCall(PetscFinalize());
 87:   return 0;
 88: }

 90: /*TEST
 91:   test:
 92:     suffix: seqaij
 93:     args: -mat_type seqaij -mat_size 12 -mat_block_size 3
 94:     nsize: 1
 95:     output_file: output/empty.out
 96:   test:
 97:     suffix: seqbaij
 98:     args: -mat_type seqbaij -mat_size 12 -mat_block_size 3
 99:     nsize: 1
100:     output_file: output/empty.out
101:   test:
102:     suffix: mpiaij
103:     args: -mat_type mpiaij -mat_size 12 -mat_block_size 3
104:     nsize: 2
105:     output_file: output/empty.out
106:   test:
107:     suffix: mpibaij
108:     args: -mat_type mpibaij -mat_size 12 -mat_block_size 3
109:     nsize: 2
110:     output_file: output/empty.out
111: TEST*/