xref: /petsc/src/snes/tutorials/network/water/waterfunctions.c (revision 6dd63270497ad23dcf16ae500a87ff2b2a0b7474)
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 {
8   PetscScalar flow_pipe;
9 
10   flow_pipe = PetscSign(hf - ht) * PetscPowScalar(PetscAbsScalar(hf - ht) / pipe->k, 1 / pipe->n);
11   return flow_pipe;
12 }
13 
14 PetscScalar Flow_Pump(Pump *pump, PetscScalar hf, PetscScalar ht)
15 {
16   PetscScalar flow_pump;
17   flow_pump = PetscPowScalar((hf - ht + pump->h0) / pump->r, 1 / pump->n);
18   return flow_pump;
19 }
20 
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;
34 
35   PetscFunctionBegin;
36   /* Get arrays for the vectors */
37   PetscCall(VecGetArrayRead(localX, &xarr));
38   PetscCall(VecGetArray(localF, &farr));
39 
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));
43 
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];
48 
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));
53 
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));
57 
58     /* Variables at node1 and node 2 */
59     hf = xarr[offsetnode1];
60     ht = xarr[offsetnode2];
61 
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   }
74 
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;
79 
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));
83 
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   }
94 
95   PetscCall(VecRestoreArrayRead(localX, &xarr));
96   PetscCall(VecRestoreArray(localF, &farr));
97   PetscFunctionReturn(PETSC_SUCCESS);
98 }
99 
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;
106 
107   PetscFunctionBegin;
108   /* Get the DM attached with the SNES */
109   PetscCall(SNESGetDM(snes, &networkdm));
110 
111   /* Get local vertices and edges */
112   PetscCall(DMNetworkGetSubnetwork(networkdm, 0, &nv, &ne, &v, &e));
113 
114   /* Get local vectors */
115   PetscCall(DMGetLocalVector(networkdm, &localX));
116   PetscCall(DMGetLocalVector(networkdm, &localF));
117 
118   /* Scatter values from global to local vector */
119   PetscCall(DMGlobalToLocalBegin(networkdm, X, INSERT_VALUES, localX));
120   PetscCall(DMGlobalToLocalEnd(networkdm, X, INSERT_VALUES, localX));
121 
122   /* Initialize residual */
123   PetscCall(VecSet(F, 0.0));
124   PetscCall(VecSet(localF, 0.0));
125 
126   PetscCall(FormFunction_Water(networkdm, localX, localF, nv, ne, v, e, NULL));
127 
128   PetscCall(DMRestoreLocalVector(networkdm, &localX));
129   PetscCall(DMLocalToGlobalBegin(networkdm, localF, ADD_VALUES, F));
130   PetscCall(DMLocalToGlobalEnd(networkdm, localF, ADD_VALUES, F));
131 
132   PetscCall(DMRestoreLocalVector(networkdm, &localF));
133   PetscFunctionReturn(PETSC_SUCCESS);
134 }
135 
136 PetscErrorCode WaterSetInitialGuess(DM networkdm, Vec X)
137 {
138   PetscInt        nv, ne;
139   const PetscInt *vtx, *edges;
140   Vec             localX;
141 
142   PetscFunctionBegin;
143   PetscCall(DMGetLocalVector(networkdm, &localX));
144 
145   PetscCall(VecSet(X, 0.0));
146   PetscCall(DMGlobalToLocalBegin(networkdm, X, INSERT_VALUES, localX));
147   PetscCall(DMGlobalToLocalEnd(networkdm, X, INSERT_VALUES, localX));
148 
149   /* Get subnetwork */
150   PetscCall(DMNetworkGetSubnetwork(networkdm, 0, &nv, &ne, &vtx, &edges));
151   PetscCall(SetInitialGuess_Water(networkdm, localX, nv, ne, vtx, edges, NULL));
152 
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 }
158 
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;
165 
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     }
180 
181     for (j = 0; j < water->nvertex; j++) {
182       if (water->vertex[j].id == node1) {
183         edgelist[2 * i] = j;
184         break;
185       }
186     }
187 
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 }
197 
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;
204 
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;
211 
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 }
225