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

  4: /*T
  5:    Concepts: partitioning
  6:    Processors: 4
  7: T*/

  9: /*
 10:   Include "petscmat.h" so that we can use matrices.  Note that this file
 11:   automatically includes:
 12:      petscsys.h       - base PETSc routines   petscvec.h - vectors
 13:      petscmat.h - matrices
 14:      petscis.h     - index sets
 15:      petscviewer.h - viewers
 16: */
 17:  #include <petscmat.h>

 19: int main(int argc,char **args)
 20: {
 21:   Mat            A;
 22:   PetscInt       *ia,*ja, bs = 2;
 23:   PetscInt       N = 9, n;
 24:   PetscInt       rstart, rend, row, col;
 25:   PetscInt       i;
 27:   PetscMPIInt    rank,size;
 28:   Vec            v;

 30:   PetscInitialize(&argc,&args,(char*)0,help);if (ierr) return ierr;
 31:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
 32:   if ( size > 4 ) SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"Can only use at most 4 processors.");
 33:   MPI_Comm_rank(PETSC_COMM_WORLD,&rank);

 35:   /* Get a partition range based on the vector size */
 36:   VecCreateMPI(PETSC_COMM_WORLD, PETSC_DECIDE, N, &v);
 37:   VecGetLocalSize(v, &n);
 38:   VecGetOwnershipRange(v, &rstart, &rend);
 39:   VecDestroy(&v);

 41:   PetscMalloc1(n+1,&ia);
 42:   PetscMalloc1(3*n,&ja);

 44:   /* Construct a tri-diagonal CSR indexing */
 45:   i = 1;
 46:   ia[0] = 0;
 47:   for ( row = rstart; row < rend; row++ )
 48:   {
 49:     ia[i] = ia[i-1];

 51:     /* diagonal */
 52:     col = row;
 53:     {
 54:       ja[ia[i]] = col;
 55:       ia[i]++;
 56:     }

 58:     /* lower diagonal */
 59:     col = row-1;
 60:     if (col >= 0)
 61:     {
 62:       ja[ia[i]] = col;
 63:       ia[i]++;
 64:     }

 66:     /* upper diagonal */
 67:     col = row+1;
 68:     if (col < N )
 69:     {
 70:       ja[ia[i]] = col;
 71:       ia[i]++;
 72:     }
 73:     i++;
 74:   }

 76:   MatCreate(PETSC_COMM_WORLD, &A);
 77:   MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE);
 78:   MatSetType(A,MATMPIAIJ);
 79:   MatMPIAIJSetPreallocationCSR(A, ia, ja, NULL);
 80:   MatView(A,PETSC_VIEWER_STDOUT_WORLD);
 81:   MatDestroy(&A);

 83:   MatCreate(PETSC_COMM_WORLD, &A);
 84:   MatSetSizes(A, bs*n, bs*n, PETSC_DETERMINE, PETSC_DETERMINE);
 85:   MatSetType(A,MATMPIBAIJ);
 86:   MatMPIBAIJSetPreallocationCSR(A, bs, ia, ja, NULL);
 87:   MatView(A,PETSC_VIEWER_STDOUT_WORLD);
 88:   MatDestroy(&A);

 90:   PetscFree(ia);
 91:   PetscFree(ja);
 92:   PetscFinalize();
 93:   return ierr;
 94: }

 96: /*TEST

 98:     test:
 99:       nsize: 4

101: TEST*/