xref: /petsc/src/snes/tutorials/network/water/waterfunctions.c (revision 750b007cd8d816cecd9de99077bb0a703b4cf61a)
1 /* function subroutines used by water.c */
2 
3 #include "water.h"
4 #include <petscdmnetwork.h>
5 
6 PetscScalar Flow_Pipe(Pipe *pipe, PetscScalar hf, PetscScalar ht) {
7   PetscScalar flow_pipe;
8 
9   flow_pipe = PetscSign(hf - ht) * PetscPowScalar(PetscAbsScalar(hf - ht) / pipe->k, (1 / pipe->n));
10   return flow_pipe;
11 }
12 
13 PetscScalar Flow_Pump(Pump *pump, PetscScalar hf, PetscScalar ht) {
14   PetscScalar flow_pump;
15   flow_pump = PetscPowScalar((hf - ht + pump->h0) / pump->r, (1 / pump->n));
16   return flow_pump;
17 }
18 
19 PetscErrorCode FormFunction_Water(DM networkdm, Vec localX, Vec localF, PetscInt nv, PetscInt ne, const PetscInt *vtx, const PetscInt *edges, void *appctx) {
20   const PetscScalar *xarr;
21   const PetscInt    *cone;
22   PetscScalar       *farr, hf, ht, flow;
23   PetscInt           i, key, vnode1, vnode2, offsetnode1, offsetnode2, offset, ncomp;
24   PetscBool          ghostvtex;
25   VERTEX_Water       vertex, vertexnode1, vertexnode2;
26   EDGE_Water         edge;
27   Pipe              *pipe;
28   Pump              *pump;
29   Reservoir         *reservoir;
30   Tank              *tank;
31 
32   PetscFunctionBegin;
33   /* Get arrays for the vectors */
34   PetscCall(VecGetArrayRead(localX, &xarr));
35   PetscCall(VecGetArray(localF, &farr));
36 
37   for (i = 0; i < ne; i++) { /* for each edge */
38     /* Get the offset and the key for the component for edge number e[i] */
39     PetscCall(DMNetworkGetComponent(networkdm, edges[i], 0, &key, (void **)&edge, NULL));
40 
41     /* Get the numbers for the vertices covering this edge */
42     PetscCall(DMNetworkGetConnectedVertices(networkdm, edges[i], &cone));
43     vnode1 = cone[0];
44     vnode2 = cone[1];
45 
46     /* Get the components at the two vertices, their variable offsets */
47     PetscCall(DMNetworkGetNumComponents(networkdm, vnode1, &ncomp));
48     PetscCall(DMNetworkGetComponent(networkdm, vnode1, ncomp - 1, &key, (void **)&vertexnode1, NULL));
49     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vnode1, ncomp - 1, &offsetnode1));
50 
51     PetscCall(DMNetworkGetNumComponents(networkdm, vnode2, &ncomp));
52     PetscCall(DMNetworkGetComponent(networkdm, vnode2, ncomp - 1, &key, (void **)&vertexnode2, NULL));
53     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vnode2, ncomp - 1, &offsetnode2));
54 
55     /* Variables at node1 and node 2 */
56     hf = xarr[offsetnode1];
57     ht = xarr[offsetnode2];
58 
59     flow = 0.0;
60     if (edge->type == EDGE_TYPE_PIPE) {
61       pipe = &edge->pipe;
62       flow = Flow_Pipe(pipe, hf, ht);
63     } else if (edge->type == EDGE_TYPE_PUMP) {
64       pump = &edge->pump;
65       flow = Flow_Pump(pump, hf, ht);
66     }
67     /* Convention: Node 1 has outgoing flow and Node 2 has incoming flow */
68     if (vertexnode1->type == VERTEX_TYPE_JUNCTION) farr[offsetnode1] -= flow;
69     if (vertexnode2->type == VERTEX_TYPE_JUNCTION) farr[offsetnode2] += flow;
70   }
71 
72   /* Subtract Demand flows at the vertices */
73   for (i = 0; i < nv; i++) {
74     PetscCall(DMNetworkIsGhostVertex(networkdm, vtx[i], &ghostvtex));
75     if (ghostvtex) continue;
76 
77     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vtx[i], ALL_COMPONENTS, &offset));
78     PetscCall(DMNetworkGetNumComponents(networkdm, vtx[i], &ncomp));
79     PetscCall(DMNetworkGetComponent(networkdm, vtx[i], ncomp - 1, &key, (void **)&vertex, NULL));
80 
81     if (vertex->type == VERTEX_TYPE_JUNCTION) {
82       farr[offset] -= vertex->junc.demand;
83     } else if (vertex->type == VERTEX_TYPE_RESERVOIR) {
84       reservoir    = &vertex->res;
85       farr[offset] = xarr[offset] - reservoir->head;
86     } else if (vertex->type == VERTEX_TYPE_TANK) {
87       tank         = &vertex->tank;
88       farr[offset] = xarr[offset] - (tank->elev + tank->initlvl);
89     }
90   }
91 
92   PetscCall(VecRestoreArrayRead(localX, &xarr));
93   PetscCall(VecRestoreArray(localF, &farr));
94   PetscFunctionReturn(0);
95 }
96 
97 PetscErrorCode WaterFormFunction(SNES snes, Vec X, Vec F, void *user) {
98   DM              networkdm;
99   Vec             localX, localF;
100   const PetscInt *v, *e;
101   PetscInt        nv, ne;
102 
103   PetscFunctionBegin;
104   /* Get the DM attached with the SNES */
105   PetscCall(SNESGetDM(snes, &networkdm));
106 
107   /* Get local vertices and edges */
108   PetscCall(DMNetworkGetSubnetwork(networkdm, 0, &nv, &ne, &v, &e));
109 
110   /* Get local vectors */
111   PetscCall(DMGetLocalVector(networkdm, &localX));
112   PetscCall(DMGetLocalVector(networkdm, &localF));
113 
114   /* Scatter values from global to local vector */
115   PetscCall(DMGlobalToLocalBegin(networkdm, X, INSERT_VALUES, localX));
116   PetscCall(DMGlobalToLocalEnd(networkdm, X, INSERT_VALUES, localX));
117 
118   /* Initialize residual */
119   PetscCall(VecSet(F, 0.0));
120   PetscCall(VecSet(localF, 0.0));
121 
122   PetscCall(FormFunction_Water(networkdm, localX, localF, nv, ne, v, e, NULL));
123 
124   PetscCall(DMRestoreLocalVector(networkdm, &localX));
125   PetscCall(DMLocalToGlobalBegin(networkdm, localF, ADD_VALUES, F));
126   PetscCall(DMLocalToGlobalEnd(networkdm, localF, ADD_VALUES, F));
127 
128   PetscCall(DMRestoreLocalVector(networkdm, &localF));
129   PetscFunctionReturn(0);
130 }
131 
132 PetscErrorCode WaterSetInitialGuess(DM networkdm, Vec X) {
133   PetscInt        nv, ne;
134   const PetscInt *vtx, *edges;
135   Vec             localX;
136 
137   PetscFunctionBegin;
138   PetscCall(DMGetLocalVector(networkdm, &localX));
139 
140   PetscCall(VecSet(X, 0.0));
141   PetscCall(DMGlobalToLocalBegin(networkdm, X, INSERT_VALUES, localX));
142   PetscCall(DMGlobalToLocalEnd(networkdm, X, INSERT_VALUES, localX));
143 
144   /* Get subnetwork */
145   PetscCall(DMNetworkGetSubnetwork(networkdm, 0, &nv, &ne, &vtx, &edges));
146   PetscCall(SetInitialGuess_Water(networkdm, localX, nv, ne, vtx, edges, NULL));
147 
148   PetscCall(DMLocalToGlobalBegin(networkdm, localX, ADD_VALUES, X));
149   PetscCall(DMLocalToGlobalEnd(networkdm, localX, ADD_VALUES, X));
150   PetscCall(DMRestoreLocalVector(networkdm, &localX));
151   PetscFunctionReturn(0);
152 }
153 
154 PetscErrorCode GetListofEdges_Water(WATERDATA *water, PetscInt *edgelist) {
155   PetscInt  i, j, node1, node2;
156   Pipe     *pipe;
157   Pump     *pump;
158   PetscBool netview = PETSC_FALSE;
159 
160   PetscFunctionBegin;
161   PetscCall(PetscOptionsHasName(NULL, NULL, "-water_view", &netview));
162   for (i = 0; i < water->nedge; i++) {
163     if (water->edge[i].type == EDGE_TYPE_PIPE) {
164       pipe  = &water->edge[i].pipe;
165       node1 = pipe->node1;
166       node2 = pipe->node2;
167       if (netview) PetscCall(PetscPrintf(PETSC_COMM_SELF, "edge %" PetscInt_FMT ", pipe v[%" PetscInt_FMT "] -> v[%" PetscInt_FMT "]\n", i, node1, node2));
168     } else {
169       pump  = &water->edge[i].pump;
170       node1 = pump->node1;
171       node2 = pump->node2;
172       if (netview) PetscCall(PetscPrintf(PETSC_COMM_SELF, "edge %" PetscInt_FMT ", pump v[%" PetscInt_FMT "] -> v[%" PetscInt_FMT "]\n", i, node1, node2));
173     }
174 
175     for (j = 0; j < water->nvertex; j++) {
176       if (water->vertex[j].id == node1) {
177         edgelist[2 * i] = j;
178         break;
179       }
180     }
181 
182     for (j = 0; j < water->nvertex; j++) {
183       if (water->vertex[j].id == node2) {
184         edgelist[2 * i + 1] = j;
185         break;
186       }
187     }
188   }
189   PetscFunctionReturn(0);
190 }
191 
192 PetscErrorCode SetInitialGuess_Water(DM networkdm, Vec localX, PetscInt nv, PetscInt ne, const PetscInt *vtx, const PetscInt *edges, void *appctx) {
193   PetscInt     i, offset, key;
194   PetscBool    ghostvtex, sharedv;
195   VERTEX_Water vertex;
196   PetscScalar *xarr;
197 
198   PetscFunctionBegin;
199   PetscCall(VecGetArray(localX, &xarr));
200   for (i = 0; i < nv; i++) {
201     PetscCall(DMNetworkIsGhostVertex(networkdm, vtx[i], &ghostvtex));
202     PetscCall(DMNetworkIsSharedVertex(networkdm, vtx[i], &sharedv));
203     if (ghostvtex || sharedv) continue;
204 
205     PetscCall(DMNetworkGetComponent(networkdm, vtx[i], 0, &key, (void **)&vertex, NULL));
206     PetscCall(DMNetworkGetLocalVecOffset(networkdm, vtx[i], 0, &offset));
207     if (vertex->type == VERTEX_TYPE_JUNCTION) {
208       xarr[offset] = 100;
209     } else if (vertex->type == VERTEX_TYPE_RESERVOIR) {
210       xarr[offset] = vertex->res.head;
211     } else {
212       xarr[offset] = vertex->tank.initlvl + vertex->tank.elev;
213     }
214   }
215   PetscCall(VecRestoreArray(localX, &xarr));
216   PetscFunctionReturn(0);
217 }
218