Date: 2026-06-01T23:13:40+00:00

Actual source code: waterfunctions.c

      1: /* function subroutines used by water.c */

      3: #include "water.h"
      4: #include <petscdmnetwork.h>

      6: PetscScalar Flow_Pipe(Pipe *pipe, PetscScalar hf, PetscScalar ht)
      7: {
      8:   PetscScalar flow_pipe;

     10:   flow_pipe = PetscSign(hf - ht) * PetscPowScalar(PetscAbsScalar(hf - ht) / pipe->k, 1 / pipe->n);
     11:   return flow_pipe;
     12: }

     14: PetscScalar Flow_Pump(Pump *pump, PetscScalar hf, PetscScalar ht)
     15: {
     16:   PetscScalar flow_pump;
     17:   flow_pump = PetscSign(hf - ht + pump->h0) * PetscPowScalar(PetscAbsScalar(hf - ht + pump->h0) / pump->r, 1 / pump->n);
     18:   return flow_pump;
     19: }

     21: PetscErrorCode FormFunction_Water(DM networkdm, Vec localX, Vec localF, PetscInt nv, PetscInt ne, const PetscInt *vtx, const PetscInt *edges, void *appctx)
     22: {
     23:   const PetscScalar *xarr;
     24:   const PetscInt    *cone;
     25:   PetscScalar       *farr, hf, ht, flow;
     26:   PetscInt           i, key, vnode1, vnode2, offsetnode1, offsetnode2, offset, ncomp;
     27:   PetscBool          ghostvtex;
     28:   VERTEX_Water       vertex, vertexnode1, vertexnode2;
     29:   EDGE_Water         edge;
     30:   Pipe              *pipe;
     31:   Pump              *pump;
     32:   Reservoir         *reservoir;
     33:   Tank              *tank;

     35:   PetscFunctionBegin;
     36:   /* Get arrays for the vectors */
     37:   PetscCall(VecGetArrayRead(localX, &xarr));
     38:   PetscCall(VecGetArray(localF, &farr));

     40:   for (i = 0; i < ne; i++) { /* for each edge */
     41:     /* Get the offset and the key for the component for edge number e[i] */
     42:     PetscCall(DMNetworkGetComponent(networkdm, edges[i], 0, &key, (void **)&edge, NULL));

     44:     /* Get the numbers for the vertices covering this edge */
     45:     PetscCall(DMNetworkGetConnectedVertices(networkdm, edges[i], &cone));
     46:     vnode1 = cone[0];
     47:     vnode2 = cone[1];

     49:     /* Get the components at the two vertices, their variable offsets */
     50:     PetscCall(DMNetworkGetNumComponents(networkdm, vnode1, &ncomp));
     51:     PetscCall(DMNetworkGetComponent(networkdm, vnode1, ncomp - 1, &key, (void **)&vertexnode1, NULL));
     52:     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vnode1, ncomp - 1, &offsetnode1));

     54:     PetscCall(DMNetworkGetNumComponents(networkdm, vnode2, &ncomp));
     55:     PetscCall(DMNetworkGetComponent(networkdm, vnode2, ncomp - 1, &key, (void **)&vertexnode2, NULL));
     56:     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vnode2, ncomp - 1, &offsetnode2));

     58:     /* Variables at node1 and node 2 */
     59:     hf = xarr[offsetnode1];
     60:     ht = xarr[offsetnode2];

     62:     flow = 0.0;
     63:     if (edge->type == EDGE_TYPE_PIPE) {
     64:       pipe = &edge->pipe;
     65:       flow = Flow_Pipe(pipe, hf, ht);
     66:     } else if (edge->type == EDGE_TYPE_PUMP) {
     67:       pump = &edge->pump;
     68:       flow = Flow_Pump(pump, hf, ht);
     69:     }
     70:     /* Convention: Node 1 has outgoing flow and Node 2 has incoming flow */
     71:     if (vertexnode1->type == VERTEX_TYPE_JUNCTION) farr[offsetnode1] -= flow;
     72:     if (vertexnode2->type == VERTEX_TYPE_JUNCTION) farr[offsetnode2] += flow;
     73:   }

     75:   /* Subtract Demand flows at the vertices */
     76:   for (i = 0; i < nv; i++) {
     77:     PetscCall(DMNetworkIsGhostVertex(networkdm, vtx[i], &ghostvtex));
     78:     if (ghostvtex) continue;

     80:     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vtx[i], ALL_COMPONENTS, &offset));
     81:     PetscCall(DMNetworkGetNumComponents(networkdm, vtx[i], &ncomp));
     82:     PetscCall(DMNetworkGetComponent(networkdm, vtx[i], ncomp - 1, &key, (void **)&vertex, NULL));

     84:     if (vertex->type == VERTEX_TYPE_JUNCTION) {
     85:       farr[offset] -= vertex->junc.demand;
     86:     } else if (vertex->type == VERTEX_TYPE_RESERVOIR) {
     87:       reservoir    = &vertex->res;
     88:       farr[offset] = xarr[offset] - reservoir->head;
     89:     } else if (vertex->type == VERTEX_TYPE_TANK) {
     90:       tank         = &vertex->tank;
     91:       farr[offset] = xarr[offset] - (tank->elev + tank->initlvl);
     92:     }
     93:   }

     95:   PetscCall(VecRestoreArrayRead(localX, &xarr));
     96:   PetscCall(VecRestoreArray(localF, &farr));
     97:   PetscFunctionReturn(PETSC_SUCCESS);
     98: }

    100: PetscErrorCode WaterFormFunction(SNES snes, Vec X, Vec F, void *user)
    101: {
    102:   DM              networkdm;
    103:   Vec             localX, localF;
    104:   const PetscInt *v, *e;
    105:   PetscInt        nv, ne;

    107:   PetscFunctionBegin;
    108:   /* Get the DM attached with the SNES */
    109:   PetscCall(SNESGetDM(snes, &networkdm));

    111:   /* Get local vertices and edges */
    112:   PetscCall(DMNetworkGetSubnetwork(networkdm, 0, &nv, &ne, &v, &e));

    114:   /* Get local vectors */
    115:   PetscCall(DMGetLocalVector(networkdm, &localX));
    116:   PetscCall(DMGetLocalVector(networkdm, &localF));

    118:   /* Scatter values from global to local vector */
    119:   PetscCall(DMGlobalToLocalBegin(networkdm, X, INSERT_VALUES, localX));
    120:   PetscCall(DMGlobalToLocalEnd(networkdm, X, INSERT_VALUES, localX));

    122:   /* Initialize residual */
    123:   PetscCall(VecSet(F, 0.0));
    124:   PetscCall(VecSet(localF, 0.0));

    126:   PetscCall(FormFunction_Water(networkdm, localX, localF, nv, ne, v, e, NULL));

    128:   PetscCall(DMRestoreLocalVector(networkdm, &localX));
    129:   PetscCall(DMLocalToGlobalBegin(networkdm, localF, ADD_VALUES, F));
    130:   PetscCall(DMLocalToGlobalEnd(networkdm, localF, ADD_VALUES, F));

    132:   PetscCall(DMRestoreLocalVector(networkdm, &localF));
    133:   PetscFunctionReturn(PETSC_SUCCESS);
    134: }

    136: PetscErrorCode WaterSetInitialGuess(DM networkdm, Vec X)
    137: {
    138:   PetscInt        nv, ne;
    139:   const PetscInt *vtx, *edges;
    140:   Vec             localX;

    142:   PetscFunctionBegin;
    143:   PetscCall(DMGetLocalVector(networkdm, &localX));

    145:   PetscCall(VecSet(X, 0.0));
    146:   PetscCall(DMGlobalToLocalBegin(networkdm, X, INSERT_VALUES, localX));
    147:   PetscCall(DMGlobalToLocalEnd(networkdm, X, INSERT_VALUES, localX));

    149:   /* Get subnetwork */
    150:   PetscCall(DMNetworkGetSubnetwork(networkdm, 0, &nv, &ne, &vtx, &edges));
    151:   PetscCall(SetInitialGuess_Water(networkdm, localX, nv, ne, vtx, edges, NULL));

    153:   PetscCall(DMLocalToGlobalBegin(networkdm, localX, ADD_VALUES, X));
    154:   PetscCall(DMLocalToGlobalEnd(networkdm, localX, ADD_VALUES, X));
    155:   PetscCall(DMRestoreLocalVector(networkdm, &localX));
    156:   PetscFunctionReturn(PETSC_SUCCESS);
    157: }

    159: PetscErrorCode GetListofEdges_Water(WATERDATA *water, PetscInt *edgelist)
    160: {
    161:   PetscInt  i, j, node1, node2;
    162:   Pipe     *pipe;
    163:   Pump     *pump;
    164:   PetscBool netview = PETSC_FALSE;

    166:   PetscFunctionBegin;
    167:   PetscCall(PetscOptionsHasName(NULL, NULL, "-water_view", &netview));
    168:   for (i = 0; i < water->nedge; i++) {
    169:     if (water->edge[i].type == EDGE_TYPE_PIPE) {
    170:       pipe  = &water->edge[i].pipe;
    171:       node1 = pipe->node1;
    172:       node2 = pipe->node2;
    173:       if (netview) PetscCall(PetscPrintf(PETSC_COMM_SELF, "edge %" PetscInt_FMT ", pipe v[%" PetscInt_FMT "] -> v[%" PetscInt_FMT "]\n", i, node1, node2));
    174:     } else {
    175:       pump  = &water->edge[i].pump;
    176:       node1 = pump->node1;
    177:       node2 = pump->node2;
    178:       if (netview) PetscCall(PetscPrintf(PETSC_COMM_SELF, "edge %" PetscInt_FMT ", pump v[%" PetscInt_FMT "] -> v[%" PetscInt_FMT "]\n", i, node1, node2));
    179:     }

    181:     for (j = 0; j < water->nvertex; j++) {
    182:       if (water->vertex[j].id == node1) {
    183:         edgelist[2 * i] = j;
    184:         break;
    185:       }
    186:     }

    188:     for (j = 0; j < water->nvertex; j++) {
    189:       if (water->vertex[j].id == node2) {
    190:         edgelist[2 * i + 1] = j;
    191:         break;
    192:       }
    193:     }
    194:   }
    195:   PetscFunctionReturn(PETSC_SUCCESS);
    196: }

    198: PetscErrorCode SetInitialGuess_Water(DM networkdm, Vec localX, PetscInt nv, PetscInt ne, const PetscInt *vtx, const PetscInt *edges, void *appctx)
    199: {
    200:   PetscInt     i, offset, key;
    201:   PetscBool    ghostvtex, sharedv;
    202:   VERTEX_Water vertex;
    203:   PetscScalar *xarr;

    205:   PetscFunctionBegin;
    206:   PetscCall(VecGetArray(localX, &xarr));
    207:   for (i = 0; i < nv; i++) {
    208:     PetscCall(DMNetworkIsGhostVertex(networkdm, vtx[i], &ghostvtex));
    209:     PetscCall(DMNetworkIsSharedVertex(networkdm, vtx[i], &sharedv));
    210:     if (ghostvtex || sharedv) continue;

    212:     PetscCall(DMNetworkGetComponent(networkdm, vtx[i], 0, &key, (void **)&vertex, NULL));
    213:     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vtx[i], 0, &offset));
    214:     if (vertex->type == VERTEX_TYPE_JUNCTION) {
    215:       xarr[offset] = 100;
    216:     } else if (vertex->type == VERTEX_TYPE_RESERVOIR) {
    217:       xarr[offset] = vertex->res.head;
    218:     } else {
    219:       xarr[offset] = vertex->tank.initlvl + vertex->tank.elev;
    220:     }
    221:   }
    222:   PetscCall(VecRestoreArray(localX, &xarr));
    223:   PetscFunctionReturn(PETSC_SUCCESS);
    224: }