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        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;

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

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

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

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

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

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

    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));

    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));

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

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

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

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

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

    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));

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

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

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

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

    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;

    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));

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

    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]));

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

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

    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));

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

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

    297:   PetscCall(DMNetworkAssembleGraphStructures(networkdm));

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

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

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

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

    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));

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

    330: /*TEST

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

    335:    test:
    336:       args: -ksp_converged_reason

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

    343: TEST*/