Actual source code: ex1_nest.c

  1: static char help[] = "This example is based on ex1 using MATNEST. \n";

  3: #include <petscdmnetwork.h>
  4: #include <petscksp.h>

  6: /* The topology looks like:

  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

 23:   Nodes:          (1), ... (4)
 24:   Branches:       b1, ... b6
 25:   Resistances:    R
 26:   Voltage source: V

 28:   Additionally, there is a current source from (2) to (1).
 29: */

 31: /*
 32:   Structures containing physical data of circuit.
 33:   Note that no topology is defined
 34: */

 36: typedef struct {
 37:   PetscInt    id;  /* node id */
 38:   PetscScalar inj; /* current injection (A) */
 39:   PetscBool   gr;  /* grounded ? */
 40: } Node;

 42: typedef struct {
 43:   PetscInt    id;  /* branch id */
 44:   PetscScalar r;   /* resistance (ohms) */
 45:   PetscScalar bat; /* battery (V) */
 46: } Branch;

 48: /*
 49:   read_data: this routine fills data structures with problem data.
 50:   This can be substituted by an external parser.
 51: */

 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;

 60:   PetscFunctionBeginUser;
 61:   nnode   = 4;
 62:   nbranch = 6;

 64:   PetscCall(PetscCalloc1(nnode, &node));
 65:   PetscCall(PetscCalloc1(nbranch, &branch));

 67:   for (i = 0; i < nnode; i++) {
 68:     node[i].id  = i;
 69:     node[i].inj = 0;
 70:     node[i].gr  = PETSC_FALSE;
 71:   }

 73:   for (i = 0; i < nbranch; i++) {
 74:     branch[i].id  = i;
 75:     branch[i].r   = 1.0;
 76:     branch[i].bat = 0;
 77:   }

 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;

 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:   */

 96:   PetscCall(PetscCalloc1(2 * nbranch, &edgelist));

 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:   }

129:   /* assign pointers */
130:   *pnnode    = nnode;
131:   *pnbranch  = nbranch;
132:   *pedgelist = edgelist;
133:   *pbranch   = branch;
134:   *pnode     = node;
135:   PetscFunctionReturn(PETSC_SUCCESS);
136: }

138: PetscErrorCode FormOperator(DM networkdm, Mat A, Vec b)
139: {
140:   Vec             localb;
141:   Branch         *branch;
142:   Node           *node;
143:   PetscInt        v, vStart, vEnd;
144:   PetscInt        eStart, eEnd;
145:   PetscBool       ghost;
146:   const PetscInt *cone;
147:   PetscScalar    *barr;
148:   PetscInt        lofst, lofst_to, lofst_fr;
149:   PetscInt        key;
150:   PetscInt        row[2], col[6];
151:   PetscScalar     val[6];
152:   Mat             e11, c12, c21, v22;
153:   Mat           **mats;

155:   PetscFunctionBegin;
156:   PetscCall(DMGetLocalVector(networkdm, &localb));
157:   PetscCall(VecSet(b, 0.0));
158:   PetscCall(VecSet(localb, 0.0));

160:   PetscCall(VecGetArray(localb, &barr));

162:   /* Get nested submatrices */
163:   PetscCall(MatNestGetSubMats(A, NULL, NULL, &mats));
164:   e11 = mats[0][0]; /* edges */
165:   c12 = mats[0][1]; /* couplings */
166:   c21 = mats[1][0]; /* couplings */
167:   v22 = mats[1][1]; /* vertices */

169:   /* Get vertices and edge range */
170:   PetscCall(DMNetworkGetEdgeRange(networkdm, &eStart, &eEnd));
171:   PetscCall(DMNetworkGetVertexRange(networkdm, &vStart, &vEnd));

173:   for (PetscInt e = 0; e < eEnd; e++) {
174:     PetscCall(DMNetworkGetComponent(networkdm, e, 0, &key, (void **)&branch, NULL));
175:     PetscCall(DMNetworkGetEdgeOffset(networkdm, e, &lofst));

177:     PetscCall(DMNetworkGetConnectedVertices(networkdm, e, &cone));
178:     PetscCall(DMNetworkGetVertexOffset(networkdm, cone[0], &lofst_fr));
179:     PetscCall(DMNetworkGetVertexOffset(networkdm, cone[1], &lofst_to));

181:     /* These are edge-edge and go to e11 */
182:     row[0] = lofst;
183:     col[0] = lofst;
184:     val[0] = 1;
185:     PetscCall(MatSetValuesLocal(e11, 1, row, 1, col, val, INSERT_VALUES));

187:     /* These are edge-vertex and go to c12 */
188:     col[0] = lofst_to;
189:     val[0] = 1;
190:     col[1] = lofst_fr;
191:     val[1] = -1;
192:     PetscCall(MatSetValuesLocal(c12, 1, row, 2, col, val, INSERT_VALUES));

194:     /* These are edge-vertex and go to c12 */
195:     /* from node */
196:     PetscCall(DMNetworkGetComponent(networkdm, cone[0], 0, &key, (void **)&node, NULL));

198:     if (!node->gr) {
199:       row[0] = lofst_fr;
200:       col[0] = lofst;
201:       val[0] = 1;
202:       PetscCall(MatSetValuesLocal(c21, 1, row, 1, col, val, INSERT_VALUES));
203:     }

205:     /* to node */
206:     PetscCall(DMNetworkGetComponent(networkdm, cone[1], 0, &key, (void **)&node, NULL));

208:     if (!node->gr) {
209:       row[0] = lofst_to;
210:       col[0] = lofst;
211:       val[0] = -1;
212:       PetscCall(MatSetValuesLocal(c21, 1, row, 1, col, val, INSERT_VALUES));
213:     }

215:     /* TODO: this is not a nested vector. Need to implement nested vector */
216:     PetscCall(DMNetworkGetLocalVecOffset(networkdm, e, ALL_COMPONENTS, &lofst));
217:     barr[lofst] = branch->bat;
218:   }

220:   for (v = vStart; v < vEnd; v++) {
221:     PetscCall(DMNetworkIsGhostVertex(networkdm, v, &ghost));
222:     if (!ghost) {
223:       PetscCall(DMNetworkGetComponent(networkdm, v, 0, &key, (void **)&node, NULL));
224:       PetscCall(DMNetworkGetVertexOffset(networkdm, v, &lofst));

226:       if (node->gr) {
227:         row[0] = lofst;
228:         col[0] = lofst;
229:         val[0] = 1;
230:         PetscCall(MatSetValuesLocal(v22, 1, row, 1, col, val, INSERT_VALUES));
231:       } else {
232:         /* TODO: this is not a nested vector. Need to implement nested vector */
233:         PetscCall(DMNetworkGetLocalVecOffset(networkdm, v, ALL_COMPONENTS, &lofst));
234:         barr[lofst] -= node->inj;
235:       }
236:     }
237:   }

239:   PetscCall(VecRestoreArray(localb, &barr));

241:   PetscCall(DMLocalToGlobalBegin(networkdm, localb, ADD_VALUES, b));
242:   PetscCall(DMLocalToGlobalEnd(networkdm, localb, ADD_VALUES, b));
243:   PetscCall(DMRestoreLocalVector(networkdm, &localb));

245:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
246:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
247:   PetscFunctionReturn(PETSC_SUCCESS);
248: }

