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

  5:   implicit none

  7:   Mat A
  8:   MatPartitioning part
  9:   IS is
 10:   PetscInt   ::     i, m, N
 11:   PetscInt   ::     rstart, rend
 12:   PetscInt, pointer, dimension(:) ::   emptyranks, bigranks, cols
 13:   PetscScalar, pointer, dimension(:) :: vals
 14:   PetscInt :: &
 15:     nbigranks = 10, &
 16:     nemptyranks = 10
 17:   PetscMPIInt   ::  rank, sizef
 18:   PetscErrorCode ierr
 19:   PetscBool set

 21:   PetscCallA(PetscInitialize(ierr))

 23:   PetscCallMPIA(MPI_Comm_rank(PETSC_COMM_WORLD, rank, ierr))
 24:   PetscCallMPIA(MPI_Comm_size(PETSC_COMM_WORLD, sizef, ierr))

 26:   allocate (emptyranks(nemptyranks))
 27:   allocate (bigranks(nbigranks))

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

 32:   m = 1
 33:   do i = 1, nemptyranks
 34:     if (rank == emptyranks(i)) m = 0
 35:   end do
 36:   do i = 1, nbigranks
 37:     if (rank == bigranks(i)) m = 5
 38:   end do

 40:   deallocate (emptyranks)
 41:   deallocate (bigranks)

 43:   PetscCallA(MatCreate(PETSC_COMM_WORLD, A, ierr))
 44:   PetscCallA(MatSetsizes(A, m, m, PETSC_DECIDE, PETSC_DECIDE, ierr))
 45:   PetscCallA(MatSetFromOptions(A, ierr))
 46:   PetscCallA(MatSeqAIJSetPreallocation(A, 3_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, ierr))
 47:   PetscCallA(MatMPIAIJSetPreallocation(A, 3_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, 2_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, ierr))
 48:   PetscCallA(MatSeqBAIJSetPreallocation(A, 1_PETSC_INT_KIND, 3_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, ierr))
 49:   PetscCallA(MatMPIBAIJSetPreallocation(A, 1_PETSC_INT_KIND, 3_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, 2_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, ierr))
 50:   PetscCallA(MatSeqSBAIJSetPreallocation(A, 1_PETSC_INT_KIND, 2_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, ierr))
 51:   PetscCallA(MatMPISBAIJSetPreallocation(A, 1_PETSC_INT_KIND, 2_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, 1_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, ierr))

 53:   PetscCallA(MatGetSize(A, PETSC_NULL_INTEGER, N, ierr))
 54:   PetscCallA(MatGetOwnershipRange(A, rstart, rend, ierr))

 56:   allocate (cols(1:3))
 57:   allocate (vals(1:3))
 58:   do i = rstart, rend - 1
 59:     cols = [mod(i + N - 1, N), i, mod(i + 1, N)]
 60:     vals = [1.0, 1.0, 1.0]
 61:     PetscCallA(MatSetValues(A, 1_PETSC_INT_KIND, [i], 3_PETSC_INT_KIND, cols, vals, INSERT_VALUES, ierr))
 62:   end do
 63:   deallocate (cols)
 64:   deallocate (vals)
 65:   PetscCallA(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY, ierr))
 66:   PetscCallA(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY, ierr))
 67:   PetscCallA(MatView(A, PETSC_VIEWER_STDOUT_WORLD, ierr))

 69:   PetscCallA(MatPartitioningCreate(PETSC_COMM_WORLD, part, ierr))
 70:   PetscCallA(MatPartitioningSetAdjacency(part, A, ierr))
 71:   PetscCallA(MatPartitioningSetFromOptions(part, ierr))
 72:   PetscCallA(MatPartitioningApply(part, is, ierr))
 73:   PetscCallA(ISView(is, PETSC_VIEWER_STDOUT_WORLD, ierr))
 74:   PetscCallA(ISDestroy(is, ierr))
 75:   PetscCallA(MatPartitioningDestroy(part, ierr))
 76:   PetscCallA(MatDestroy(A, ierr))
 77:   PetscCallA(PetscFinalize(ierr))

 79: end program

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