xref: /petsc/src/ksp/ksp/tutorials/network/ex1.c (revision 72c1e49ee3d896ae53f9f914a80b983ec63c03c7)
1 static char help[] = "This example demonstrates the use of DMNetwork interface for solving a simple electric circuit. \n\
2                       The example can be found in p.150 of 'Strang, Gilbert. Computational Science and Engineering. Wellesley, MA'.\n\n";
3 
4 #include <petscdmnetwork.h>
5 #include <petscksp.h>
6 
7 /* The topology looks like:
8 
9             (0)
10             /|\
11            / | \
12           /  |  \
13          R   R   V
14         /    |b3  \
15     b0 /    (3)    \ b1
16       /    /   \    R
17      /   R       R   \
18     /  /           \  \
19    / / b4        b5  \ \
20   //                   \\
21 (1)--------- R -------- (2)
22              b2
23 
24   Nodes:          (0), ... (3)
25   Branches:       b0, ... b5
26   Resistances:    R
27   Voltage source: V
28 
29   Additionally, there is a current source from (1) to (0).
30 */
31 
32 /*
33   Structures containing physical data of circuit.
34   Note that no topology is defined
35 */
36 
37 typedef struct {
38   PetscInt    id;  /* node id */
39   PetscScalar inj; /* current injection (A) */
40   PetscBool   gr;  /* boundary node */
41 } PETSC_ATTRIBUTEALIGNED(PetscMax(sizeof(PetscInt), sizeof(PetscScalar))) Node;
42 
43 typedef struct {
44   PetscInt    id;  /* branch id */
45   PetscScalar r;   /* resistance (ohms) */
46   PetscScalar bat; /* battery (V) */
47 } PETSC_ATTRIBUTEALIGNED(PetscMax(sizeof(PetscInt), sizeof(PetscScalar))) Branch;
48 
49 /*
50   read_data: this routine fills data structures with problem data.
51   This can be substituted by an external parser.
52 */
53 
read_data(PetscInt * pnnode,PetscInt * pnbranch,Node ** pnode,Branch ** pbranch,PetscInt ** pedgelist)54 PetscErrorCode read_data(PetscInt *pnnode, PetscInt *pnbranch, Node **pnode, Branch **pbranch, PetscInt **pedgelist)
55 {
56   PetscInt  nnode, nbranch, i;
57   Branch   *branch;
58   Node     *node;
59   PetscInt *edgelist;
60 
61   PetscFunctionBeginUser;
62   nnode   = 4;
63   nbranch = 6;
64 
65   PetscCall(PetscCalloc2(nnode, &node, nbranch, &branch));
66 
67   for (i = 0; i < nnode; i++) {
68     node[i].id  = i;
69     node[i].inj = 0;
70     node[i].gr  = PETSC_FALSE;
71   }
72 
73   for (i = 0; i < nbranch; i++) {
74     branch[i].id  = i;
75     branch[i].r   = 1.0;
76     branch[i].bat = 0;
77   }
78 
79   /*
80     Branch 1 contains a voltage source of 12.0 V
81     From node 0 to 1 there exists a current source of 4.0 A
82     Node 3 is grounded, hence the voltage is 0.
83   */
84   branch[1].bat = 12.0;
85   node[0].inj   = -4.0;
86   node[1].inj   = 4.0;
87   node[3].gr    = PETSC_TRUE;
88 
89   /*
90     edgelist is an array of len = 2*nbranch that defines the
91     topology of the network. For branch i we have:
92       edgelist[2*i]     = from node
93       edgelist[2*i + 1] = to node.
94   */
95   PetscCall(PetscCalloc1(2 * nbranch, &edgelist));
96 
97   for (i = 0; i < nbranch; i++) {
98     switch (i) {
99     case 0:
100       edgelist[2 * i]     = 0;
101       edgelist[2 * i + 1] = 1;
102       break;
103     case 1:
104       edgelist[2 * i]     = 0;
105       edgelist[2 * i + 1] = 2;
106       break;
107     case 2:
108       edgelist[2 * i]     = 1;
109       edgelist[2 * i + 1] = 2;
110       break;
111     case 3:
112       edgelist[2 * i]     = 0;
113       edgelist[2 * i + 1] = 3;
114       break;
115     case 4:
116       edgelist[2 * i]     = 1;
117       edgelist[2 * i + 1] = 3;
118       break;
119     case 5:
120       edgelist[2 * i]     = 2;
121       edgelist[2 * i + 1] = 3;
122       break;
123     default:
124       break;
125     }
126   }
127 
128   /* assign pointers */
129   *pnnode    = nnode;
130   *pnbranch  = nbranch;
131   *pedgelist = edgelist;
132   *pbranch   = branch;
133   *pnode     = node;
134   PetscFunctionReturn(PETSC_SUCCESS);
135 }
136 
FormOperator(DM dmnetwork,Mat A,Vec b)137 PetscErrorCode FormOperator(DM dmnetwork, Mat A, Vec b)
138 {
139   Branch         *branch;
140   Node           *node;
141   PetscInt        e, v, vStart, vEnd, eStart, eEnd;
142   PetscInt        lofst, lofst_to, lofst_fr, row[2], col[6];
143   PetscBool       ghost;
144   const PetscInt *cone;
145   PetscScalar    *barr, val[6];
146 
147   PetscFunctionBegin;
148   PetscCall(MatZeroEntries(A));
149 
150   PetscCall(VecSet(b, 0.0));
151   PetscCall(VecGetArray(b, &barr));
152 
153   /*
154     We define the current i as an "edge characteristic" and the voltage v as a "vertex characteristic".
155     We iterate the list of edges and vertices, query the associated voltages and currents
156     and use them to write the Kirchoff equations:
157 
158     Branch equations: i/r + v_to - v_from     = v_source (battery)
159     Node equations:   sum(i_to) - sum(i_from) = i_source
160    */
161   PetscCall(DMNetworkGetEdgeRange(dmnetwork, &eStart, &eEnd));
162   for (e = 0; e < eEnd; e++) {
163     PetscCall(DMNetworkGetComponent(dmnetwork, e, 0, NULL, (void **)&branch, NULL));
164     PetscCall(DMNetworkGetLocalVecOffset(dmnetwork, e, ALL_COMPONENTS, &lofst));
165 
166     PetscCall(DMNetworkGetConnectedVertices(dmnetwork, e, &cone));
167     PetscCall(DMNetworkGetLocalVecOffset(dmnetwork, cone[0], ALL_COMPONENTS, &lofst_fr));
168     PetscCall(DMNetworkGetLocalVecOffset(dmnetwork, cone[1], ALL_COMPONENTS, &lofst_to));
169 
170     /* set rhs b for Branch equation */
171     barr[lofst] = branch->bat;
172 
173     /* set Branch equation */
174     row[0] = lofst;
175     col[0] = lofst;
176     val[0] = 1. / branch->r;
177     col[1] = lofst_to;
178     val[1] = 1;
179     col[2] = lofst_fr;
180     val[2] = -1;
181     PetscCall(MatSetValuesLocal(A, 1, row, 3, col, val, ADD_VALUES));
182 
183     /* set Node equation */
184     PetscCall(DMNetworkGetComponent(dmnetwork, cone[0], 0, NULL, (void **)&node, NULL));
185 
186     /* from node */
187     if (!node->gr) { /* not a boundary node */
188       row[0] = lofst_fr;
189       col[0] = lofst;
190       val[0] = -1;
191       PetscCall(MatSetValuesLocal(A, 1, row, 1, col, val, ADD_VALUES));
192     }
193 
194     /* to node */
195     PetscCall(DMNetworkGetComponent(dmnetwork, cone[1], 0, NULL, (void **)&node, NULL));
196 
197     if (!node->gr) { /* not a boundary node */
198       row[0] = lofst_to;
199       col[0] = lofst;
200       val[0] = 1;
201       PetscCall(MatSetValuesLocal(A, 1, row, 1, col, val, ADD_VALUES));
202     }
203   }
204 
205   /* set rhs b for Node equation */
206   PetscCall(DMNetworkGetVertexRange(dmnetwork, &vStart, &vEnd));
207   for (v = vStart; v < vEnd; v++) {
208     PetscCall(DMNetworkIsGhostVertex(dmnetwork, v, &ghost));
209     if (!ghost) {
210       PetscCall(DMNetworkGetComponent(dmnetwork, v, 0, NULL, (void **)&node, NULL));
211       PetscCall(DMNetworkGetLocalVecOffset(dmnetwork, v, ALL_COMPONENTS, &lofst));
212 
213       if (node->gr) { /* a boundary node */
214         row[0] = lofst;
215         col[0] = lofst;
216         val[0] = 1;
217         PetscCall(MatSetValuesLocal(A, 1, row, 1, col, val, ADD_VALUES));
218       } else { /* not a boundary node */
219         barr[lofst] += node->inj;
220       }
221     }
222   }
223 
224   PetscCall(VecRestoreArray(b, &barr));
225 
226   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
227   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
228   PetscFunctionReturn(PETSC_SUCCESS);
229 }
230 
main(int argc,char ** argv)231 int main(int argc, char **argv)
232 {
233   PetscInt    i, nnode = 0, nbranch = 0, eStart, eEnd, vStart, vEnd;
234   PetscMPIInt size, rank;
235   DM          dmnetwork;
236   Vec         x, b;
237   Mat         A;
238   KSP         ksp;
239   PetscInt   *edgelist = NULL;
240   PetscInt    componentkey[2];
241   Node       *node;
242   Branch     *branch;
243   PetscInt    nE[1];
244 
245   PetscFunctionBeginUser;
246   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
247   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
248   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
249 
250   /* "Read" data only for processor 0 */
251   if (rank == 0) PetscCall(read_data(&nnode, &nbranch, &node, &branch, &edgelist));
252 
253   PetscCall(DMNetworkCreate(PETSC_COMM_WORLD, &dmnetwork));
254   PetscCall(DMNetworkRegisterComponent(dmnetwork, "nstr", sizeof(Node), &componentkey[0]));
255   PetscCall(DMNetworkRegisterComponent(dmnetwork, "bsrt", sizeof(Branch), &componentkey[1]));
256 
257   /* Set local number of nodes/edges, add edge connectivity */
258   nE[0] = nbranch;
259   PetscCall(DMNetworkSetNumSubNetworks(dmnetwork, PETSC_DECIDE, 1));
260   PetscCall(DMNetworkAddSubnetwork(dmnetwork, "", nE[0], edgelist, NULL));
261 
262   /* Set up the network layout */
263   PetscCall(DMNetworkLayoutSetUp(dmnetwork));
264 
265   /* Add network components (physical parameters of nodes and branches) and num of variables */
266   if (rank == 0) {
267     PetscCall(DMNetworkGetEdgeRange(dmnetwork, &eStart, &eEnd));
268     for (i = eStart; i < eEnd; i++) PetscCall(DMNetworkAddComponent(dmnetwork, i, componentkey[1], &branch[i - eStart], 1));
269 
270     PetscCall(DMNetworkGetVertexRange(dmnetwork, &vStart, &vEnd));
271     for (i = vStart; i < vEnd; i++) PetscCall(DMNetworkAddComponent(dmnetwork, i, componentkey[0], &node[i - vStart], 1));
272   }
273 
274   /* Network partitioning and distribution of data */
275   PetscCall(DMSetUp(dmnetwork));
276   PetscCall(DMNetworkDistribute(&dmnetwork, 0));
277 
278   /* We do not use these data structures anymore since they have been copied to dmnetwork */
279   if (rank == 0) {
280     PetscCall(PetscFree(edgelist));
281     PetscCall(PetscFree2(node, branch));
282   }
283 
284   /* Create vectors and matrix */
285   PetscCall(DMCreateGlobalVector(dmnetwork, &x));
286   PetscCall(VecDuplicate(x, &b));
287   PetscCall(DMCreateMatrix(dmnetwork, &A));
288 
289   /* Assembly system of equations */
290   PetscCall(FormOperator(dmnetwork, A, b));
291 
292   /* Solve linear system: A x = b */
293   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
294   PetscCall(KSPSetOperators(ksp, A, A));
295   PetscCall(KSPSetFromOptions(ksp));
296   PetscCall(KSPSolve(ksp, b, x));
297   PetscCall(VecView(x, PETSC_VIEWER_STDOUT_WORLD));
298 
299   /* Free work space */
300   PetscCall(VecDestroy(&x));
301   PetscCall(VecDestroy(&b));
302   PetscCall(MatDestroy(&A));
303   PetscCall(KSPDestroy(&ksp));
304   PetscCall(DMDestroy(&dmnetwork));
305   PetscCall(PetscFinalize());
306   return 0;
307 }
308 
309 /*TEST
310 
311    build:
312       requires: !complex double defined(PETSC_HAVE_ATTRIBUTEALIGNED)
313 
314    test:
315       diff_args: -j
316       args: -ksp_monitor_short
317 
318    test:
319       diff_args: -j
320       suffix: 2
321       nsize: 2
322       args: -petscpartitioner_type simple -ksp_converged_reason
323 
324 TEST*/
325