2: program main
  3: #include <petsc/finclude/petscvec.h>
  4: #include <petsc/finclude/petscmat.h>

  6: use petscvec
  7: use petscmat

  9: implicit none

 11:   Mat             A
 12:   MatPartitioning   part
 13:   IS              is
 14:   PetscInt   ::     i,m,N
 15:   PetscInt   ::     rstart,rend
 16:   PetscInt,pointer,dimension(:) ::   emptyranks,bigranks,cols
 17:   PetscScalar,pointer,dimension(:) :: vals
 18:   PetscInt :: &
 19:     nbigranks   = 10, &
 20:     nemptyranks = 10
 21:   PetscMPIInt   ::  rank,sizef
 22:   PetscErrorCode  ierr
 23:   PetscBool  set
 24:   PetscInt,parameter :: zero = 0, one = 1, two = 2, three = 3

 26:   PetscCallA(PetscInitialize(ierr))

 28:   PetscCallMPIA(MPI_Comm_rank(PETSC_COMM_WORLD,rank,ierr))
 29:   PetscCallMPIA(MPI_Comm_size(PETSC_COMM_WORLD,sizef,ierr))

 31:   allocate(emptyranks(nemptyranks))
 32:   allocate(bigranks(nbigranks))

 34:   PetscCallA(PetscOptionsGetIntArray(PETSC_NULL_OPTIONS,PETSC_NULL_CHARACTER,"-emptyranks",emptyranks,nemptyranks,set,ierr))
 35:   PetscCallA(PetscOptionsGetIntArray(PETSC_NULL_OPTIONS,PETSC_NULL_CHARACTER,"-bigranks",bigranks,nbigranks,set,ierr))

 37:   m = 1
 38:   do i=1,nemptyranks
 39:     if (rank == emptyranks(i)) m = 0
 40:   end do
 41:   do i=1,nbigranks
 42:     if (rank == bigranks(i)) m = 5
 43:   end do

 45:   deallocate(emptyranks)
 46:   deallocate(bigranks)

 48:   PetscCallA(MatCreate(PETSC_COMM_WORLD,A,ierr))
 49:   PetscCallA(MatSetsizes(A,m,m,PETSC_DECIDE,PETSC_DECIDE,ierr))
 50:   PetscCallA(MatSetFromOptions(A,ierr))
 51:   PetscCallA(MatSeqAIJSetPreallocation(A,three,PETSC_NULL_INTEGER,ierr))
 52:   PetscCallA(MatMPIAIJSetPreallocation(A,three,PETSC_NULL_INTEGER,two,PETSC_NULL_INTEGER,ierr))
 53:   PetscCallA(MatSeqBAIJSetPreallocation(A,one,three,PETSC_NULL_INTEGER,ierr))
 54:   PetscCallA(MatMPIBAIJSetPreallocation(A,one,three,PETSC_NULL_INTEGER,2,PETSC_NULL_INTEGER,ierr))
 55:   PetscCallA(MatSeqSBAIJSetPreallocation(A,one,two,PETSC_NULL_INTEGER,ierr))
 56:   PetscCallA(MatMPISBAIJSetPreallocation(A,one,two,PETSC_NULL_INTEGER,1,PETSC_NULL_INTEGER,ierr))

 58:   PetscCallA(MatGetSize(A,PETSC_NULL_INTEGER,N,ierr))
 59:   PetscCallA(MatGetOwnershipRange(A,rstart,rend,ierr))

 61:   allocate(cols(0:3))
 62:   allocate(vals(0:3))
 63:   do i=rstart,rend-1

 65:     cols = (/mod((i+N-1),N),i,mod((i+1),N)/)
 66:     vals = [1.0,1.0,1.0]
 67:     PetscCallA(MatSetValues(A,one,i,three,cols,vals,INSERT_VALUES,ierr))
 68:   end do
 69:   deallocate(cols)
 70:   deallocate(vals)
 71:   PetscCallA(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY,ierr))
 72:   PetscCallA(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY,ierr))
 73:   PetscCallA(MatView(A,PETSC_VIEWER_STDOUT_WORLD,ierr))

 75:   PetscCallA(MatPartitioningCreate(PETSC_COMM_WORLD,part,ierr))
 76:   PetscCallA(MatPartitioningSetAdjacency(part,A,ierr))
 77:   PetscCallA(MatPartitioningSetFromOptions(part,ierr))
 78:   PetscCallA(MatPartitioningApply(part,is,ierr))
 79:   PetscCallA(ISView(is,PETSC_VIEWER_STDOUT_WORLD,ierr))
 80:   PetscCallA(ISDestroy(is,ierr))
 81:   PetscCallA(MatPartitioningDestroy(part,ierr))
 82:   PetscCallA(MatDestroy(A,ierr))
 83:   PetscCallA(PetscFinalize(ierr))

 85: end program

 87: !/*TEST
 88: !
 89: !   test:
 90: !      nsize: 8
 91: !      args: -emptyranks 0,2,4 -bigranks 1,3,7 -mat_partitioning_type average
 92: !      output_file: output/ex17_1.out
 93: !      # cannot test with external package partitioners since they produce different results on different systems
 94: !
 95: !TEST*/