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

 25:   call PetscInitialize(PETSC_NULL_CHARACTER,ierr)
 26:   if (ierr /= 0) then
 27:     print*,'PetscInitialize failed'
 28:     stop
 29:   endif
 30: 
 31:   call MPI_Comm_rank(PETSC_COMM_WORLD,rank,ierr);CHKERRA(ierr)
 32:   call MPI_Comm_size(PETSC_COMM_WORLD,sizef,ierr);CHKERRA(ierr)

 34:   allocate(emptyranks(0:nemptyranks-1))
 35:   allocate(bigranks(0:nbigranks-1))

 37:   call PetscOptionsGetIntArray(PETSC_NULL_OPTIONS,"Ranks to be skipped by partition","-emptyranks",emptyranks,nemptyranks,set,ierr);CHKERRA(ierr)
 38:   call PetscOptionsGetIntArray(PETSC_NULL_OPTIONS,"Ranks to be overloaded","-bigranks",bigranks,nbigranks,set,ierr);CHKERRA(ierr)
 39: 
 40:   m = 1
 41:   do i=0,nemptyranks-1
 42:     if (rank == emptyranks(i)) m = 0
 43:   end do
 44:   do i=0,nbigranks-1
 45:     if (rank == bigranks(i)) m = 5
 46:   end do

 48:   call MatCreate(PETSC_COMM_WORLD,A,ierr);CHKERRA(ierr)
 49:   call MatSetsizes(A,m,m,PETSC_DECIDE,PETSC_DECIDE,ierr);CHKERRA(ierr)
 50:   call MatSetFromOptions(A,ierr);CHKERRA(ierr)
 51:   call MatSeqAIJSetPreallocation(A,3,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 52:   call MatMPIAIJSetPreallocation(A,3,PETSC_NULL_INTEGER,2,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 53:   call MatSeqBAIJSetPreallocation(A,1,3,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 54:   call MatMPIBAIJSetPreallocation(A,1,3,PETSC_NULL_INTEGER,2,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 55:   call MatSeqSBAIJSetPreallocation(A,1,2,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)
 56:   call MatMPISBAIJSetPreallocation(A,1,2,PETSC_NULL_INTEGER,1,PETSC_NULL_INTEGER,ierr);CHKERRA(ierr)

 58:   call MatGetSize(A,PETSC_NULL_INTEGER,N,ierr);CHKERRA(ierr)
 59:   call MatGetOwnershipRange(A,rstart,rend,ierr);CHKERRA(ierr)
 60: 
 61:   do i=rstart,rend-1
 62:     allocate(cols(0:3))
 63:     allocate(vals(0:3))
 64: 
 65:     cols = (/mod((i+N-1),N),i,mod((i+1),N)/)
 66:     vals = [1.0,1.0,1.0]
 67:     call MatSetValues(A,1,i,3,cols,vals,INSERT_VALUES,ierr);CHKERRA(ierr)
 68:   end do

 70:   call MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY,ierr);CHKERRA(ierr)
 71:   call MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY,ierr);CHKERRA(ierr)
 72:   call MatView(A,PETSC_VIEWER_STDOUT_WORLD,ierr);CHKERRA(ierr)

 74:   call MatPartitioningCreate(PETSC_COMM_WORLD,part,ierr);CHKERRA(ierr)
 75:   call MatPartitioningSetAdjacency(part,A,ierr);CHKERRA(ierr)
 76:   call MatPartitioningSetFromOptions(part,ierr);CHKERRA(ierr)
 77:   call MatPartitioningApply(part,is,ierr);CHKERRA(ierr)
 78:   call ISView(is,PETSC_VIEWER_STDOUT_WORLD,ierr);CHKERRA(ierr)
 79:   call ISDestroy(is,ierr);CHKERRA(ierr)
 80:   call MatPartitioningDestroy(part,ierr);CHKERRA(ierr)
 81:   call MatDestroy(A,ierr);CHKERRA(ierr)
 82:   deallocate(emptyranks)
 83:   deallocate(bigranks)
 84:   deallocate(cols)
 85:   deallocate(vals)
 86:   call PetscFinalize(ierr);CHKERRA(ierr)
 87: 
 88: end program