Actual source code: ex213.c


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

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

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

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

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

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

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

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

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

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

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

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

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

 88: /*TEST

 90:     test:
 91:       nsize: 4

 93: TEST*/