xref: /petsc/src/ksp/ksp/tutorials/network/ex2.c (revision 732aec7a18f2199fb53bb9a2f3aef439a834ce31)
1 static char help[] = "This example is based on ex1.c, but generates a random network of chosen sizes and parameters. \n\
2   Usage: -n determines number of nodes. The nonnegative seed can be specified with the flag -seed, otherwise the program generates a random seed.\n\n";
3 
4 #include <petscdmnetwork.h>
5 #include <petscksp.h>
6 #include <time.h>
7 
8 typedef struct {
9   PetscInt    id;  /* node id */
10   PetscScalar inj; /* current injection (A) */
11   PetscBool   gr;  /* grounded ? */
12 } Node;
13 
14 typedef struct {
15   PetscInt    id;  /* branch id */
16   PetscScalar r;   /* resistance (ohms) */
17   PetscScalar bat; /* battery (V) */
18 } Branch;
19 
20 typedef struct Edge {
21   PetscInt     n;
22   PetscInt     i;
23   PetscInt     j;
24   struct Edge *next;
25 } Edge;
26 
findDistance(PetscReal x1,PetscReal x2,PetscReal y1,PetscReal y2)27 PetscReal findDistance(PetscReal x1, PetscReal x2, PetscReal y1, PetscReal y2)
28 {
29   return PetscSqrtReal(PetscPowReal(x2 - x1, 2.0) + PetscPowReal(y2 - y1, 2.0));
30 }
31 
32 /*
33   The algorithm for network formation is based on the paper:
34   Routing of Multipoint Connections, Bernard M. Waxman. 1988
35 */
36 
random_network(PetscInt nvertex,PetscInt * pnbranch,Node ** pnode,Branch ** pbranch,PetscInt ** pedgelist,PetscInt seed)37 PetscErrorCode random_network(PetscInt nvertex, PetscInt *pnbranch, Node **pnode, Branch **pbranch, PetscInt **pedgelist, PetscInt seed)
38 {
39   PetscInt    i, j, nedges = 0;
40   PetscInt   *edgelist;
41   PetscInt    nbat, ncurr, fr, to;
42   PetscReal  *x, *y, value, xmax = 10.0; /* generate points in square */
43   PetscReal   maxdist = 0.0, dist, alpha, beta, prob;
44   PetscRandom rnd;
45   Branch     *branch;
46   Node       *node;
47   Edge       *head = NULL, *nnew = NULL, *aux = NULL;
48 
49   PetscFunctionBeginUser;
50   PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &rnd));
51   PetscCall(PetscRandomSetFromOptions(rnd));
52 
53   PetscCall(PetscRandomSetSeed(rnd, seed));
54   PetscCall(PetscRandomSeed(rnd));
55 
56   /* These parameters might be modified for experimentation */
57   nbat  = (PetscInt)(0.1 * nvertex);
58   ncurr = (PetscInt)(0.1 * nvertex);
59   alpha = 0.6;
60   beta  = 0.2;
61 
62   PetscCall(PetscMalloc2(nvertex, &x, nvertex, &y));
63 
64   PetscCall(PetscRandomSetInterval(rnd, 0.0, xmax));
65   for (i = 0; i < nvertex; i++) {
66     PetscCall(PetscRandomGetValueReal(rnd, &x[i]));
67     PetscCall(PetscRandomGetValueReal(rnd, &y[i]));
68   }
69 
70   /* find maximum distance */
71   for (i = 0; i < nvertex; i++) {
72     for (j = 0; j < nvertex; j++) {
73       dist = findDistance(x[i], x[j], y[i], y[j]);
74       if (dist >= maxdist) maxdist = dist;
75     }
76   }
77 
78   PetscCall(PetscRandomSetInterval(rnd, 0.0, 1.0));
79   for (i = 0; i < nvertex; i++) {
80     for (j = 0; j < nvertex; j++) {
81       if (j != i) {
82         dist = findDistance(x[i], x[j], y[i], y[j]);
83         prob = beta * PetscExpScalar(-dist / (maxdist * alpha));
84         PetscCall(PetscRandomGetValueReal(rnd, &value));
85         if (value <= prob) {
86           PetscCall(PetscMalloc1(1, &nnew));
87           if (head == NULL) {
88             head       = nnew;
89             head->next = NULL;
90             head->n    = nedges;
91             head->i    = i;
92             head->j    = j;
93           } else {
94             aux        = head;
95             head       = nnew;
96             head->n    = nedges;
97             head->next = aux;
98             head->i    = i;
99             head->j    = j;
100           }
101           nedges += 1;
102         }
103       }
104     }
105   }
106 
107   PetscCall(PetscMalloc1(2 * nedges, &edgelist));
108 
109   for (aux = head; aux; aux = aux->next) {
110     edgelist[(aux->n) * 2]     = aux->i;
111     edgelist[(aux->n) * 2 + 1] = aux->j;
112   }
113 
114   aux = head;
115   while (aux != NULL) {
116     nnew = aux;
117     aux  = aux->next;
118     PetscCall(PetscFree(nnew));
119   }
120 
121   PetscCall(PetscCalloc2(nvertex, &node, nedges, &branch));
122 
123   for (i = 0; i < nvertex; i++) {
124     node[i].id  = i;
125     node[i].inj = 0;
126     node[i].gr  = PETSC_FALSE;
127   }
128 
129   for (i = 0; i < nedges; i++) {
130     branch[i].id  = i;
131     branch[i].r   = 1.0;
132     branch[i].bat = 0;
133   }
134 
135   /* Chose random node as ground voltage */
136   PetscCall(PetscRandomSetInterval(rnd, 0.0, nvertex));
137   PetscCall(PetscRandomGetValueReal(rnd, &value));
138   node[(int)value].gr = PETSC_TRUE;
139 
140   /* Create random current and battery injectionsa */
141   for (i = 0; i < ncurr; i++) {
142     PetscCall(PetscRandomSetInterval(rnd, 0.0, nvertex));
143     PetscCall(PetscRandomGetValueReal(rnd, &value));
144     fr = edgelist[(int)value * 2];
145     to = edgelist[(int)value * 2 + 1];
146     node[fr].inj += 1.0;
147     node[to].inj -= 1.0;
148   }
149 
150   for (i = 0; i < nbat; i++) {
151     PetscCall(PetscRandomSetInterval(rnd, 0.0, nedges));
152     PetscCall(PetscRandomGetValueReal(rnd, &value));
153     branch[(int)value].bat += 1.0;
154   }
155 
156   PetscCall(PetscFree2(x, y));
157   PetscCall(PetscRandomDestroy(&rnd));
158 
159   /* assign pointers */
160   *pnbranch  = nedges;
161   *pedgelist = edgelist;
162   *pbranch   = branch;
163   *pnode     = node;
164   PetscFunctionReturn(PETSC_SUCCESS);
165 }
166 
FormOperator(DM networkdm,Mat A,Vec b)167 PetscErrorCode FormOperator(DM networkdm, Mat A, Vec b)
168 {
169   Vec             localb;
170   Branch         *branch;
171   Node           *node;
172   PetscInt        e, v, vStart, vEnd, eStart, eEnd;
173   PetscInt        lofst, lofst_to, lofst_fr, row[2], col[6];
174   PetscBool       ghost;
175   const PetscInt *cone;
176   PetscScalar    *barr, val[6];
177 
178   PetscFunctionBegin;
179   PetscCall(DMGetLocalVector(networkdm, &localb));
180   PetscCall(VecSet(b, 0.0));
181   PetscCall(VecSet(localb, 0.0));
182   PetscCall(MatZeroEntries(A));
183 
184   PetscCall(VecGetArray(localb, &barr));
185 
186   /*
187     We can define the current as a "edge characteristic" and the voltage
188     and the voltage as a "vertex characteristic". With that, we can iterate
189     the list of edges and vertices, query the associated voltages and currents
190     and use them to write the Kirchoff equations.
191   */
192 
193   /* Branch equations: i/r + uj - ui = battery */
194   PetscCall(DMNetworkGetEdgeRange(networkdm, &eStart, &eEnd));
195   for (e = 0; e < eEnd; e++) {
196     PetscCall(DMNetworkGetComponent(networkdm, e, 0, NULL, (void **)&branch, NULL));
197     PetscCall(DMNetworkGetLocalVecOffset(networkdm, e, ALL_COMPONENTS, &lofst));
198 
199     PetscCall(DMNetworkGetConnectedVertices(networkdm, e, &cone));
200     PetscCall(DMNetworkGetLocalVecOffset(networkdm, cone[0], ALL_COMPONENTS, &lofst_fr));
201     PetscCall(DMNetworkGetLocalVecOffset(networkdm, cone[1], ALL_COMPONENTS, &lofst_to));
202 
203     barr[lofst] = branch->bat;
204 
205     row[0] = lofst;
206     col[0] = lofst;
207     val[0] = 1;
208     col[1] = lofst_to;
209     val[1] = 1;
210     col[2] = lofst_fr;
211     val[2] = -1;
212     PetscCall(MatSetValuesLocal(A, 1, row, 3, col, val, ADD_VALUES));
213 
214     /* from node */
215     PetscCall(DMNetworkGetComponent(networkdm, cone[0], 0, NULL, (void **)&node, NULL));
216 
217     if (!node->gr) {
218       row[0] = lofst_fr;
219       col[0] = lofst;
220       val[0] = 1;
221       PetscCall(MatSetValuesLocal(A, 1, row, 1, col, val, ADD_VALUES));
222     }
223 
224     /* to node */
225     PetscCall(DMNetworkGetComponent(networkdm, cone[1], 0, NULL, (void **)&node, NULL));
226 
227     if (!node->gr) {
228       row[0] = lofst_to;
229       col[0] = lofst;
230       val[0] = -1;
231       PetscCall(MatSetValuesLocal(A, 1, row, 1, col, val, ADD_VALUES));
232     }
233   }
234 
235   PetscCall(DMNetworkGetVertexRange(networkdm, &vStart, &vEnd));
236   for (v = vStart; v < vEnd; v++) {
237     PetscCall(DMNetworkIsGhostVertex(networkdm, v, &ghost));
238     if (!ghost) {
239       PetscCall(DMNetworkGetComponent(networkdm, v, 0, NULL, (void **)&node, NULL));
240       PetscCall(DMNetworkGetLocalVecOffset(networkdm, v, ALL_COMPONENTS, &lofst));
241 
242       if (node->gr) {
243         row[0] = lofst;
244         col[0] = lofst;
245         val[0] = 1;
246         PetscCall(MatSetValuesLocal(A, 1, row, 1, col, val, ADD_VALUES));
247       } else {
248         barr[lofst] -= node->inj;
249       }
250     }
251   }
252 
253   PetscCall(VecRestoreArray(localb, &barr));
254 
255   PetscCall(DMLocalToGlobalBegin(networkdm, localb, ADD_VALUES, b));
256   PetscCall(DMLocalToGlobalEnd(networkdm, localb, ADD_VALUES, b));
257   PetscCall(DMRestoreLocalVector(networkdm, &localb));
258 
259   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
260   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
261   PetscFunctionReturn(PETSC_SUCCESS);
262 }
263 
main(int argc,char ** argv)264 int main(int argc, char **argv)
265 {
266   PetscInt      i, nbranch = 0, eStart, eEnd, vStart, vEnd;
267   PetscInt      seed = 0, nnode = 0;
268   PetscMPIInt   size, rank;
269   DM            networkdm;
270   Vec           x, b;
271   Mat           A;
272   KSP           ksp;
273   PetscInt     *edgelist = NULL;
274   PetscInt      componentkey[2];
275   Node         *node;
276   Branch       *branch;
277   PetscLogStage stage[3];
278 
279   PetscFunctionBeginUser;
280   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
281   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
282   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
283 
284   PetscCall(PetscOptionsGetInt(NULL, NULL, "-seed", &seed, NULL));
285 
286   PetscCall(PetscLogStageRegister("Network Creation", &stage[0]));
287   PetscCall(PetscLogStageRegister("DMNetwork data structures", &stage[1]));
288   PetscCall(PetscLogStageRegister("KSP", &stage[2]));
289 
290   PetscCall(PetscLogStagePush(stage[0]));
291   /* "read" data only for processor 0 */
292   if (rank == 0) {
293     nnode = 100;
294     PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &nnode, NULL));
295     PetscCall(random_network(nnode, &nbranch, &node, &branch, &edgelist, seed));
296   }
297   PetscCall(PetscLogStagePop());
298 
299   PetscCall(PetscLogStagePush(stage[1]));
300   PetscCall(DMNetworkCreate(PETSC_COMM_WORLD, &networkdm));
301   PetscCall(DMNetworkRegisterComponent(networkdm, "nstr", sizeof(Node), &componentkey[0]));
302   PetscCall(DMNetworkRegisterComponent(networkdm, "bsrt", sizeof(Branch), &componentkey[1]));
303 
304   /* Set number of nodes/edges and edge connectivity */
305   PetscCall(DMNetworkSetNumSubNetworks(networkdm, PETSC_DECIDE, 1));
306   PetscCall(DMNetworkAddSubnetwork(networkdm, "", nbranch, edgelist, NULL));
307 
308   /* Set up the network layout */
309   PetscCall(DMNetworkLayoutSetUp(networkdm));
310 
311   /* Add network components (physical parameters of nodes and branches) and num of variables */
312   if (rank == 0) {
313     PetscCall(DMNetworkGetEdgeRange(networkdm, &eStart, &eEnd));
314     for (i = eStart; i < eEnd; i++) PetscCall(DMNetworkAddComponent(networkdm, i, componentkey[1], &branch[i - eStart], 1));
315 
316     PetscCall(DMNetworkGetVertexRange(networkdm, &vStart, &vEnd));
317     for (i = vStart; i < vEnd; i++) PetscCall(DMNetworkAddComponent(networkdm, i, componentkey[0], &node[i - vStart], 1));
318   }
319 
320   /* Network partitioning and distribution of data */
321   PetscCall(DMSetUp(networkdm));
322   PetscCall(DMNetworkDistribute(&networkdm, 0));
323   PetscCall(DMNetworkAssembleGraphStructures(networkdm));
324 
325   /* We don't use these data structures anymore since they have been copied to networkdm */
326   if (rank == 0) {
327     PetscCall(PetscFree(edgelist));
328     PetscCall(PetscFree2(node, branch));
329   }
330 
331   /* Create vectors and matrix */
332   PetscCall(DMCreateGlobalVector(networkdm, &x));
333   PetscCall(VecDuplicate(x, &b));
334   PetscCall(DMCreateMatrix(networkdm, &A));
335 
336   PetscCall(PetscLogStagePop());
337 
338   PetscCall(PetscLogStagePush(stage[2]));
339   /* Assembly system of equations */
340   PetscCall(FormOperator(networkdm, A, b));
341 
342   /* Solve linear system: A x = b */
343   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
344   PetscCall(KSPSetOperators(ksp, A, A));
345   PetscCall(KSPSetFromOptions(ksp));
346   PetscCall(KSPSolve(ksp, b, x));
347 
348   PetscCall(PetscLogStagePop());
349 
350   /* Free work space */
351   PetscCall(VecDestroy(&x));
352   PetscCall(VecDestroy(&b));
353   PetscCall(MatDestroy(&A));
354   PetscCall(KSPDestroy(&ksp));
355   PetscCall(DMDestroy(&networkdm));
356   PetscCall(PetscFinalize());
357   return 0;
358 }
359 
360 /*TEST
361 
362    build:
363       requires: !single double defined(PETSC_HAVE_ATTRIBUTEALIGNED) 64bitptr
364 
365    test:
366       args: -ksp_converged_reason
367 
368    test:
369       suffix: 2
370       nsize: 2
371       args: -petscpartitioner_type simple -pc_type asm -sub_pc_type ilu -ksp_converged_reason
372 
373    test:
374       suffix: 3
375       nsize: 4
376       args: -petscpartitioner_type simple -pc_type asm -sub_pc_type lu -sub_pc_factor_shift_type nonzero -ksp_converged_reason
377 
378    test:
379       suffix: graphindex
380       args: -n 20 -vertex_global_section_view -edge_global_section_view
381 
382    test:
383       suffix: graphindex_2
384       nsize: 2
385       args: -petscpartitioner_type simple -n 20 -vertex_global_section_view -edge_global_section_view
386 
387 TEST*/
388