Actual source code: ex213.c

  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   petscvec.h - vectors
  7:      petscmat.h - matrices
  8:      petscis.h     - index sets
  9:      petscviewer.h - viewers
 10: */
 11: #include <petscmat.h>

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

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

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

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

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

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

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

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

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

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

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

 87: /*TEST

 89:     test:
 90:       nsize: 4

 92: TEST*/