1 static char help[] = "Example of using graph partitioning to partition a graph\n\n"; 2 3 #include <petscmat.h> 4 5 int main(int argc, char **args) { 6 Mat A; 7 MatPartitioning part; 8 IS is; 9 PetscInt r, N = 10, start, end, *vweights; 10 PetscBool set_vweights = PETSC_FALSE, use_edge_weights = PETSC_FALSE; 11 PetscMPIInt rank; 12 MPI_Comm comm; 13 14 PetscFunctionBeginUser; 15 PetscCall(PetscInitialize(&argc, &args, (char *)0, help)); 16 comm = PETSC_COMM_WORLD; 17 PetscCall(PetscOptionsGetInt(NULL, NULL, "-N", &N, NULL)); 18 PetscCallMPI(MPI_Comm_rank(comm, &rank)); 19 PetscCall(MatCreate(comm, &A)); 20 PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, N, N)); 21 PetscCall(MatSetFromOptions(A)); 22 PetscCall(MatSeqAIJSetPreallocation(A, 3, NULL)); 23 PetscCall(MatMPIAIJSetPreallocation(A, 3, NULL, 2, NULL)); 24 PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_vertex_weights", &set_vweights, NULL)); 25 PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_use_edge_weights", &use_edge_weights, NULL)); 26 /* Create a linear mesh */ 27 PetscCall(MatGetOwnershipRange(A, &start, &end)); 28 if (set_vweights) { 29 PetscCall(PetscMalloc1(end - start, &vweights)); 30 for (r = start; r < end; ++r) vweights[r - start] = rank + 1; 31 } 32 for (r = start; r < end; ++r) { 33 if (r == 0) { 34 PetscInt cols[2]; 35 PetscScalar vals[2]; 36 37 cols[0] = r; 38 cols[1] = r + 1; 39 vals[0] = 1.0; 40 vals[1] = use_edge_weights ? 2.0 : 1.0; 41 42 PetscCall(MatSetValues(A, 1, &r, 2, cols, vals, INSERT_VALUES)); 43 } else if (r == N - 1) { 44 PetscInt cols[2]; 45 PetscScalar vals[2]; 46 47 cols[0] = r - 1; 48 cols[1] = r; 49 vals[0] = use_edge_weights ? 3.0 : 1.0; 50 vals[1] = 1.0; 51 52 PetscCall(MatSetValues(A, 1, &r, 2, cols, vals, INSERT_VALUES)); 53 } else { 54 PetscInt cols[3]; 55 PetscScalar vals[3]; 56 57 cols[0] = r - 1; 58 cols[1] = r; 59 cols[2] = r + 1; 60 /* ADJ matrix needs to be symmetric */ 61 vals[0] = use_edge_weights ? (cols[0] == 0 ? 2.0 : 5.0) : 1.0; 62 vals[1] = 1.0; 63 vals[2] = use_edge_weights ? (cols[2] == N - 1 ? 3.0 : 5.0) : 1.0; 64 65 PetscCall(MatSetValues(A, 1, &r, 3, cols, vals, INSERT_VALUES)); 66 } 67 } 68 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 69 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 70 71 PetscCall(MatPartitioningCreate(comm, &part)); 72 PetscCall(MatPartitioningSetAdjacency(part, A)); 73 if (set_vweights) PetscCall(MatPartitioningSetVertexWeights(part, vweights)); 74 if (use_edge_weights) { 75 PetscCall(MatPartitioningSetUseEdgeWeights(part, use_edge_weights)); 76 77 PetscCall(MatPartitioningGetUseEdgeWeights(part, &use_edge_weights)); 78 PetscCheck(use_edge_weights, comm, PETSC_ERR_ARG_INCOMP, "use_edge_weights flag does not setup correctly "); 79 } 80 PetscCall(MatPartitioningSetFromOptions(part)); 81 PetscCall(MatPartitioningApply(part, &is)); 82 PetscCall(ISView(is, PETSC_VIEWER_STDOUT_WORLD)); 83 PetscCall(ISDestroy(&is)); 84 PetscCall(MatPartitioningDestroy(&part)); 85 86 PetscCall(MatDestroy(&A)); 87 PetscCall(PetscFinalize()); 88 return 0; 89 } 90 91 /*TEST 92 93 test: 94 nsize: 3 95 requires: parmetis 96 args: -mat_partitioning_type parmetis 97 98 test: 99 suffix: 2 100 nsize: 3 101 requires: ptscotch 102 args: -mat_partitioning_type ptscotch 103 104 test: 105 suffix: 3 106 nsize: 4 107 requires: party 108 args: -mat_partitioning_type party 109 110 test: 111 suffix: 4 112 nsize: 3 113 requires: chaco 114 args: -mat_partitioning_type chaco 115 116 test: 117 suffix: 5 118 nsize: 3 119 requires: parmetis 120 args: -mat_partitioning_type hierarch -mat_partitioning_hierarchical_nfineparts 3 -mat_partitioning_nparts 10 -N 100 121 122 test: 123 suffix: 6 124 nsize: 3 125 requires: parmetis 126 args: -mat_partitioning_type hierarch -mat_partitioning_hierarchical_nfineparts 3 -mat_partitioning_nparts 10 -N 100 -test_vertex_weights 1 -mat_partitioning_use_edge_weights 1 127 128 test: 129 suffix: 7 130 nsize: 2 131 requires: parmetis 132 args: -mat_partitioning_type hierarch -mat_partitioning_hierarchical_nfineparts 2 -mat_partitioning_nparts 10 -mat_partitioning_hierarchical_fineparttype hierarch -malloc_dump -N 100 -mat_partitioning_improve 1 133 134 test: 135 suffix: 8 136 nsize: 2 137 requires: parmetis 138 args: -mat_partitioning_type parmetis -mat_partitioning_nparts 3 -test_use_edge_weights 1 139 140 test: 141 suffix: 9 142 nsize: 2 143 requires: ptscotch 144 args: -mat_partitioning_type ptscotch -mat_partitioning_nparts 3 -test_use_edge_weights 1 -mat_partitioning_ptscotch_proc_weight 0 145 146 TEST*/ 147