xref: /petsc/src/snes/tutorials/network/ex1.c (revision 40badf4fbc550ac1f60bd080eaff6de6d55b946d)
1 static char help[] = "This example demonstrates the use of DMNetwork interface with subnetworks for solving a coupled nonlinear \n\
2                       electric power grid and water pipe problem.\n\
3                       The available solver options are in the ex1options file \n\
4                       and the data files are in the datafiles of subdirectories.\n\
5                       This example shows the use of subnetwork feature in DMNetwork. \n\
6                       Run this program: mpiexec -n <n> ./ex1 \n\\n";
7 
8 /* T
9    Concepts: DMNetwork
10    Concepts: PETSc SNES solver
11 */
12 
13 #include "power/power.h"
14 #include "water/water.h"
15 
16 typedef struct{
17   UserCtx_Power appctx_power;
18   AppCtx_Water  appctx_water;
19   PetscInt      subsnes_id; /* snes solver id */
20   PetscInt      it;         /* iteration number */
21   Vec           localXold;  /* store previous solution, used by FormFunction_Dummy() */
22 } UserCtx;
23 
24 PetscErrorCode UserMonitor(SNES snes,PetscInt its,PetscReal fnorm ,void *appctx)
25 {
26   UserCtx        *user = (UserCtx*)appctx;
27   Vec            X,localXold = user->localXold;
28   DM             networkdm;
29   PetscMPIInt    rank;
30   MPI_Comm       comm;
31 
32   PetscFunctionBegin;
33   CHKERRQ(PetscObjectGetComm((PetscObject)snes,&comm));
34   CHKERRMPI(MPI_Comm_rank(comm,&rank));
35 #if 0
36   if (rank == 0) {
37     PetscInt       subsnes_id = user->subsnes_id;
38     if (subsnes_id == 2) {
39       CHKERRQ(PetscPrintf(PETSC_COMM_SELF," it %D, subsnes_id %D, fnorm %g\n",user->it,user->subsnes_id,(double)fnorm));
40     } else {
41       CHKERRQ(PetscPrintf(PETSC_COMM_SELF,"       subsnes_id %D, fnorm %g\n",user->subsnes_id,(double)fnorm));
42     }
43   }
44 #endif
45   CHKERRQ(SNESGetSolution(snes,&X));
46   CHKERRQ(SNESGetDM(snes,&networkdm));
47   CHKERRQ(DMGlobalToLocalBegin(networkdm,X,INSERT_VALUES,localXold));
48   CHKERRQ(DMGlobalToLocalEnd(networkdm,X,INSERT_VALUES,localXold));
49   PetscFunctionReturn(0);
50 }
51 
52 PetscErrorCode FormJacobian_subPower(SNES snes,Vec X, Mat J,Mat Jpre,void *appctx)
53 {
54   DM             networkdm;
55   Vec            localX;
56   PetscInt       nv,ne,i,j,offset,nvar,row;
57   const PetscInt *vtx,*edges;
58   PetscBool      ghostvtex;
59   PetscScalar    one = 1.0;
60   PetscMPIInt    rank;
61   MPI_Comm       comm;
62 
63   PetscFunctionBegin;
64   CHKERRQ(SNESGetDM(snes,&networkdm));
65   CHKERRQ(DMGetLocalVector(networkdm,&localX));
66 
67   CHKERRQ(PetscObjectGetComm((PetscObject)networkdm,&comm));
68   CHKERRMPI(MPI_Comm_rank(comm,&rank));
69 
70   CHKERRQ(DMGlobalToLocalBegin(networkdm,X,INSERT_VALUES,localX));
71   CHKERRQ(DMGlobalToLocalEnd(networkdm,X,INSERT_VALUES,localX));
72 
73   CHKERRQ(MatZeroEntries(J));
74 
75   /* Power subnetwork: copied from power/FormJacobian_Power() */
76   CHKERRQ(DMNetworkGetSubnetwork(networkdm,0,&nv,&ne,&vtx,&edges));
77   CHKERRQ(FormJacobian_Power_private(networkdm,localX,J,nv,ne,vtx,edges,appctx));
78 
79   /* Water subnetwork: Identity */
80   CHKERRQ(DMNetworkGetSubnetwork(networkdm,1,&nv,&ne,&vtx,&edges));
81   for (i=0; i<nv; i++) {
82     CHKERRQ(DMNetworkIsGhostVertex(networkdm,vtx[i],&ghostvtex));
83     if (ghostvtex) continue;
84 
85     CHKERRQ(DMNetworkGetGlobalVecOffset(networkdm,vtx[i],ALL_COMPONENTS,&offset));
86     CHKERRQ(DMNetworkGetComponent(networkdm,vtx[i],ALL_COMPONENTS,NULL,NULL,&nvar));
87     for (j=0; j<nvar; j++) {
88       row = offset + j;
89       CHKERRQ(MatSetValues(J,1,&row,1,&row,&one,ADD_VALUES));
90     }
91   }
92   CHKERRQ(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY));
93   CHKERRQ(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY));
94 
95   CHKERRQ(DMRestoreLocalVector(networkdm,&localX));
96   PetscFunctionReturn(0);
97 }
98 
99 /* Dummy equation localF(X) = localX - localXold */
100 PetscErrorCode FormFunction_Dummy(DM networkdm,Vec localX, Vec localF,PetscInt nv,PetscInt ne,const PetscInt* vtx,const PetscInt* edges,void* appctx)
101 {
102   const PetscScalar *xarr,*xoldarr;
103   PetscScalar       *farr;
104   PetscInt          i,j,offset,nvar;
105   PetscBool         ghostvtex;
106   UserCtx           *user = (UserCtx*)appctx;
107   Vec               localXold = user->localXold;
108 
109   PetscFunctionBegin;
110   CHKERRQ(VecGetArrayRead(localX,&xarr));
111   CHKERRQ(VecGetArrayRead(localXold,&xoldarr));
112   CHKERRQ(VecGetArray(localF,&farr));
113 
114   for (i=0; i<nv; i++) {
115     CHKERRQ(DMNetworkIsGhostVertex(networkdm,vtx[i],&ghostvtex));
116     if (ghostvtex) continue;
117 
118     CHKERRQ(DMNetworkGetLocalVecOffset(networkdm,vtx[i],ALL_COMPONENTS,&offset));
119     CHKERRQ(DMNetworkGetComponent(networkdm,vtx[i],ALL_COMPONENTS,NULL,NULL,&nvar));
120     for (j=0; j<nvar; j++) {
121       farr[offset+j] = xarr[offset+j] - xoldarr[offset+j];
122     }
123   }
124 
125   CHKERRQ(VecRestoreArrayRead(localX,&xarr));
126   CHKERRQ(VecRestoreArrayRead(localXold,&xoldarr));
127   CHKERRQ(VecRestoreArray(localF,&farr));
128   PetscFunctionReturn(0);
129 }
130 
131 PetscErrorCode FormFunction(SNES snes,Vec X,Vec F,void *appctx)
132 {
133   DM             networkdm;
134   Vec            localX,localF;
135   PetscInt       nv,ne,v;
136   const PetscInt *vtx,*edges;
137   PetscMPIInt    rank;
138   MPI_Comm       comm;
139   UserCtx        *user = (UserCtx*)appctx;
140   UserCtx_Power  appctx_power = (*user).appctx_power;
141   AppCtx_Water   appctx_water = (*user).appctx_water;
142 
143   PetscFunctionBegin;
144   CHKERRQ(SNESGetDM(snes,&networkdm));
145   CHKERRQ(PetscObjectGetComm((PetscObject)networkdm,&comm));
146   CHKERRMPI(MPI_Comm_rank(comm,&rank));
147 
148   CHKERRQ(DMGetLocalVector(networkdm,&localX));
149   CHKERRQ(DMGetLocalVector(networkdm,&localF));
150   CHKERRQ(VecSet(F,0.0));
151   CHKERRQ(VecSet(localF,0.0));
152 
153   CHKERRQ(DMGlobalToLocalBegin(networkdm,X,INSERT_VALUES,localX));
154   CHKERRQ(DMGlobalToLocalEnd(networkdm,X,INSERT_VALUES,localX));
155 
156   /* Form Function for power subnetwork */
157   CHKERRQ(DMNetworkGetSubnetwork(networkdm,0,&nv,&ne,&vtx,&edges));
158   if (user->subsnes_id == 1) { /* snes_water only */
159     CHKERRQ(FormFunction_Dummy(networkdm,localX,localF,nv,ne,vtx,edges,user));
160   } else {
161     CHKERRQ(FormFunction_Power(networkdm,localX,localF,nv,ne,vtx,edges,&appctx_power));
162   }
163 
164   /* Form Function for water subnetwork */
165   CHKERRQ(DMNetworkGetSubnetwork(networkdm,1,&nv,&ne,&vtx,&edges));
166   if (user->subsnes_id == 0) { /* snes_power only */
167     CHKERRQ(FormFunction_Dummy(networkdm,localX,localF,nv,ne,vtx,edges,user));
168   } else {
169     CHKERRQ(FormFunction_Water(networkdm,localX,localF,nv,ne,vtx,edges,NULL));
170   }
171 
172   /* Illustrate how to access the coupling vertex of the subnetworks without doing anything to F yet */
173   CHKERRQ(DMNetworkGetSharedVertices(networkdm,&nv,&vtx));
174   for (v=0; v<nv; v++) {
175     PetscInt       key,ncomp,nvar,nconnedges,k,e,keye,goffset[3];
176     void*          component;
177     const PetscInt *connedges;
178 
179     CHKERRQ(DMNetworkGetComponent(networkdm,vtx[v],ALL_COMPONENTS,NULL,NULL,&nvar));
180     CHKERRQ(DMNetworkGetNumComponents(networkdm,vtx[v],&ncomp));
181     /* printf("  [%d] coupling vertex[%D]: v %D, ncomp %D; nvar %D\n",rank,v,vtx[v], ncomp,nvar); */
182 
183     for (k=0; k<ncomp; k++) {
184       CHKERRQ(DMNetworkGetComponent(networkdm,vtx[v],k,&key,&component,&nvar));
185       CHKERRQ(DMNetworkGetGlobalVecOffset(networkdm,vtx[v],k,&goffset[k]));
186 
187       /* Verify the coupling vertex is a powernet load vertex or a water vertex */
188       switch (k) {
189       case 0:
190         PetscCheckFalse(key != appctx_power.compkey_bus || nvar != 2,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"key %D not a power bus vertex or nvar %D != 2",key,nvar);
191         break;
192       case 1:
193         PetscCheckFalse(key != appctx_power.compkey_load || nvar != 0 || goffset[1] != goffset[0]+2,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Not a power load vertex");
194         break;
195       case 2:
196         PetscCheckFalse(key != appctx_water.compkey_vtx || nvar != 1 || goffset[2] != goffset[1],PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Not a water vertex");
197         break;
198       default: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "k %D is wrong",k);
199       }
200       /* printf("  [%d] coupling vertex[%D]: key %D; nvar %D, goffset %D\n",rank,v,key,nvar,goffset[k]); */
201     }
202 
203     /* Get its supporting edges */
204     CHKERRQ(DMNetworkGetSupportingEdges(networkdm,vtx[v],&nconnedges,&connedges));
205     /* printf("\n[%d] coupling vertex: nconnedges %D\n",rank,nconnedges);CHKERRQ(ierr); */
206     for (k=0; k<nconnedges; k++) {
207       e = connedges[k];
208       CHKERRQ(DMNetworkGetNumComponents(networkdm,e,&ncomp));
209       /* printf("\n  [%d] connected edge[%D]=%D has ncomp %D\n",rank,k,e,ncomp); */
210       CHKERRQ(DMNetworkGetComponent(networkdm,e,0,&keye,&component,NULL));
211       if (keye == appctx_water.compkey_edge) { /* water_compkey_edge */
212         EDGE_Water        edge=(EDGE_Water)component;
213         if (edge->type == EDGE_TYPE_PUMP) {
214           /* printf("  connected edge[%D]=%D has keye=%D, is appctx_water.compkey_edge with EDGE_TYPE_PUMP\n",k,e,keye); */
215         }
216       } else { /* ower->compkey_branch */
217         PetscCheckFalse(keye != appctx_power.compkey_branch,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Not a power branch");
218       }
219     }
220   }
221 
222   CHKERRQ(DMRestoreLocalVector(networkdm,&localX));
223 
224   CHKERRQ(DMLocalToGlobalBegin(networkdm,localF,ADD_VALUES,F));
225   CHKERRQ(DMLocalToGlobalEnd(networkdm,localF,ADD_VALUES,F));
226   CHKERRQ(DMRestoreLocalVector(networkdm,&localF));
227 #if 0
228   if (rank == 0) printf("F:\n");
229   CHKERRQ(VecView(F,PETSC_VIEWER_STDOUT_WORLD));
230 #endif
231   PetscFunctionReturn(0);
232 }
233 
234 PetscErrorCode SetInitialGuess(DM networkdm,Vec X,void* appctx)
235 {
236   PetscInt       nv,ne,i,j,ncomp,offset,key;
237   const PetscInt *vtx,*edges;
238   UserCtx        *user = (UserCtx*)appctx;
239   Vec            localX = user->localXold;
240   UserCtx_Power  appctx_power = (*user).appctx_power;
241   AppCtx_Water   appctx_water = (*user).appctx_water;
242   PetscBool      ghost;
243   PetscScalar    *xarr;
244   VERTEX_Power   bus;
245   VERTEX_Water   vertex;
246   void*          component;
247   GEN            gen;
248 
249   PetscFunctionBegin;
250   CHKERRQ(VecSet(X,0.0));
251   CHKERRQ(VecSet(localX,0.0));
252 
253   /* Set initial guess for power subnetwork */
254   CHKERRQ(DMNetworkGetSubnetwork(networkdm,0,&nv,&ne,&vtx,&edges));
255   CHKERRQ(SetInitialGuess_Power(networkdm,localX,nv,ne,vtx,edges,&appctx_power));
256 
257   /* Set initial guess for water subnetwork */
258   CHKERRQ(DMNetworkGetSubnetwork(networkdm,1,&nv,&ne,&vtx,&edges));
259   CHKERRQ(SetInitialGuess_Water(networkdm,localX,nv,ne,vtx,edges,NULL));
260 
261   /* Set initial guess at the coupling vertex */
262   CHKERRQ(VecGetArray(localX,&xarr));
263   CHKERRQ(DMNetworkGetSharedVertices(networkdm,&nv,&vtx));
264   for (i=0; i<nv; i++) {
265     CHKERRQ(DMNetworkIsGhostVertex(networkdm,vtx[i],&ghost));
266     if (ghost) continue;
267 
268     CHKERRQ(DMNetworkGetNumComponents(networkdm,vtx[i],&ncomp));
269     for (j=0; j<ncomp; j++) {
270       CHKERRQ(DMNetworkGetLocalVecOffset(networkdm,vtx[i],j,&offset));
271       CHKERRQ(DMNetworkGetComponent(networkdm,vtx[i],j,&key,(void**)&component,NULL));
272       if (key == appctx_power.compkey_bus) {
273         bus = (VERTEX_Power)(component);
274         xarr[offset]   = bus->va*PETSC_PI/180.0;
275         xarr[offset+1] = bus->vm;
276       } else if (key == appctx_power.compkey_gen) {
277         gen = (GEN)(component);
278         if (!gen->status) continue;
279         xarr[offset+1] = gen->vs;
280       } else if (key == appctx_water.compkey_vtx) {
281         vertex = (VERTEX_Water)(component);
282         if (vertex->type == VERTEX_TYPE_JUNCTION) {
283           xarr[offset] = 100;
284         } else if (vertex->type == VERTEX_TYPE_RESERVOIR) {
285           xarr[offset] = vertex->res.head;
286         } else {
287           xarr[offset] = vertex->tank.initlvl + vertex->tank.elev;
288         }
289       }
290     }
291   }
292   CHKERRQ(VecRestoreArray(localX,&xarr));
293 
294   CHKERRQ(DMLocalToGlobalBegin(networkdm,localX,ADD_VALUES,X));
295   CHKERRQ(DMLocalToGlobalEnd(networkdm,localX,ADD_VALUES,X));
296   PetscFunctionReturn(0);
297 }
298 
299 int main(int argc,char **argv)
300 {
301   PetscErrorCode   ierr;
302   DM               networkdm;
303   PetscLogStage    stage[4];
304   PetscMPIInt      rank,size;
305   PetscInt         Nsubnet=2,numVertices[2],numEdges[2],i,j,nv,ne,it_max=10;
306   const PetscInt   *vtx,*edges;
307   Vec              X,F;
308   SNES             snes,snes_power,snes_water;
309   Mat              Jac;
310   PetscBool        ghost,viewJ=PETSC_FALSE,viewX=PETSC_FALSE,viewDM=PETSC_FALSE,test=PETSC_FALSE,distribute=PETSC_TRUE,flg;
311   UserCtx          user;
312   SNESConvergedReason reason;
313 
314   /* Power subnetwork */
315   UserCtx_Power       *appctx_power  = &user.appctx_power;
316   char                pfdata_file[PETSC_MAX_PATH_LEN] = "power/case9.m";
317   PFDATA              *pfdata = NULL;
318   PetscInt            genj,loadj,*edgelist_power = NULL,power_netnum;
319   PetscScalar         Sbase = 0.0;
320 
321   /* Water subnetwork */
322   AppCtx_Water        *appctx_water = &user.appctx_water;
323   WATERDATA           *waterdata = NULL;
324   char                waterdata_file[PETSC_MAX_PATH_LEN] = "water/sample1.inp";
325   PetscInt            *edgelist_water = NULL,water_netnum;
326 
327   /* Shared vertices between subnetworks */
328   PetscInt           power_svtx,water_svtx;
329 
330   ierr = PetscInitialize(&argc,&argv,"ex1options",help);if (ierr) return ierr;
331   CHKERRMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&rank));
332   CHKERRMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
333 
334   /* (1) Read Data - Only rank 0 reads the data */
335   /*--------------------------------------------*/
336   CHKERRQ(PetscLogStageRegister("Read Data",&stage[0]));
337   CHKERRQ(PetscLogStagePush(stage[0]));
338 
339   for (i=0; i<Nsubnet; i++) {
340     numVertices[i] = 0;
341     numEdges[i]    = 0;
342   }
343 
344   /* All processes READ THE DATA FOR THE FIRST SUBNETWORK: Electric Power Grid */
345   /* Used for shared vertex, because currently the coupling info must be available in all processes!!! */
346   CHKERRQ(PetscOptionsGetString(NULL,NULL,"-pfdata",pfdata_file,PETSC_MAX_PATH_LEN-1,NULL));
347   CHKERRQ(PetscNew(&pfdata));
348   CHKERRQ(PFReadMatPowerData(pfdata,pfdata_file));
349   if (rank == 0) {
350     CHKERRQ(PetscPrintf(PETSC_COMM_SELF,"Power network: nb = %D, ngen = %D, nload = %D, nbranch = %D\n",pfdata->nbus,pfdata->ngen,pfdata->nload,pfdata->nbranch));
351   }
352   Sbase = pfdata->sbase;
353   if (rank == 0) { /* proc[0] will create Electric Power Grid */
354     numEdges[0]    = pfdata->nbranch;
355     numVertices[0] = pfdata->nbus;
356 
357     CHKERRQ(PetscMalloc1(2*numEdges[0],&edgelist_power));
358     CHKERRQ(GetListofEdges_Power(pfdata,edgelist_power));
359   }
360   /* Broadcast power Sbase to all processors */
361   CHKERRMPI(MPI_Bcast(&Sbase,1,MPIU_SCALAR,0,PETSC_COMM_WORLD));
362   appctx_power->Sbase     = Sbase;
363   appctx_power->jac_error = PETSC_FALSE;
364   /* If external option activated. Introduce error in jacobian */
365   CHKERRQ(PetscOptionsHasName(NULL,NULL, "-jac_error", &appctx_power->jac_error));
366 
367   /* All processes READ THE DATA FOR THE SECOND SUBNETWORK: Water */
368   /* Used for shared vertex, because currently the coupling info must be available in all processes!!! */
369   CHKERRQ(PetscNew(&waterdata));
370   CHKERRQ(PetscOptionsGetString(NULL,NULL,"-waterdata",waterdata_file,PETSC_MAX_PATH_LEN-1,NULL));
371   CHKERRQ(WaterReadData(waterdata,waterdata_file));
372   if (size == 1 || (size > 1 && rank == 1)) {
373     CHKERRQ(PetscCalloc1(2*waterdata->nedge,&edgelist_water));
374     CHKERRQ(GetListofEdges_Water(waterdata,edgelist_water));
375     numEdges[1]    = waterdata->nedge;
376     numVertices[1] = waterdata->nvertex;
377   }
378   PetscLogStagePop();
379 
380   /* (2) Create a network consist of two subnetworks */
381   /*-------------------------------------------------*/
382   CHKERRQ(PetscLogStageRegister("Net Setup",&stage[1]));
383   CHKERRQ(PetscLogStagePush(stage[1]));
384 
385   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-viewDM",&viewDM,NULL));
386   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-test",&test,NULL));
387   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-distribute",&distribute,NULL));
388 
389   /* Create an empty network object */
390   CHKERRQ(DMNetworkCreate(PETSC_COMM_WORLD,&networkdm));
391 
392   /* Register the components in the network */
393   CHKERRQ(DMNetworkRegisterComponent(networkdm,"branchstruct",sizeof(struct _p_EDGE_Power),&appctx_power->compkey_branch));
394   CHKERRQ(DMNetworkRegisterComponent(networkdm,"busstruct",sizeof(struct _p_VERTEX_Power),&appctx_power->compkey_bus));
395   CHKERRQ(DMNetworkRegisterComponent(networkdm,"genstruct",sizeof(struct _p_GEN),&appctx_power->compkey_gen));
396   CHKERRQ(DMNetworkRegisterComponent(networkdm,"loadstruct",sizeof(struct _p_LOAD),&appctx_power->compkey_load));
397 
398   CHKERRQ(DMNetworkRegisterComponent(networkdm,"edge_water",sizeof(struct _p_EDGE_Water),&appctx_water->compkey_edge));
399   CHKERRQ(DMNetworkRegisterComponent(networkdm,"vertex_water",sizeof(struct _p_VERTEX_Water),&appctx_water->compkey_vtx));
400 #if 0
401   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"power->compkey_branch %d\n",appctx_power->compkey_branch));
402   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"power->compkey_bus    %d\n",appctx_power->compkey_bus));
403   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"power->compkey_gen    %d\n",appctx_power->compkey_gen));
404   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"power->compkey_load   %d\n",appctx_power->compkey_load));
405   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"water->compkey_edge   %d\n",appctx_water->compkey_edge));
406   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"water->compkey_vtx    %d\n",appctx_water->compkey_vtx));
407 #endif
408   CHKERRQ(PetscSynchronizedPrintf(PETSC_COMM_WORLD,"[%d] Total local nvertices %D + %D = %D, nedges %D + %D = %D\n",rank,numVertices[0],numVertices[1],numVertices[0]+numVertices[1],numEdges[0],numEdges[1],numEdges[0]+numEdges[1]));
409   CHKERRQ(PetscSynchronizedFlush(PETSC_COMM_WORLD,PETSC_STDOUT));
410 
411   CHKERRQ(DMNetworkSetNumSubNetworks(networkdm,PETSC_DECIDE,Nsubnet));
412   CHKERRQ(DMNetworkAddSubnetwork(networkdm,"power",numEdges[0],edgelist_power,&power_netnum));
413   CHKERRQ(DMNetworkAddSubnetwork(networkdm,"water",numEdges[1],edgelist_water,&water_netnum));
414 
415   /* vertex subnet[0].4 shares with vertex subnet[1].0 */
416   power_svtx = 4; water_svtx = 0;
417   CHKERRQ(DMNetworkAddSharedVertices(networkdm,power_netnum,water_netnum,1,&power_svtx,&water_svtx));
418 
419   /* Set up the network layout */
420   CHKERRQ(DMNetworkLayoutSetUp(networkdm));
421 
422   /* ADD VARIABLES AND COMPONENTS FOR THE POWER SUBNETWORK */
423   /*-------------------------------------------------------*/
424   genj = 0; loadj = 0;
425   CHKERRQ(DMNetworkGetSubnetwork(networkdm,power_netnum,&nv,&ne,&vtx,&edges));
426 
427   for (i = 0; i < ne; i++) {
428     CHKERRQ(DMNetworkAddComponent(networkdm,edges[i],appctx_power->compkey_branch,&pfdata->branch[i],0));
429   }
430 
431   for (i = 0; i < nv; i++) {
432     CHKERRQ(DMNetworkIsSharedVertex(networkdm,vtx[i],&flg));
433     if (flg) continue;
434 
435     CHKERRQ(DMNetworkAddComponent(networkdm,vtx[i],appctx_power->compkey_bus,&pfdata->bus[i],2));
436     if (pfdata->bus[i].ngen) {
437       for (j = 0; j < pfdata->bus[i].ngen; j++) {
438         CHKERRQ(DMNetworkAddComponent(networkdm,vtx[i],appctx_power->compkey_gen,&pfdata->gen[genj++],0));
439       }
440     }
441     if (pfdata->bus[i].nload) {
442       for (j=0; j < pfdata->bus[i].nload; j++) {
443         CHKERRQ(DMNetworkAddComponent(networkdm,vtx[i],appctx_power->compkey_load,&pfdata->load[loadj++],0));
444       }
445     }
446   }
447 
448   /* ADD VARIABLES AND COMPONENTS FOR THE WATER SUBNETWORK */
449   /*-------------------------------------------------------*/
450   CHKERRQ(DMNetworkGetSubnetwork(networkdm,water_netnum,&nv,&ne,&vtx,&edges));
451   for (i = 0; i < ne; i++) {
452     CHKERRQ(DMNetworkAddComponent(networkdm,edges[i],appctx_water->compkey_edge,&waterdata->edge[i],0));
453   }
454 
455   for (i = 0; i < nv; i++) {
456     CHKERRQ(DMNetworkIsSharedVertex(networkdm,vtx[i],&flg));
457     if (flg) continue;
458 
459     CHKERRQ(DMNetworkAddComponent(networkdm,vtx[i],appctx_water->compkey_vtx,&waterdata->vertex[i],1));
460   }
461 
462   /* ADD VARIABLES AND COMPONENTS AT THE SHARED VERTEX: net[0].4 coupls with net[1].0 -- owning and all ghost ranks of the vertex do this */
463   /*----------------------------------------------------------------------------------------------------------------------------*/
464   CHKERRQ(DMNetworkGetSharedVertices(networkdm,&nv,&vtx));
465   for (i = 0; i < nv; i++) {
466     /* power */
467     CHKERRQ(DMNetworkAddComponent(networkdm,vtx[i],appctx_power->compkey_bus,&pfdata->bus[4],2));
468     /* bus[4] is a load, add its component */
469     CHKERRQ(DMNetworkAddComponent(networkdm,vtx[i],appctx_power->compkey_load,&pfdata->load[0],0));
470 
471     /* water */
472     CHKERRQ(DMNetworkAddComponent(networkdm,vtx[i],appctx_water->compkey_vtx,&waterdata->vertex[0],1));
473   }
474 
475   /* Set up DM for use */
476   CHKERRQ(DMSetUp(networkdm));
477   if (viewDM) {
478     CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"\nAfter DMSetUp, DMView:\n"));
479     CHKERRQ(DMView(networkdm,PETSC_VIEWER_STDOUT_WORLD));
480   }
481 
482   /* Free user objects */
483   CHKERRQ(PetscFree(edgelist_power));
484   CHKERRQ(PetscFree(pfdata->bus));
485   CHKERRQ(PetscFree(pfdata->gen));
486   CHKERRQ(PetscFree(pfdata->branch));
487   CHKERRQ(PetscFree(pfdata->load));
488   CHKERRQ(PetscFree(pfdata));
489 
490   CHKERRQ(PetscFree(edgelist_water));
491   CHKERRQ(PetscFree(waterdata->vertex));
492   CHKERRQ(PetscFree(waterdata->edge));
493   CHKERRQ(PetscFree(waterdata));
494 
495   /* Re-distribute networkdm to multiple processes for better job balance */
496   if (size >1 && distribute) {
497     CHKERRQ(DMNetworkDistribute(&networkdm,0));
498     if (viewDM) {
499       CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"\nAfter DMNetworkDistribute, DMView:\n"));
500       CHKERRQ(DMView(networkdm,PETSC_VIEWER_STDOUT_WORLD));
501     }
502   }
503 
504   /* Test DMNetworkGetSubnetwork() and DMNetworkGetSubnetworkSharedVertices() */
505   if (test) {
506     PetscInt  v,gidx;
507     CHKERRMPI(MPI_Barrier(PETSC_COMM_WORLD));
508     for (i=0; i<Nsubnet; i++) {
509       CHKERRQ(DMNetworkGetSubnetwork(networkdm,i,&nv,&ne,&vtx,&edges));
510       CHKERRQ(PetscPrintf(PETSC_COMM_SELF,"[%d] After distribute, subnet[%d] ne %d, nv %d\n",rank,i,ne,nv));
511       CHKERRMPI(MPI_Barrier(PETSC_COMM_WORLD));
512 
513       for (v=0; v<nv; v++) {
514         CHKERRQ(DMNetworkIsGhostVertex(networkdm,vtx[v],&ghost));
515         CHKERRQ(DMNetworkGetGlobalVertexIndex(networkdm,vtx[v],&gidx));
516         CHKERRQ(PetscPrintf(PETSC_COMM_SELF,"[%d] subnet[%d] v %d %d; ghost %d\n",rank,i,vtx[v],gidx,ghost));
517       }
518       CHKERRMPI(MPI_Barrier(PETSC_COMM_WORLD));
519     }
520     CHKERRMPI(MPI_Barrier(PETSC_COMM_WORLD));
521 
522     CHKERRQ(DMNetworkGetSharedVertices(networkdm,&nv,&vtx));
523     CHKERRQ(PetscPrintf(PETSC_COMM_SELF,"[%d] After distribute, num of shared vertices nsv = %d\n",rank,nv));
524     for (v=0; v<nv; v++) {
525       CHKERRQ(DMNetworkGetGlobalVertexIndex(networkdm,vtx[v],&gidx));
526       CHKERRQ(PetscPrintf(PETSC_COMM_SELF,"[%d] sv %d, gidx=%d\n",rank,vtx[v],gidx));
527     }
528     CHKERRMPI(MPI_Barrier(PETSC_COMM_WORLD));
529   }
530 
531   /* Create solution vector X */
532   CHKERRQ(DMCreateGlobalVector(networkdm,&X));
533   CHKERRQ(VecDuplicate(X,&F));
534   CHKERRQ(DMGetLocalVector(networkdm,&user.localXold));
535   PetscLogStagePop();
536 
537   /* (3) Setup Solvers */
538   /*-------------------*/
539   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-viewJ",&viewJ,NULL));
540   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-viewX",&viewX,NULL));
541 
542   CHKERRQ(PetscLogStageRegister("SNES Setup",&stage[2]));
543   CHKERRQ(PetscLogStagePush(stage[2]));
544 
545   CHKERRQ(SetInitialGuess(networkdm,X,&user));
546 
547   /* Create coupled snes */
548   /*-------------------- */
549   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"SNES_coupled setup ......\n"));
550   user.subsnes_id = Nsubnet;
551   CHKERRQ(SNESCreate(PETSC_COMM_WORLD,&snes));
552   CHKERRQ(SNESSetDM(snes,networkdm));
553   CHKERRQ(SNESSetOptionsPrefix(snes,"coupled_"));
554   CHKERRQ(SNESSetFunction(snes,F,FormFunction,&user));
555   CHKERRQ(SNESMonitorSet(snes,UserMonitor,&user,NULL));
556   CHKERRQ(SNESSetFromOptions(snes));
557 
558   if (viewJ) {
559     /* View Jac structure */
560     CHKERRQ(SNESGetJacobian(snes,&Jac,NULL,NULL,NULL));
561     CHKERRQ(MatView(Jac,PETSC_VIEWER_DRAW_WORLD));
562   }
563 
564   if (viewX) {
565     CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Solution:\n"));
566     CHKERRQ(VecView(X,PETSC_VIEWER_STDOUT_WORLD));
567   }
568 
569   if (viewJ) {
570     /* View assembled Jac */
571     CHKERRQ(MatView(Jac,PETSC_VIEWER_DRAW_WORLD));
572   }
573 
574   /* Create snes_power */
575   /*-------------------*/
576   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"SNES_power setup ......\n"));
577 
578   user.subsnes_id = 0;
579   CHKERRQ(SNESCreate(PETSC_COMM_WORLD,&snes_power));
580   CHKERRQ(SNESSetDM(snes_power,networkdm));
581   CHKERRQ(SNESSetOptionsPrefix(snes_power,"power_"));
582   CHKERRQ(SNESSetFunction(snes_power,F,FormFunction,&user));
583   CHKERRQ(SNESMonitorSet(snes_power,UserMonitor,&user,NULL));
584 
585   /* Use user-provide Jacobian */
586   CHKERRQ(DMCreateMatrix(networkdm,&Jac));
587   CHKERRQ(SNESSetJacobian(snes_power,Jac,Jac,FormJacobian_subPower,&user));
588   CHKERRQ(SNESSetFromOptions(snes_power));
589 
590   if (viewX) {
591     CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Power Solution:\n"));
592     CHKERRQ(VecView(X,PETSC_VIEWER_STDOUT_WORLD));
593   }
594   if (viewJ) {
595     CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Power Jac:\n"));
596     CHKERRQ(SNESGetJacobian(snes_power,&Jac,NULL,NULL,NULL));
597     CHKERRQ(MatView(Jac,PETSC_VIEWER_DRAW_WORLD));
598     /* CHKERRQ(MatView(Jac,PETSC_VIEWER_STDOUT_WORLD)); */
599   }
600 
601   /* Create snes_water */
602   /*-------------------*/
603   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"SNES_water setup......\n"));
604 
605   user.subsnes_id = 1;
606   CHKERRQ(SNESCreate(PETSC_COMM_WORLD,&snes_water));
607   CHKERRQ(SNESSetDM(snes_water,networkdm));
608   CHKERRQ(SNESSetOptionsPrefix(snes_water,"water_"));
609   CHKERRQ(SNESSetFunction(snes_water,F,FormFunction,&user));
610   CHKERRQ(SNESMonitorSet(snes_water,UserMonitor,&user,NULL));
611   CHKERRQ(SNESSetFromOptions(snes_water));
612 
613   if (viewX) {
614     CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Water Solution:\n"));
615     CHKERRQ(VecView(X,PETSC_VIEWER_STDOUT_WORLD));
616   }
617   CHKERRQ(PetscLogStagePop());
618 
619   /* (4) Solve */
620   /*-----------*/
621   CHKERRQ(PetscLogStageRegister("SNES Solve",&stage[3]));
622   CHKERRQ(PetscLogStagePush(stage[3]));
623   user.it = 0;
624   reason  = SNES_DIVERGED_DTOL;
625   while (user.it < it_max && (PetscInt)reason<0) {
626 #if 0
627     user.subsnes_id = 0;
628     CHKERRQ(SNESSolve(snes_power,NULL,X));
629 
630     user.subsnes_id = 1;
631     CHKERRQ(SNESSolve(snes_water,NULL,X));
632 #endif
633     user.subsnes_id = Nsubnet;
634     CHKERRQ(SNESSolve(snes,NULL,X));
635 
636     CHKERRQ(SNESGetConvergedReason(snes,&reason));
637     user.it++;
638   }
639   CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Coupled_SNES converged in %D iterations\n",user.it));
640   if (viewX) {
641     CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Final Solution:\n"));
642     CHKERRQ(VecView(X,PETSC_VIEWER_STDOUT_WORLD));
643   }
644   CHKERRQ(PetscLogStagePop());
645 
646   /* Free objects */
647   /* -------------*/
648   CHKERRQ(VecDestroy(&X));
649   CHKERRQ(VecDestroy(&F));
650   CHKERRQ(DMRestoreLocalVector(networkdm,&user.localXold));
651 
652   CHKERRQ(SNESDestroy(&snes));
653   CHKERRQ(MatDestroy(&Jac));
654   CHKERRQ(SNESDestroy(&snes_power));
655   CHKERRQ(SNESDestroy(&snes_water));
656 
657   CHKERRQ(DMDestroy(&networkdm));
658   ierr = PetscFinalize();
659   return ierr;
660 }
661 
662 /*TEST
663 
664    build:
665      requires: !complex double defined(PETSC_HAVE_ATTRIBUTEALIGNED)
666      depends: power/PFReadData.c power/pffunctions.c water/waterreaddata.c water/waterfunctions.c
667 
668    test:
669       args: -coupled_snes_converged_reason -options_left no -viewDM
670       localrunfiles: ex1options power/case9.m water/sample1.inp
671       output_file: output/ex1.out
672 
673    test:
674       suffix: 2
675       nsize: 3
676       args: -coupled_snes_converged_reason -options_left no -petscpartitioner_type parmetis
677       localrunfiles: ex1options power/case9.m water/sample1.inp
678       output_file: output/ex1_2.out
679       requires: parmetis
680 
681 #   test:
682 #      suffix: 3
683 #      nsize: 3
684 #      args: -coupled_snes_converged_reason -options_left no -distribute false
685 #      localrunfiles: ex1options power/case9.m water/sample1.inp
686 #      output_file: output/ex1_2.out
687 
688    test:
689       suffix: 4
690       nsize: 4
691       args: -coupled_snes_converged_reason -options_left no -petscpartitioner_type simple -viewDM
692       localrunfiles: ex1options power/case9.m water/sample1.inp
693       output_file: output/ex1_4.out
694 
695 TEST*/
696