xref: /petsc/src/ksp/ksp/tutorials/network/ex1_nest.c (revision 732aec7a18f2199fb53bb9a2f3aef439a834ce31)
1 static char help[] = "This example is based on ex1 using MATNEST. \n";
2 
3 #include <petscdmnetwork.h>
4 #include <petscksp.h>
5 
6 /* The topology looks like:
7 
8             (1)
9             /|\
10            / | \
11           /  |  \
12          R   R   V
13         /    |b4  \
14     b1 /    (4)    \ b2
15       /    /   \    R
16      /   R       R   \
17     /  /           \  \
18    / / b5        b6  \ \
19   //                   \\
20 (2)--------- R -------- (3)
21              b3
22 
23   Nodes:          (1), ... (4)
24   Branches:       b1, ... b6
25   Resistances:    R
26   Voltage source: V
27 
28   Additionally, there is a current source from (2) to (1).
29 */
30 
31 /*
32   Structures containing physical data of circuit.
33   Note that no topology is defined
34 */
35 
36 typedef struct {
37   PetscInt    id;  /* node id */
38   PetscScalar inj; /* current injection (A) */
39   PetscBool   gr;  /* grounded ? */
40 } Node;
41 
42 typedef struct {
43   PetscInt    id;  /* branch id */
44   PetscScalar r;   /* resistance (ohms) */
45   PetscScalar bat; /* battery (V) */
46 } Branch;
47 
48 /*
49   read_data: this routine fills data structures with problem data.
50   This can be substituted by an external parser.
51 */
52 
read_data(PetscInt * pnnode,PetscInt * pnbranch,Node ** pnode,Branch ** pbranch,PetscInt ** pedgelist)53 PetscErrorCode read_data(PetscInt *pnnode, PetscInt *pnbranch, Node **pnode, Branch **pbranch, PetscInt **pedgelist)
54 {
55   PetscInt  nnode, nbranch, i;
56   Branch   *branch;
57   Node     *node;
58   PetscInt *edgelist;
59 
60   PetscFunctionBeginUser;
61   nnode   = 4;
62   nbranch = 6;
63 
64   PetscCall(PetscCalloc1(nnode, &node));
65   PetscCall(PetscCalloc1(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 would have that:
92       edgelist[2*i]     = from node
93       edgelist[2*i + 1] = to node
94   */
95 
96   PetscCall(PetscCalloc1(2 * nbranch, &edgelist));
97 
98   for (i = 0; i < nbranch; i++) {
99     switch (i) {
100     case 0:
101       edgelist[2 * i]     = 0;
102       edgelist[2 * i + 1] = 1;
103       break;
104     case 1:
105       edgelist[2 * i]     = 0;
106       edgelist[2 * i + 1] = 2;
107       break;
108     case 2:
109       edgelist[2 * i]     = 1;
110       edgelist[2 * i + 1] = 2;
111       break;
112     case 3:
113       edgelist[2 * i]     = 0;
114       edgelist[2 * i + 1] = 3;
115       break;
116     case 4:
117       edgelist[2 * i]     = 1;
118       edgelist[2 * i + 1] = 3;
119       break;
120     case 5:
121       edgelist[2 * i]     = 2;
122       edgelist[2 * i + 1] = 3;
123       break;
124     default:
125       break;
126     }
127   }
128 
129   /* assign pointers */
130   *pnnode    = nnode;
131   *pnbranch  = nbranch;
132   *pedgelist = edgelist;
133   *pbranch   = branch;
134   *pnode     = node;
135   PetscFunctionReturn(PETSC_SUCCESS);
136 }
137 
FormOperator(DM networkdm,Mat A,Vec b)138 PetscErrorCode FormOperator(DM networkdm, Mat A, Vec b)
139 {
140   Vec             localb;
141   Branch         *branch;
142   Node           *node;
143   PetscInt        e;
144   PetscInt        v, vStart, vEnd;
145   PetscInt        eStart, eEnd;
146   PetscBool       ghost;
147   const PetscInt *cone;
148   PetscScalar    *barr;
149   PetscInt        lofst, lofst_to, lofst_fr;
150   PetscInt        key;
151   PetscInt        row[2], col[6];
152   PetscScalar     val[6];
153   Mat             e11, c12, c21, v22;
154   Mat           **mats;
155 
156   PetscFunctionBegin;
157   PetscCall(DMGetLocalVector(networkdm, &localb));
158   PetscCall(VecSet(b, 0.0));
159   PetscCall(VecSet(localb, 0.0));
160 
161   PetscCall(VecGetArray(localb, &barr));
162 
163   /* Get nested submatrices */
164   PetscCall(MatNestGetSubMats(A, NULL, NULL, &mats));
165   e11 = mats[0][0]; /* edges */
166   c12 = mats[0][1]; /* couplings */
167   c21 = mats[1][0]; /* couplings */
168   v22 = mats[1][1]; /* vertices */
169 
170   /* Get vertices and edge range */
171   PetscCall(DMNetworkGetEdgeRange(networkdm, &eStart, &eEnd));
172   PetscCall(DMNetworkGetVertexRange(networkdm, &vStart, &vEnd));
173 
174   for (e = 0; e < eEnd; e++) {
175     PetscCall(DMNetworkGetComponent(networkdm, e, 0, &key, (void **)&branch, NULL));
176     PetscCall(DMNetworkGetEdgeOffset(networkdm, e, &lofst));
177 
178     PetscCall(DMNetworkGetConnectedVertices(networkdm, e, &cone));
179     PetscCall(DMNetworkGetVertexOffset(networkdm, cone[0], &lofst_fr));
180     PetscCall(DMNetworkGetVertexOffset(networkdm, cone[1], &lofst_to));
181 
182     /* These are edge-edge and go to e11 */
183     row[0] = lofst;
184     col[0] = lofst;
185     val[0] = 1;
186     PetscCall(MatSetValuesLocal(e11, 1, row, 1, col, val, INSERT_VALUES));
187 
188     /* These are edge-vertex and go to c12 */
189     col[0] = lofst_to;
190     val[0] = 1;
191     col[1] = lofst_fr;
192     val[1] = -1;
193     PetscCall(MatSetValuesLocal(c12, 1, row, 2, col, val, INSERT_VALUES));
194 
195     /* These are edge-vertex and go to c12 */
196     /* from node */
197     PetscCall(DMNetworkGetComponent(networkdm, cone[0], 0, &key, (void **)&node, NULL));
198 
199     if (!node->gr) {
200       row[0] = lofst_fr;
201       col[0] = lofst;
202       val[0] = 1;
203       PetscCall(MatSetValuesLocal(c21, 1, row, 1, col, val, INSERT_VALUES));
204     }
205 
206     /* to node */
207     PetscCall(DMNetworkGetComponent(networkdm, cone[1], 0, &key, (void **)&node, NULL));
208 
209     if (!node->gr) {
210       row[0] = lofst_to;
211       col[0] = lofst;
212       val[0] = -1;
213       PetscCall(MatSetValuesLocal(c21, 1, row, 1, col, val, INSERT_VALUES));
214     }
215 
216     /* TODO: this is not a nested vector. Need to implement nested vector */
217     PetscCall(DMNetworkGetLocalVecOffset(networkdm, e, ALL_COMPONENTS, &lofst));
218     barr[lofst] = branch->bat;
219   }
220 
221   for (v = vStart; v < vEnd; v++) {
222     PetscCall(DMNetworkIsGhostVertex(networkdm, v, &ghost));
223     if (!ghost) {
224       PetscCall(DMNetworkGetComponent(networkdm, v, 0, &key, (void **)&node, NULL));
225       PetscCall(DMNetworkGetVertexOffset(networkdm, v, &lofst));
226 
227       if (node->gr) {
228         row[0] = lofst;
229         col[0] = lofst;
230         val[0] = 1;
231         PetscCall(MatSetValuesLocal(v22, 1, row, 1, col, val, INSERT_VALUES));
232       } else {
233         /* TODO: this is not a nested vector. Need to implement nested vector */
234         PetscCall(DMNetworkGetLocalVecOffset(networkdm, v, ALL_COMPONENTS, &lofst));
235         barr[lofst] -= node->inj;
236       }
237     }
238   }
239 
240   PetscCall(VecRestoreArray(localb, &barr));
241 
242   PetscCall(DMLocalToGlobalBegin(networkdm, localb, ADD_VALUES, b));
243   PetscCall(DMLocalToGlobalEnd(networkdm, localb, ADD_VALUES, b));
244   PetscCall(DMRestoreLocalVector(networkdm, &localb));
245 
246   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
247   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
248   PetscFunctionReturn(PETSC_SUCCESS);
249 }
250 
main(int argc,char ** argv)251 int main(int argc, char **argv)
252 {
253   PetscInt    i, nnode = 0, nbranch = 0;
254   PetscInt    eStart, eEnd, vStart, vEnd;
255   PetscMPIInt size, rank;
256   DM          networkdm;
257   Vec         x, b;
258   Mat         A;
259   KSP         ksp;
260   PetscInt   *edgelist = NULL;
261   PetscInt    componentkey[2];
262   Node       *node;
263   Branch     *branch;
264 
265   PetscFunctionBeginUser;
266   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
267   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
268   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
269 
270   /* "read" data only for processor 0 */
271   if (rank == 0) PetscCall(read_data(&nnode, &nbranch, &node, &branch, &edgelist));
272 
273   PetscCall(DMNetworkCreate(PETSC_COMM_WORLD, &networkdm));
274   PetscCall(DMNetworkRegisterComponent(networkdm, "nstr", sizeof(Node), &componentkey[0]));
275   PetscCall(DMNetworkRegisterComponent(networkdm, "bsrt", sizeof(Branch), &componentkey[1]));
276 
277   /* Set number of nodes/edges, add edge connectivity */
278   PetscCall(DMNetworkSetNumSubNetworks(networkdm, PETSC_DECIDE, 1));
279   PetscCall(DMNetworkAddSubnetwork(networkdm, "", nbranch, edgelist, NULL));
280 
281   /* Set up the network layout */
282   PetscCall(DMNetworkLayoutSetUp(networkdm));
283 
284   /* Add network components (physical parameters of nodes and branches) and num of variables */
285   if (rank == 0) {
286     PetscCall(DMNetworkGetEdgeRange(networkdm, &eStart, &eEnd));
287     for (i = eStart; i < eEnd; i++) PetscCall(DMNetworkAddComponent(networkdm, i, componentkey[1], &branch[i - eStart], 1));
288 
289     PetscCall(DMNetworkGetVertexRange(networkdm, &vStart, &vEnd));
290     for (i = vStart; i < vEnd; i++) PetscCall(DMNetworkAddComponent(networkdm, i, componentkey[0], &node[i - vStart], 1));
291   }
292 
293   /* Network partitioning and distribution of data */
294   PetscCall(DMSetUp(networkdm));
295   PetscCall(DMNetworkDistribute(&networkdm, 0));
296 
297   PetscCall(DMNetworkAssembleGraphStructures(networkdm));
298 
299   /* We don't use these data structures anymore since they have been copied to networkdm */
300   if (rank == 0) {
301     PetscCall(PetscFree(edgelist));
302     PetscCall(PetscFree(node));
303     PetscCall(PetscFree(branch));
304   }
305 
306   PetscCall(DMCreateGlobalVector(networkdm, &x));
307   PetscCall(VecDuplicate(x, &b));
308 
309   PetscCall(DMSetMatType(networkdm, MATNEST));
310   PetscCall(DMCreateMatrix(networkdm, &A));
311 
312   /* Assembly system of equations */
313   PetscCall(FormOperator(networkdm, A, b));
314 
315   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
316   PetscCall(KSPSetOperators(ksp, A, A));
317   PetscCall(KSPSetFromOptions(ksp));
318   PetscCall(KSPSolve(ksp, b, x));
319   PetscCall(VecView(x, 0));
320 
321   PetscCall(VecDestroy(&x));
322   PetscCall(VecDestroy(&b));
323   PetscCall(MatDestroy(&A));
324   PetscCall(KSPDestroy(&ksp));
325   PetscCall(DMDestroy(&networkdm));
326   PetscCall(PetscFinalize());
327   return 0;
328 }
329 
330 /*TEST
331 
332    build:
333       requires: !complex double defined(PETSC_HAVE_ATTRIBUTEALIGNED) 64bitptr
334 
335    test:
336       args: -ksp_converged_reason
337 
338    test:
339       suffix: 2
340       nsize: 2
341       args: -petscpartitioner_type simple -ksp_converged_reason
342 
343 TEST*/
344