1: static char help[] = "Tests MatMPIBAIJSetPreallocationCSR()\n\n";

  3: /*
  4:   Include "petscmat.h" so that we can use matrices.  Note that this file
  5:   automatically includes:
  6:      petscsys.h - base PETSc routines
  7:      petscvec.h - vectors
  8:      petscis.h  - index sets
  9: */
 10: #include <petscmat.h>

 12: int main(int argc, char **args)
 13: {
 14:   Mat         A;
 15:   PetscInt   *ia, *ja, bs = 2;
 16:   PetscInt    N = 9, n;
 17:   PetscInt    rstart, rend, row, col;
 18:   PetscInt    i;
 19:   PetscMPIInt rank, size;
 20:   Vec         v;

 22:   PetscFunctionBeginUser;
 23:   PetscCall(PetscInitialize(&argc, &args, NULL, help));
 24:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 25:   PetscCheck(size < 5, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Can only use at most 4 processors.");
 26:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));

 28:   /* Get a partition range based on the vector size */
 29:   PetscCall(VecCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, PETSC_DECIDE, N, &v));
 30:   PetscCall(VecGetLocalSize(v, &n));
 31:   PetscCall(VecGetOwnershipRange(v, &rstart, &rend));
 32:   PetscCall(VecDestroy(&v));

 34:   PetscCall(PetscMalloc1(n + 1, &ia));
 35:   PetscCall(PetscMalloc1(3 * n, &ja));

 37:   /* Construct a tri-diagonal CSR indexing */
 38:   i     = 1;
 39:   ia[0] = 0;
 40:   for (row = rstart; row < rend; row++) {
 41:     ia[i] = ia[i - 1];

 43:     /* diagonal */
 44:     col = row;
 45:     {
 46:       ja[ia[i]] = col;
 47:       ia[i]++;
 48:     }

 50:     /* lower diagonal */
 51:     col = row - 1;
 52:     if (col >= 0) {
 53:       ja[ia[i]] = col;
 54:       ia[i]++;
 55:     }

 57:     /* upper diagonal */
 58:     col = row + 1;
 59:     if (col < N) {
 60:       ja[ia[i]] = col;
 61:       ia[i]++;
 62:     }
 63:     i++;
 64:   }

 66:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 67:   PetscCall(MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE));
 68:   PetscCall(MatSetType(A, MATMPIAIJ));
 69:   PetscCall(MatMPIAIJSetPreallocationCSR(A, ia, ja, NULL));
 70:   PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
 71:   PetscCall(MatDestroy(&A));

 73:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 74:   PetscCall(MatSetSizes(A, bs * n, bs * n, PETSC_DETERMINE, PETSC_DETERMINE));
 75:   PetscCall(MatSetType(A, MATMPIBAIJ));
 76:   PetscCall(MatMPIBAIJSetPreallocationCSR(A, bs, ia, ja, NULL));
 77:   PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
 78:   PetscCall(MatDestroy(&A));

 80:   PetscCall(PetscFree(ia));
 81:   PetscCall(PetscFree(ja));
 82:   PetscCall(PetscFinalize());
 83:   return 0;
 84: }

 86: /*TEST

 88:     test:
 89:       nsize: 4

 91: TEST*/