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

  5: use petscvec
  6: use petscmat

  8: implicit none

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

 25:   PetscCallA(PetscInitialize(ierr))

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

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

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

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

 44:   deallocate(emptyranks)
 45:   deallocate(bigranks)

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

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

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

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

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

 84: end program

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