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