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:   call PetscInitialize(PETSC_NULL_CHARACTER,ierr)
 27:   if (ierr /= 0) then
 28:     print*,'PetscInitialize failed'
 29:     stop
 30:   endif

 32:   call MPI_Comm_rank(PETSC_COMM_WORLD,rank,ierr);CHKERRA(ierr)
 33:   call MPI_Comm_size(PETSC_COMM_WORLD,sizef,ierr);CHKERRA(ierr)

 35:   allocate(emptyranks(nemptyranks))
 36:   allocate(bigranks(nbigranks))

 38:   call PetscOptionsGetIntArray(PETSC_NULL_OPTIONS,PETSC_NULL_CHARACTER,"-emptyranks",emptyranks,nemptyranks,set,ierr);CHKERRA(ierr)
 39:   call PetscOptionsGetIntArray(PETSC_NULL_OPTIONS,PETSC_NULL_CHARACTER,"-bigranks",bigranks,nbigranks,set,ierr);CHKERRA(ierr)

 41:   m = 1
 42:   do i=1,nemptyranks
 43:     if (rank == emptyranks(i)) m = 0
 44:   end do
 45:   do i=1,nbigranks
 46:     if (rank == bigranks(i)) m = 5
 47:   end do

 49:   deallocate(emptyranks)
 50:   deallocate(bigranks)

 52:   call MatCreate(PETSC_COMM_WORLD,A,ierr);CHKERRA(ierr)
 53:   call MatSetsizes(A,m,m,PETSC_DECIDE,PETSC_DECIDE,ierr);CHKERRA(ierr)
 54:   call MatSetFromOptions(A,ierr);CHKERRA(ierr)
 55:   call MatSeqAIJSetPreallocation(A,three,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 56:   call MatMPIAIJSetPreallocation(A,three,PETSC_NULL_INTEGER,two,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 57:   call MatSeqBAIJSetPreallocation(A,one,three,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 58:   call MatMPIBAIJSetPreallocation(A,one,three,PETSC_NULL_INTEGER,2,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 59:   call MatSeqSBAIJSetPreallocation(A,one,two,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 60:   call MatMPISBAIJSetPreallocation(A,one,two,PETSC_NULL_INTEGER,1,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)

 62:   call MatGetSize(A,PETSC_NULL_INTEGER,N,ierr);CHKERRA(ierr)
 63:   call MatGetOwnershipRange(A,rstart,rend,ierr);CHKERRA(ierr)

 65:   allocate(cols(0:3))
 66:   allocate(vals(0:3))
 67:   do i=rstart,rend-1

 69:     cols = (/mod((i+N-1),N),i,mod((i+1),N)/)
 70:     vals = [1.0,1.0,1.0]
 71:     call MatSetValues(A,one,i,three,cols,vals,INSERT_VALUES,ierr);CHKERRA(ierr)
 72:   end do
 73:   deallocate(cols)
 74:   deallocate(vals)
 75:   call MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY,ierr);CHKERRA(ierr)
 76:   call MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY,ierr);CHKERRA(ierr)
 77:   call MatView(A,PETSC_VIEWER_STDOUT_WORLD,ierr);CHKERRA(ierr)

 79:   call MatPartitioningCreate(PETSC_COMM_WORLD,part,ierr);CHKERRA(ierr)
 80:   call MatPartitioningSetAdjacency(part,A,ierr);CHKERRA(ierr)
 81:   call MatPartitioningSetFromOptions(part,ierr);CHKERRA(ierr)
 82:   call MatPartitioningApply(part,is,ierr);CHKERRA(ierr)
 83:   call ISView(is,PETSC_VIEWER_STDOUT_WORLD,ierr);CHKERRA(ierr)
 84:   call ISDestroy(is,ierr);CHKERRA(ierr)
 85:   call MatPartitioningDestroy(part,ierr);CHKERRA(ierr)
 86:   call MatDestroy(A,ierr);CHKERRA(ierr)
 87:   call PetscFinalize(ierr);CHKERRA(ierr)

 89: end program

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