xref: /petsc/src/mat/tutorials/ex15.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a)
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