250: int main(int argc, char **argv)
251: {
252:   PetscInt    i, nnode = 0, nbranch = 0;
253:   PetscInt    eStart, eEnd, vStart, vEnd;
254:   PetscMPIInt size, rank;
255:   DM          networkdm;
256:   Vec         x, b;
257:   Mat         A;
258:   KSP         ksp;
259:   PetscInt   *edgelist = NULL;
260:   PetscInt    componentkey[2];
261:   Node       *node;
262:   Branch     *branch;

264:   PetscFunctionBeginUser;
265:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
266:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
267:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));

269:   /* "read" data only for processor 0 */
270:   if (rank == 0) PetscCall(read_data(&nnode, &nbranch, &node, &branch, &edgelist));

272:   PetscCall(DMNetworkCreate(PETSC_COMM_WORLD, &networkdm));
273:   PetscCall(DMNetworkRegisterComponent(networkdm, "nstr", sizeof(Node), &componentkey[0]));
274:   PetscCall(DMNetworkRegisterComponent(networkdm, "bsrt", sizeof(Branch), &componentkey[1]));

276:   /* Set number of nodes/edges, add edge connectivity */
277:   PetscCall(DMNetworkSetNumSubNetworks(networkdm, PETSC_DECIDE, 1));
278:   PetscCall(DMNetworkAddSubnetwork(networkdm, "", nbranch, edgelist, NULL));

280:   /* Set up the network layout */
281:   PetscCall(DMNetworkLayoutSetUp(networkdm));

283:   /* Add network components (physical parameters of nodes and branches) and num of variables */
284:   if (rank == 0) {
285:     PetscCall(DMNetworkGetEdgeRange(networkdm, &eStart, &eEnd));
286:     for (i = eStart; i < eEnd; i++) PetscCall(DMNetworkAddComponent(networkdm, i, componentkey[1], &branch[i - eStart], 1));

288:     PetscCall(DMNetworkGetVertexRange(networkdm, &vStart, &vEnd));
289:     for (i = vStart; i < vEnd; i++) PetscCall(DMNetworkAddComponent(networkdm, i, componentkey[0], &node[i - vStart], 1));
290:   }

292:   /* Network partitioning and distribution of data */
293:   PetscCall(DMSetUp(networkdm));
294:   PetscCall(DMNetworkDistribute(&networkdm, 0));

296:   PetscCall(DMNetworkAssembleGraphStructures(networkdm));

298:   /* We don't use these data structures anymore since they have been copied to networkdm */
299:   if (rank == 0) {
300:     PetscCall(PetscFree(edgelist));
301:     PetscCall(PetscFree(node));
302:     PetscCall(PetscFree(branch));
303:   }

305:   PetscCall(DMCreateGlobalVector(networkdm, &x));
306:   PetscCall(VecDuplicate(x, &b));

308:   PetscCall(DMSetMatType(networkdm, MATNEST));
309:   PetscCall(DMCreateMatrix(networkdm, &A));

311:   /* Assembly system of equations */
312:   PetscCall(FormOperator(networkdm, A, b));

314:   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
315:   PetscCall(KSPSetOperators(ksp, A, A));
316:   PetscCall(KSPSetFromOptions(ksp));
317:   PetscCall(KSPSolve(ksp, b, x));
318:   PetscCall(VecView(x, 0));

320:   PetscCall(VecDestroy(&x));
321:   PetscCall(VecDestroy(&b));
322:   PetscCall(MatDestroy(&A));
323:   PetscCall(KSPDestroy(&ksp));
324:   PetscCall(DMDestroy(&networkdm));
325:   PetscCall(PetscFinalize());
326:   return 0;
327: }

329: /*TEST

331:    build:
332:       requires: !complex double defined(PETSC_HAVE_ATTRIBUTEALIGNED) 64bitptr

334:    test:
335:       args: -ksp_converged_reason

337:    test:
338:       suffix: 2
339:       nsize: 2
340:       args: -petscpartitioner_type simple -ksp_converged_reason

342: TEST*/