1 /* -----------------------------------------------------------------
2  * Programmer(s): Scott D. Cohen, Alan C. Hindmarsh and
3  *                Radu Serban @ LLNL
4  * -----------------------------------------------------------------
5  * SUNDIALS Copyright Start
6  * Copyright (c) 2002-2021, Lawrence Livermore National Security
7  * and Southern Methodist University.
8  * All rights reserved.
9  *
10  * See the top-level LICENSE and NOTICE files for details.
11  *
12  * SPDX-License-Identifier: BSD-3-Clause
13  * SUNDIALS Copyright End
14  * -----------------------------------------------------------------
15  * Example problem:
16  *
17  * An ODE system is generated from the following 2-species diurnal
18  * kinetics advection-diffusion PDE system in 2 space dimensions:
19  *
20  * dc(i)/dt = Kh*(d/dx)^2 c(i) + V*dc(i)/dx + (d/dy)(Kv(y)*dc(i)/dy)
21  *                 + Ri(c1,c2,t)      for i = 1,2,   where
22  *   R1(c1,c2,t) = -q1*c1*c3 - q2*c1*c2 + 2*q3(t)*c3 + q4(t)*c2 ,
23  *   R2(c1,c2,t) =  q1*c1*c3 - q2*c1*c2 - q4(t)*c2 ,
24  *   Kv(y) = Kv0*exp(y/5) ,
25  * Kh, V, Kv0, q1, q2, and c3 are constants, and q3(t) and q4(t)
26  * vary diurnally. The problem is posed on the square
27  *   0 <= x <= 20,    30 <= y <= 50   (all in km),
28  * with homogeneous Neumann boundary conditions, and for time t in
29  *   0 <= t <= 86400 sec (1 day).
30  * The PDE system is treated by central differences on a uniform
31  * 10 x 10 mesh, with simple polynomial initial profiles.
32  * The problem is solved with CVODES, with the BDF/GMRES
33  * method (i.e. using the SUNLinSol_SPGMR linear solver) and the
34  * block-diagonal part of the Newton matrix as a left
35  * preconditioner. A copy of the block-diagonal part of the
36  * Jacobian is saved and conditionally reused within the Precond
37  * routine.
38  * -----------------------------------------------------------------*/
39 
40 #include <stdio.h>
41 #include <stdlib.h>
42 #include <math.h>
43 
44 #include <cvodes/cvodes.h>             /* main integrator header file          */
45 #include <nvector/nvector_serial.h>    /* access to serial N_Vector            */
46 #include <sunlinsol/sunlinsol_spgmr.h> /* access to SPGMR SUNLinearSolver      */
47 #include <sundials/sundials_dense.h>   /* use generic dense solver in precond. */
48 #include <sundials/sundials_types.h>   /* defs. of realtype, sunindextype      */
49 
50 /* helpful macros */
51 
52 #ifndef SQR
53 #define SQR(A) ((A)*(A))
54 #endif
55 
56 /* Problem Constants */
57 
58 #define ZERO RCONST(0.0)
59 #define ONE  RCONST(1.0)
60 #define TWO  RCONST(2.0)
61 
62 #define NUM_SPECIES  2                 /* number of species         */
63 #define KH           RCONST(4.0e-6)    /* horizontal diffusivity Kh */
64 #define VEL          RCONST(0.001)     /* advection velocity V      */
65 #define KV0          RCONST(1.0e-8)    /* coefficient in Kv(y)      */
66 #define Q1           RCONST(1.63e-16)  /* coefficients q1, q2, c3   */
67 #define Q2           RCONST(4.66e-16)
68 #define C3           RCONST(3.7e16)
69 #define A3           RCONST(22.62)     /* coefficient in expression for q3(t) */
70 #define A4           RCONST(7.601)     /* coefficient in expression for q4(t) */
71 #define C1_SCALE     RCONST(1.0e6)     /* coefficients in initial profiles    */
72 #define C2_SCALE     RCONST(1.0e12)
73 
74 #define T0           ZERO                 /* initial time */
75 #define NOUT         12                   /* number of output times */
76 #define TWOHR        RCONST(7200.0)       /* number of seconds in two hours  */
77 #define HALFDAY      RCONST(4.32e4)       /* number of seconds in a half day */
78 #define PI       RCONST(3.1415926535898)  /* pi */
79 
80 #define XMIN         ZERO                 /* grid boundaries in x  */
81 #define XMAX         RCONST(20.0)
82 #define YMIN         RCONST(30.0)         /* grid boundaries in y  */
83 #define YMAX         RCONST(50.0)
84 #define XMID         RCONST(10.0)         /* grid midpoints in x,y */
85 #define YMID         RCONST(40.0)
86 
87 #define MX           10             /* MX = number of x mesh points */
88 #define MY           10             /* MY = number of y mesh points */
89 #define NSMX         20             /* NSMX = NUM_SPECIES*MX */
90 #define MM           (MX*MY)        /* MM = MX*MY */
91 
92 /* CVodeInit Constants */
93 
94 #define RTOL    RCONST(1.0e-5)    /* scalar relative tolerance */
95 #define FLOOR   RCONST(100.0)     /* value of C1 or C2 at which tolerances */
96                                   /* change from relative to absolute      */
97 #define ATOL    (RTOL*FLOOR)      /* scalar absolute tolerance */
98 #define NEQ     (NUM_SPECIES*MM)  /* NEQ = number of equations */
99 
100 /* User-defined vector and matrix accessor macros: IJKth, IJth */
101 
102 /* IJKth is defined in order to isolate the translation from the
103    mathematical 3-dimensional structure of the dependent variable vector
104    to the underlying 1-dimensional storage. IJth is defined in order to
105    write code which indexes into small dense matrices with a (row,column)
106    pair, where 1 <= row, column <= NUM_SPECIES.
107 
108    IJKth(vdata,i,j,k) references the element in the vdata array for
109    species i at mesh point (j,k), where 1 <= i <= NUM_SPECIES,
110    0 <= j <= MX-1, 0 <= k <= MY-1. The vdata array is obtained via
111    the call vdata = N_VGetArrayPointer(v), where v is an N_Vector.
112    For each mesh point (j,k), the elements for species i and i+1 are
113    contiguous within vdata.
114 
115    IJth(a,i,j) references the (i,j)th entry of the small matrix realtype **a,
116    where 1 <= i,j <= NUM_SPECIES. The small matrix routines in sundials_dense.h
117    work with matrices stored by column in a 2-dimensional array. In C,
118    arrays are indexed starting at 0, not 1. */
119 
120 #define IJKth(vdata,i,j,k) (vdata[i-1 + (j)*NUM_SPECIES + (k)*NSMX])
121 #define IJth(a,i,j)        (a[j-1][i-1])
122 
123 /* Type : UserData
124    contains preconditioner blocks, pivot arrays, and problem constants */
125 
126 typedef struct {
127   realtype **P[MX][MY], **Jbd[MX][MY];
128   sunindextype *pivot[MX][MY];
129   realtype q4, om, dx, dy, hdco, haco, vdco;
130 } *UserData;
131 
132 /* Private Helper Functions */
133 
134 static UserData AllocUserData(void);
135 static void InitUserData(UserData data);
136 static void FreeUserData(UserData data);
137 static void SetInitialProfiles(N_Vector u, realtype dx, realtype dy);
138 static void PrintOutput(void *cvode_mem, N_Vector u, realtype t);
139 static void PrintFinalStats(void *cvode_mem);
140 static int check_retval(void *returnvalue, const char *funcname, int opt);
141 
142 /* Functions Called by the Solver */
143 
144 static int f(realtype t, N_Vector u, N_Vector udot, void *user_data);
145 
146 static int jtv(N_Vector v, N_Vector Jv, realtype t,
147                N_Vector y, N_Vector fy,
148                void *user_data, N_Vector tmp);
149 
150 static int Precond(realtype tn, N_Vector u, N_Vector fu,
151                    booleantype jok, booleantype *jcurPtr, realtype gamma,
152                    void *user_data);
153 
154 static int PSolve(realtype tn, N_Vector u, N_Vector fu,
155                   N_Vector r, N_Vector z,
156                   realtype gamma, realtype delta,
157                   int lr, void *user_data);
158 
159 
160 /*
161  *-------------------------------
162  * Main Program
163  *-------------------------------
164  */
165 
main()166 int main()
167 {
168   realtype abstol, reltol, t, tout;
169   N_Vector u;
170   UserData data;
171   SUNLinearSolver LS;
172   void *cvode_mem;
173   int iout, retval;
174 
175   u = NULL;
176   data = NULL;
177   LS = NULL;
178   cvode_mem = NULL;
179 
180   /* Allocate memory, and set problem data, initial values, tolerances */
181   u = N_VNew_Serial(NEQ);
182   if(check_retval((void *)u, "N_VNew_Serial", 0)) return(1);
183   data = AllocUserData();
184   if(check_retval((void *)data, "AllocUserData", 2)) return(1);
185   InitUserData(data);
186   SetInitialProfiles(u, data->dx, data->dy);
187   abstol=ATOL;
188   reltol=RTOL;
189 
190   /* Call CVodeCreate to create the solver memory and specify the
191    * Backward Differentiation Formula */
192   cvode_mem = CVodeCreate(CV_BDF);
193   if(check_retval((void *)cvode_mem, "CVodeCreate", 0)) return(1);
194 
195   /* Set the pointer to user-defined data */
196   retval = CVodeSetUserData(cvode_mem, data);
197   if(check_retval(&retval, "CVodeSetUserData", 1)) return(1);
198 
199   /* Call CVodeInit to initialize the integrator memory and specify the
200    * user's right hand side function in u'=f(t,u), the inital time T0, and
201    * the initial dependent variable vector u. */
202   retval = CVodeInit(cvode_mem, f, T0, u);
203   if(check_retval(&retval, "CVodeInit", 1)) return(1);
204 
205   /* Call CVodeSStolerances to specify the scalar relative tolerance
206    * and scalar absolute tolerances */
207   retval = CVodeSStolerances(cvode_mem, reltol, abstol);
208   if (check_retval(&retval, "CVodeSStolerances", 1)) return(1);
209 
210   /* Call SUNLinSol_SPGMR to specify the linear solver SUNLinSol_SPGMR
211    * with left preconditioning and the default Krylov dimension */
212   LS = SUNLinSol_SPGMR(u, PREC_LEFT, 0);
213   if(check_retval((void *)LS, "SUNLinSol_SPGMR", 0)) return(1);
214 
215   /* Call CVodeSetLinearSolver to attach the linear sovler to CVode */
216   retval = CVodeSetLinearSolver(cvode_mem, LS, NULL);
217   if (check_retval(&retval, "CVodeSetLinearSolver", 1)) return 1;
218 
219   /* set the JAcobian-times-vector function */
220   retval = CVodeSetJacTimes(cvode_mem, NULL, jtv);
221   if(check_retval(&retval, "CVodeSetJacTimes", 1)) return(1);
222 
223   /* Set the preconditioner solve and setup functions */
224   retval = CVodeSetPreconditioner(cvode_mem, Precond, PSolve);
225   if(check_retval(&retval, "CVodeSetPreconditioner", 1)) return(1);
226 
227   /* In loop over output points, call CVode, print results, test for error */
228   printf(" \n2-species diurnal advection-diffusion problem\n\n");
229   for (iout=1, tout = TWOHR; iout <= NOUT; iout++, tout += TWOHR) {
230     retval = CVode(cvode_mem, tout, u, &t, CV_NORMAL);
231     PrintOutput(cvode_mem, u, t);
232     if(check_retval(&retval, "CVode", 1)) break;
233   }
234 
235   PrintFinalStats(cvode_mem);
236 
237   /* Free memory */
238   N_VDestroy(u);
239   FreeUserData(data);
240   CVodeFree(&cvode_mem);
241   SUNLinSolFree(LS);
242 
243   return(0);
244 }
245 
246 /*
247  *-------------------------------
248  * Private helper functions
249  *-------------------------------
250  */
251 
252 /* Allocate memory for data structure of type UserData */
253 
AllocUserData(void)254 static UserData AllocUserData(void)
255 {
256   int jx, jy;
257   UserData data;
258 
259   data = (UserData) malloc(sizeof *data);
260 
261   for (jx=0; jx < MX; jx++) {
262     for (jy=0; jy < MY; jy++) {
263       (data->P)[jx][jy] = newDenseMat(NUM_SPECIES, NUM_SPECIES);
264       (data->Jbd)[jx][jy] = newDenseMat(NUM_SPECIES, NUM_SPECIES);
265       (data->pivot)[jx][jy] = newIndexArray(NUM_SPECIES);
266     }
267   }
268 
269   return(data);
270 }
271 
272 /* Load problem constants in data */
273 
InitUserData(UserData data)274 static void InitUserData(UserData data)
275 {
276   data->om = PI/HALFDAY;
277   data->dx = (XMAX-XMIN)/(MX-1);
278   data->dy = (YMAX-YMIN)/(MY-1);
279   data->hdco = KH/SQR(data->dx);
280   data->haco = VEL/(TWO*data->dx);
281   data->vdco = (ONE/SQR(data->dy))*KV0;
282 }
283 
284 /* Free data memory */
285 
FreeUserData(UserData data)286 static void FreeUserData(UserData data)
287 {
288   int jx, jy;
289 
290   for (jx=0; jx < MX; jx++) {
291     for (jy=0; jy < MY; jy++) {
292       destroyMat((data->P)[jx][jy]);
293       destroyMat((data->Jbd)[jx][jy]);
294       destroyArray((data->pivot)[jx][jy]);
295     }
296   }
297 
298   free(data);
299 }
300 
301 /* Set initial conditions in u */
302 
SetInitialProfiles(N_Vector u,realtype dx,realtype dy)303 static void SetInitialProfiles(N_Vector u, realtype dx, realtype dy)
304 {
305   int jx, jy;
306   realtype x, y, cx, cy;
307   realtype *udata;
308 
309   /* Set pointer to data array in vector u. */
310 
311   udata = N_VGetArrayPointer(u);
312 
313   /* Load initial profiles of c1 and c2 into u vector */
314 
315   for (jy=0; jy < MY; jy++) {
316     y = YMIN + jy*dy;
317     cy = SQR(RCONST(0.1)*(y - YMID));
318     cy = ONE - cy + RCONST(0.5)*SQR(cy);
319     for (jx=0; jx < MX; jx++) {
320       x = XMIN + jx*dx;
321       cx = SQR(RCONST(0.1)*(x - XMID));
322       cx = ONE - cx + RCONST(0.5)*SQR(cx);
323       IJKth(udata,1,jx,jy) = C1_SCALE*cx*cy;
324       IJKth(udata,2,jx,jy) = C2_SCALE*cx*cy;
325     }
326   }
327 }
328 
329 /* Print current t, step count, order, stepsize, and sampled c1,c2 values */
330 
PrintOutput(void * cvode_mem,N_Vector u,realtype t)331 static void PrintOutput(void *cvode_mem, N_Vector u, realtype t)
332 {
333   long int nst;
334   int qu, retval;
335   realtype hu, *udata;
336   int mxh = MX/2 - 1, myh = MY/2 - 1, mx1 = MX - 1, my1 = MY - 1;
337 
338   udata = N_VGetArrayPointer(u);
339 
340   retval = CVodeGetNumSteps(cvode_mem, &nst);
341   check_retval(&retval, "CVodeGetNumSteps", 1);
342   retval = CVodeGetLastOrder(cvode_mem, &qu);
343   check_retval(&retval, "CVodeGetLastOrder", 1);
344   retval = CVodeGetLastStep(cvode_mem, &hu);
345   check_retval(&retval, "CVodeGetLastStep", 1);
346 
347 #if defined(SUNDIALS_EXTENDED_PRECISION)
348   printf("t = %.2Le   no. steps = %ld   order = %d   stepsize = %.2Le\n",
349          t, nst, qu, hu);
350   printf("c1 (bot.left/middle/top rt.) = %12.3Le  %12.3Le  %12.3Le\n",
351          IJKth(udata,1,0,0), IJKth(udata,1,mxh,myh), IJKth(udata,1,mx1,my1));
352   printf("c2 (bot.left/middle/top rt.) = %12.3Le  %12.3Le  %12.3Le\n\n",
353          IJKth(udata,2,0,0), IJKth(udata,2,mxh,myh), IJKth(udata,2,mx1,my1));
354 #elif defined(SUNDIALS_DOUBLE_PRECISION)
355   printf("t = %.2e   no. steps = %ld   order = %d   stepsize = %.2e\n",
356          t, nst, qu, hu);
357   printf("c1 (bot.left/middle/top rt.) = %12.3e  %12.3e  %12.3e\n",
358          IJKth(udata,1,0,0), IJKth(udata,1,mxh,myh), IJKth(udata,1,mx1,my1));
359   printf("c2 (bot.left/middle/top rt.) = %12.3e  %12.3e  %12.3e\n\n",
360          IJKth(udata,2,0,0), IJKth(udata,2,mxh,myh), IJKth(udata,2,mx1,my1));
361 #else
362   printf("t = %.2e   no. steps = %ld   order = %d   stepsize = %.2e\n",
363          t, nst, qu, hu);
364   printf("c1 (bot.left/middle/top rt.) = %12.3e  %12.3e  %12.3e\n",
365          IJKth(udata,1,0,0), IJKth(udata,1,mxh,myh), IJKth(udata,1,mx1,my1));
366   printf("c2 (bot.left/middle/top rt.) = %12.3e  %12.3e  %12.3e\n\n",
367          IJKth(udata,2,0,0), IJKth(udata,2,mxh,myh), IJKth(udata,2,mx1,my1));
368 #endif
369 }
370 
371 /* Get and print final statistics */
372 
PrintFinalStats(void * cvode_mem)373 static void PrintFinalStats(void *cvode_mem)
374 {
375   long int lenrw, leniw ;
376   long int lenrwLS, leniwLS;
377   long int nst, nfe, nsetups, nni, ncfn, netf;
378   long int nli, npe, nps, ncfl, nfeLS;
379   int retval;
380 
381   retval = CVodeGetWorkSpace(cvode_mem, &lenrw, &leniw);
382   check_retval(&retval, "CVodeGetWorkSpace", 1);
383   retval = CVodeGetNumSteps(cvode_mem, &nst);
384   check_retval(&retval, "CVodeGetNumSteps", 1);
385   retval = CVodeGetNumRhsEvals(cvode_mem, &nfe);
386   check_retval(&retval, "CVodeGetNumRhsEvals", 1);
387   retval = CVodeGetNumLinSolvSetups(cvode_mem, &nsetups);
388   check_retval(&retval, "CVodeGetNumLinSolvSetups", 1);
389   retval = CVodeGetNumErrTestFails(cvode_mem, &netf);
390   check_retval(&retval, "CVodeGetNumErrTestFails", 1);
391   retval = CVodeGetNumNonlinSolvIters(cvode_mem, &nni);
392   check_retval(&retval, "CVodeGetNumNonlinSolvIters", 1);
393   retval = CVodeGetNumNonlinSolvConvFails(cvode_mem, &ncfn);
394   check_retval(&retval, "CVodeGetNumNonlinSolvConvFails", 1);
395 
396   retval = CVodeGetLinWorkSpace(cvode_mem, &lenrwLS, &leniwLS);
397   check_retval(&retval, "CVodeGetLinWorkSpace", 1);
398   retval = CVodeGetNumLinIters(cvode_mem, &nli);
399   check_retval(&retval, "CVodeGetNumLinIters", 1);
400   retval = CVodeGetNumPrecEvals(cvode_mem, &npe);
401   check_retval(&retval, "CVodeGetNumPrecEvals", 1);
402   retval = CVodeGetNumPrecSolves(cvode_mem, &nps);
403   check_retval(&retval, "CVodeGetNumPrecSolves", 1);
404   retval = CVodeGetNumLinConvFails(cvode_mem, &ncfl);
405   check_retval(&retval, "CVodeGetNumLinConvFails", 1);
406   retval = CVodeGetNumLinRhsEvals(cvode_mem, &nfeLS);
407   check_retval(&retval, "CVodeGetNumLinRhsEvals", 1);
408 
409   printf("\nFinal Statistics.. \n\n");
410   printf("lenrw   = %5ld     leniw   = %5ld\n"  , lenrw, leniw);
411   printf("lenrwLS = %5ld     leniwLS = %5ld\n"  , lenrwLS, leniwLS);
412   printf("nst     = %5ld\n"                     , nst);
413   printf("nfe     = %5ld     nfeLS   = %5ld\n"  , nfe, nfeLS);
414   printf("nni     = %5ld     nli     = %5ld\n"  , nni, nli);
415   printf("nsetups = %5ld     netf    = %5ld\n"  , nsetups, netf);
416   printf("npe     = %5ld     nps     = %5ld\n"  , npe, nps);
417   printf("ncfn    = %5ld     ncfl    = %5ld\n\n", ncfn, ncfl);
418 }
419 
420 /* Check function return value...
421      opt == 0 means SUNDIALS function allocates memory so check if
422               returned NULL pointer
423      opt == 1 means SUNDIALS function returns an integer value so check if
424               retval < 0
425      opt == 2 means function allocates memory so check if returned
426               NULL pointer */
427 
check_retval(void * returnvalue,const char * funcname,int opt)428 static int check_retval(void *returnvalue, const char *funcname, int opt)
429 {
430   int *retval;
431 
432   /* Check if SUNDIALS function returned NULL pointer - no memory allocated */
433   if (opt == 0 && returnvalue == NULL) {
434     fprintf(stderr, "\nSUNDIALS_ERROR: %s() failed - returned NULL pointer\n\n",
435             funcname);
436     return(1); }
437 
438   /* Check if retval < 0 */
439   else if (opt == 1) {
440     retval = (int *) returnvalue;
441     if (*retval < 0) {
442       fprintf(stderr, "\nSUNDIALS_ERROR: %s() failed with retval = %d\n\n",
443               funcname, *retval);
444       return(1); }}
445 
446   /* Check if function returned NULL pointer - no memory allocated */
447   else if (opt == 2 && returnvalue == NULL) {
448     fprintf(stderr, "\nMEMORY_ERROR: %s() failed - returned NULL pointer\n\n",
449             funcname);
450     return(1); }
451 
452   return(0);
453 }
454 
455 /*
456  *-------------------------------
457  * Functions called by the solver
458  *-------------------------------
459  */
460 
461 /* f routine. Compute RHS function f(t,u). */
462 
f(realtype t,N_Vector u,N_Vector udot,void * user_data)463 static int f(realtype t, N_Vector u, N_Vector udot, void *user_data)
464 {
465   realtype q3, c1, c2, c1dn, c2dn, c1up, c2up, c1lt, c2lt;
466   realtype c1rt, c2rt, cydn, cyup, hord1, hord2, horad1, horad2;
467   realtype qq1, qq2, qq3, qq4, rkin1, rkin2, s, vertd1, vertd2, ydn, yup;
468   realtype q4coef, dely, verdco, hordco, horaco;
469   realtype *udata, *dudata;
470   int jx, jy, idn, iup, ileft, iright;
471   UserData data;
472 
473   data = (UserData) user_data;
474   udata = N_VGetArrayPointer(u);
475   dudata = N_VGetArrayPointer(udot);
476 
477   /* Set diurnal rate coefficients. */
478 
479   s = sin(data->om*t);
480   if (s > ZERO) {
481     q3 = exp(-A3/s);
482     data->q4 = exp(-A4/s);
483   } else {
484       q3 = ZERO;
485       data->q4 = ZERO;
486   }
487 
488   /* Make local copies of problem variables, for efficiency. */
489 
490   q4coef = data->q4;
491   dely = data->dy;
492   verdco = data->vdco;
493   hordco  = data->hdco;
494   horaco  = data->haco;
495 
496   /* Loop over all grid points. */
497 
498   for (jy=0; jy < MY; jy++) {
499 
500     /* Set vertical diffusion coefficients at jy +- 1/2 */
501 
502     ydn = YMIN + (jy - RCONST(0.5))*dely;
503     yup = ydn + dely;
504     cydn = verdco*exp(RCONST(0.2)*ydn);
505     cyup = verdco*exp(RCONST(0.2)*yup);
506     idn = (jy == 0) ? 1 : -1;
507     iup = (jy == MY-1) ? -1 : 1;
508     for (jx=0; jx < MX; jx++) {
509 
510       /* Extract c1 and c2, and set kinetic rate terms. */
511 
512       c1 = IJKth(udata,1,jx,jy);
513       c2 = IJKth(udata,2,jx,jy);
514       qq1 = Q1*c1*C3;
515       qq2 = Q2*c1*c2;
516       qq3 = q3*C3;
517       qq4 = q4coef*c2;
518       rkin1 = -qq1 - qq2 + TWO*qq3 + qq4;
519       rkin2 = qq1 - qq2 - qq4;
520 
521       /* Set vertical diffusion terms. */
522 
523       c1dn = IJKth(udata,1,jx,jy+idn);
524       c2dn = IJKth(udata,2,jx,jy+idn);
525       c1up = IJKth(udata,1,jx,jy+iup);
526       c2up = IJKth(udata,2,jx,jy+iup);
527       vertd1 = cyup*(c1up - c1) - cydn*(c1 - c1dn);
528       vertd2 = cyup*(c2up - c2) - cydn*(c2 - c2dn);
529 
530       /* Set horizontal diffusion and advection terms. */
531 
532       ileft = (jx == 0) ? 1 : -1;
533       iright =(jx == MX-1) ? -1 : 1;
534       c1lt = IJKth(udata,1,jx+ileft,jy);
535       c2lt = IJKth(udata,2,jx+ileft,jy);
536       c1rt = IJKth(udata,1,jx+iright,jy);
537       c2rt = IJKth(udata,2,jx+iright,jy);
538       hord1 = hordco*(c1rt - TWO*c1 + c1lt);
539       hord2 = hordco*(c2rt - TWO*c2 + c2lt);
540       horad1 = horaco*(c1rt - c1lt);
541       horad2 = horaco*(c2rt - c2lt);
542 
543       /* Load all terms into udot. */
544 
545       IJKth(dudata, 1, jx, jy) = vertd1 + hord1 + horad1 + rkin1;
546       IJKth(dudata, 2, jx, jy) = vertd2 + hord2 + horad2 + rkin2;
547     }
548   }
549 
550   return(0);
551 }
552 
553 
554 /* Jacobian-times-vector routine. */
555 
jtv(N_Vector v,N_Vector Jv,realtype t,N_Vector u,N_Vector fu,void * user_data,N_Vector tmp)556 static int jtv(N_Vector v, N_Vector Jv, realtype t,
557                N_Vector u, N_Vector fu,
558                void *user_data, N_Vector tmp)
559 {
560   realtype c1, c2;
561   realtype v1, v2, v1dn, v2dn, v1up, v2up, v1lt, v2lt, v1rt, v2rt;
562   realtype Jv1, Jv2;
563   realtype cydn, cyup;
564   realtype s, ydn, yup;
565   realtype q4coef, dely, verdco, hordco, horaco;
566   int jx, jy, idn, iup, ileft, iright;
567   realtype *udata, *vdata, *Jvdata;
568   UserData data;
569 
570   data = (UserData) user_data;
571 
572   udata = N_VGetArrayPointer(u);
573   vdata = N_VGetArrayPointer(v);
574   Jvdata = N_VGetArrayPointer(Jv);
575 
576   /* Set diurnal rate coefficients. */
577 
578   s = sin(data->om*t);
579   if (s > ZERO) {
580     data->q4 = exp(-A4/s);
581   } else {
582     data->q4 = ZERO;
583   }
584 
585   /* Make local copies of problem variables, for efficiency. */
586 
587   q4coef = data->q4;
588   dely = data->dy;
589   verdco = data->vdco;
590   hordco  = data->hdco;
591   horaco  = data->haco;
592 
593   /* Loop over all grid points. */
594 
595   for (jy=0; jy < MY; jy++) {
596 
597     /* Set vertical diffusion coefficients at jy +- 1/2 */
598 
599     ydn = YMIN + (jy - RCONST(0.5))*dely;
600     yup = ydn + dely;
601 
602     cydn = verdco*exp(RCONST(0.2)*ydn);
603     cyup = verdco*exp(RCONST(0.2)*yup);
604 
605     idn = (jy == 0) ? 1 : -1;
606     iup = (jy == MY-1) ? -1 : 1;
607 
608     for (jx=0; jx < MX; jx++) {
609 
610       Jv1 = ZERO;
611       Jv2 = ZERO;
612 
613       /* Extract c1 and c2 at the current location and at neighbors */
614 
615       c1 = IJKth(udata,1,jx,jy);
616       c2 = IJKth(udata,2,jx,jy);
617 
618       v1 = IJKth(vdata,1,jx,jy);
619       v2 = IJKth(vdata,2,jx,jy);
620 
621       v1dn = IJKth(vdata,1,jx,jy+idn);
622       v2dn = IJKth(vdata,2,jx,jy+idn);
623       v1up = IJKth(vdata,1,jx,jy+iup);
624       v2up = IJKth(vdata,2,jx,jy+iup);
625 
626       ileft = (jx == 0) ? 1 : -1;
627       iright =(jx == MX-1) ? -1 : 1;
628 
629       v1lt = IJKth(vdata,1,jx+ileft,jy);
630       v2lt = IJKth(vdata,2,jx+ileft,jy);
631       v1rt = IJKth(vdata,1,jx+iright,jy);
632       v2rt = IJKth(vdata,2,jx+iright,jy);
633 
634       /* Set kinetic rate terms. */
635 
636       /*
637 	 rkin1 = -Q1*C3 * c1 - Q2 * c1*c2 + q4coef * c2  + TWO*C3*q3;
638 	 rkin2 =  Q1*C3 * c1 - Q2 * c1*c2 - q4coef * c2;
639       */
640 
641       Jv1 += -(Q1*C3 + Q2*c2) * v1  +  (q4coef - Q2*c1) * v2;
642       Jv2 +=  (Q1*C3 - Q2*c2) * v1  -  (q4coef + Q2*c1) * v2;
643 
644       /* Set vertical diffusion terms. */
645 
646       /*
647 	 vertd1 = -(cyup+cydn) * c1 + cyup * c1up + cydn * c1dn;
648 	 vertd2 = -(cyup+cydn) * c2 + cyup * c2up + cydn * c2dn;
649       */
650 
651       Jv1 += -(cyup+cydn) * v1  +  cyup * v1up  +  cydn * v1dn;
652       Jv2 += -(cyup+cydn) * v2  +  cyup * v2up  +  cydn * v2dn;
653 
654       /* Set horizontal diffusion and advection terms. */
655 
656       /*
657 	 hord1 = hordco*(c1rt - TWO*c1 + c1lt);
658 	 hord2 = hordco*(c2rt - TWO*c2 + c2lt);
659       */
660 
661       Jv1 += hordco*(v1rt - TWO*v1 + v1lt);
662       Jv2 += hordco*(v2rt - TWO*v2 + v2lt);
663 
664       /*
665 	 horad1 = horaco*(c1rt - c1lt);
666 	 horad2 = horaco*(c2rt - c2lt);
667       */
668 
669       Jv1 += horaco*(v1rt - v1lt);
670       Jv2 += horaco*(v2rt - v2lt);
671 
672       /* Load two components of J*v */
673 
674       /*
675 	 IJKth(dudata, 1, jx, jy) = vertd1 + hord1 + horad1 + rkin1;
676 	 IJKth(dudata, 2, jx, jy) = vertd2 + hord2 + horad2 + rkin2;
677       */
678 
679       IJKth(Jvdata, 1, jx, jy) = Jv1;
680       IJKth(Jvdata, 2, jx, jy) = Jv2;
681 
682     }
683 
684   }
685 
686   return(0);
687 
688 }
689 
690 
691 /* Preconditioner setup routine. Generate and preprocess P. */
692 
Precond(realtype tn,N_Vector u,N_Vector fu,booleantype jok,booleantype * jcurPtr,realtype gamma,void * user_data)693 static int Precond(realtype tn, N_Vector u, N_Vector fu,
694                    booleantype jok, booleantype *jcurPtr, realtype gamma,
695                    void *user_data)
696 {
697   realtype c1, c2, cydn, cyup, diag, ydn, yup, q4coef, dely, verdco, hordco;
698   realtype **(*P)[MY], **(*Jbd)[MY];
699   sunindextype *(*pivot)[MY], retval;
700   int jx, jy;
701   realtype *udata, **a, **j;
702   UserData data;
703 
704   /* Make local copies of pointers in user_data, and of pointer to u's data */
705 
706   data = (UserData) user_data;
707   P = data->P;
708   Jbd = data->Jbd;
709   pivot = data->pivot;
710   udata = N_VGetArrayPointer(u);
711 
712   if (jok) {
713 
714     /* jok = SUNTRUE: Copy Jbd to P */
715 
716     for (jy=0; jy < MY; jy++)
717       for (jx=0; jx < MX; jx++)
718         denseCopy(Jbd[jx][jy], P[jx][jy], NUM_SPECIES, NUM_SPECIES);
719 
720     *jcurPtr = SUNFALSE;
721 
722   }
723 
724   else {
725     /* jok = SUNFALSE: Generate Jbd from scratch and copy to P */
726 
727     /* Make local copies of problem variables, for efficiency. */
728 
729     q4coef = data->q4;
730     dely = data->dy;
731     verdco = data->vdco;
732     hordco  = data->hdco;
733 
734     /* Compute 2x2 diagonal Jacobian blocks (using q4 values
735        computed on the last f call).  Load into P. */
736 
737     for (jy=0; jy < MY; jy++) {
738       ydn = YMIN + (jy - RCONST(0.5))*dely;
739       yup = ydn + dely;
740       cydn = verdco*exp(RCONST(0.2)*ydn);
741       cyup = verdco*exp(RCONST(0.2)*yup);
742       diag = -(cydn + cyup + TWO*hordco);
743       for (jx=0; jx < MX; jx++) {
744         c1 = IJKth(udata,1,jx,jy);
745         c2 = IJKth(udata,2,jx,jy);
746         j = Jbd[jx][jy];
747         a = P[jx][jy];
748         IJth(j,1,1) = (-Q1*C3 - Q2*c2) + diag;
749         IJth(j,1,2) = -Q2*c1 + q4coef;
750         IJth(j,2,1) = Q1*C3 - Q2*c2;
751         IJth(j,2,2) = (-Q2*c1 - q4coef) + diag;
752         denseCopy(j, a, NUM_SPECIES, NUM_SPECIES);
753       }
754     }
755 
756     *jcurPtr = SUNTRUE;
757 
758   }
759 
760   /* Scale by -gamma */
761 
762   for (jy=0; jy < MY; jy++)
763     for (jx=0; jx < MX; jx++)
764       denseScale(-gamma, P[jx][jy], NUM_SPECIES, NUM_SPECIES);
765 
766   /* Add identity matrix and do LU decompositions on blocks in place. */
767 
768   for (jx=0; jx < MX; jx++) {
769     for (jy=0; jy < MY; jy++) {
770       denseAddIdentity(P[jx][jy], NUM_SPECIES);
771       retval = denseGETRF(P[jx][jy], NUM_SPECIES, NUM_SPECIES, pivot[jx][jy]);
772       if (retval != 0) return(1);
773     }
774   }
775 
776   return(0);
777 }
778 
779 /* Preconditioner solve routine */
780 
PSolve(realtype tn,N_Vector u,N_Vector fu,N_Vector r,N_Vector z,realtype gamma,realtype delta,int lr,void * user_data)781 static int PSolve(realtype tn, N_Vector u, N_Vector fu,
782                   N_Vector r, N_Vector z,
783                   realtype gamma, realtype delta,
784                   int lr, void *user_data)
785 {
786   realtype **(*P)[MY];
787   sunindextype *(*pivot)[MY];
788   int jx, jy;
789   realtype *zdata, *v;
790   UserData data;
791 
792   /* Extract the P and pivot arrays from user_data. */
793 
794   data = (UserData) user_data;
795   P = data->P;
796   pivot = data->pivot;
797   zdata = N_VGetArrayPointer(z);
798 
799   N_VScale(ONE, r, z);
800 
801   /* Solve the block-diagonal system Px = r using LU factors stored
802      in P and pivot data in pivot, and return the solution in z. */
803 
804   for (jx=0; jx < MX; jx++) {
805     for (jy=0; jy < MY; jy++) {
806       v = &(IJKth(zdata, 1, jx, jy));
807       denseGETRS(P[jx][jy], NUM_SPECIES, pivot[jx][jy], v);
808     }
809   }
810 
811   return(0);
812 }
813