Actual source code: cn.c

  1: /*$Id: cn.c,v 1.29 2001/04/04 16:56:35 balay Exp $*/
  2: /*
  3:        Code for Timestepping with implicit Crank-Nicholson method.
  4:     THIS IS NOT YET COMPLETE -- DO NOT USE!!
  5: */
 6:  #include src/ts/tsimpl.h

  8: typedef struct {
  9:   Vec  update;      /* work vector where new solution is formed */
 10:   Vec  func;        /* work vector where F(t[i],u[i]) is stored */
 11:   Vec  rhs;         /* work vector for RHS; vec_sol/dt */
 12: } TS_CN;

 14: /*------------------------------------------------------------------------------*/
 15: /*
 16:    TSComputeRHSFunctionEuler - Evaluates the right-hand-side function. 

 18:    Note: If the user did not provide a function but merely a matrix,
 19:    this routine applies the matrix.
 20: */
 21: int TSComputeRHSFunctionEuler(TS ts,double t,Vec x,Vec y)
 22: {
 23:   int    ierr;
 24:   Scalar neg_two = -2.0,neg_mdt = -1.0/ts->time_step;


 30:   if (ts->rhsfunction) {
 31:     PetscStackPush("TS user right-hand-side function");
 32:     (*ts->rhsfunction)(ts,t,x,y,ts->funP);
 33:     PetscStackPop;
 34:     return(0);
 35:   }

 37:   if (ts->rhsmatrix) { /* assemble matrix for this timestep */
 38:     MatStructure flg;
 39:     PetscStackPush("TS user right-hand-side matrix function");
 40:     (*ts->rhsmatrix)(ts,t,&ts->A,&ts->B,&flg,ts->jacP);
 41:     PetscStackPop;
 42:   }
 43:   MatMult(ts->A,x,y);
 44:   /* shift: y = y -2*x */
 45:   VecAXPY(&neg_two,x,y);
 46:   /* scale: y = y -2*x */
 47:   VecScale(&neg_mdt,y);

 49:   /* apply user-provided boundary conditions (only needed if these are time dependent) */
 50:   TSComputeRHSBoundaryConditions(ts,t,y);

 52:   return(0);
 53: }

 55: /*
 56:     Version for linear PDE where RHS does not depend on time. Has built a
 57:   single matrix that is to be used for all timesteps.
 58: */
 59: static int TSStep_CN_Linear_Constant_Matrix(TS ts,int *steps,double *ptime)
 60: {
 61:   TS_CN     *cn = (TS_CN*)ts->data;
 62:   Vec       sol = ts->vec_sol,update = cn->update;
 63:   Vec       rhs = cn->rhs;
 64:   int       ierr,i,max_steps = ts->max_steps,its;
 65:   Scalar    dt = ts->time_step,two = 2.0;
 66: 
 68:   *steps = -ts->steps;
 69:   TSMonitor(ts,ts->steps,ts->ptime,sol);

 71:   /* set initial guess to be previous solution */
 72:   VecCopy(sol,update);

 74:   for (i=0; i<max_steps; i++) {
 75:     ts->ptime += ts->time_step;
 76:     if (ts->ptime > ts->max_time) break;

 78:     /* phase 1 - explicit step */
 79:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
 80:     VecAXPBY(&dt,&two,update,sol);

 82:     /* phase 2 - implicit step */
 83:     VecCopy(sol,rhs);
 84:     /* apply user-provided boundary conditions (only needed if they are time dependent) */
 85:     TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);

 87:     SLESSolve(ts->sles,rhs,update,&its);
 88:     ts->linear_its += PetscAbsInt(its);
 89:     VecCopy(update,sol);
 90:     ts->steps++;
 91:     TSMonitor(ts,ts->steps,ts->ptime,sol);
 92:   }

 94:   *steps += ts->steps;
 95:   *ptime  = ts->ptime;
 96:   return(0);
 97: }
 98: /*
 99:       Version where matrix depends on time 
100: */
101: static int TSStep_CN_Linear_Variable_Matrix(TS ts,int *steps,double *ptime)
102: {
103:   TS_CN        *cn = (TS_CN*)ts->data;
104:   Vec          sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
105:   int          ierr,i,max_steps = ts->max_steps,its;
106:   Scalar       dt = ts->time_step,two = 2.0,neg_dt = -1.0*ts->time_step;
107:   MatStructure str;

110:   *steps = -ts->steps;
111:   TSMonitor(ts,ts->steps,ts->ptime,sol);

113:   /* set initial guess to be previous solution */
114:   VecCopy(sol,update);

116:   for (i=0; i<max_steps; i++) {
117:     ts->ptime += ts->time_step;
118:     if (ts->ptime > ts->max_time) break;
119:     /*
120:         evaluate matrix function 
121:     */
122:     (*ts->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
123:     if (!ts->Ashell) {
124:       MatScale(&neg_dt,ts->A);
125:       MatShift(&two,ts->A);
126:     }
127:     if (ts->B != ts->A && ts->Ashell != ts->B && str != SAME_PRECONDITIONER) {
128:       MatScale(&neg_dt,ts->B);
129:       MatShift(&two,ts->B);
130:     }

132:     /* phase 1 - explicit step */
133:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
134:     VecAXPBY(&dt,&two,update,sol);

136:     /* phase 2 - implicit step */
137:     VecCopy(sol,rhs);

139:     /* apply user-provided boundary conditions (only needed if they are time dependent) */
140:     TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);

142:     SLESSetOperators(ts->sles,ts->A,ts->B,str);
143:     SLESSolve(ts->sles,rhs,update,&its);
144:     ts->linear_its += PetscAbsInt(its);
145:     VecCopy(update,sol);
146:     ts->steps++;
147:     TSMonitor(ts,ts->steps,ts->ptime,sol);
148:   }

150:   *steps += ts->steps;
151:   *ptime  = ts->ptime;
152:   return(0);
153: }
154: /*
155:     Version for nonlinear PDE.
156: */
157: static int TSStep_CN_Nonlinear(TS ts,int *steps,double *ptime)
158: {
159:   Vec       sol = ts->vec_sol;
160:   int       ierr,i,max_steps = ts->max_steps,its,lits;
161:   TS_CN *cn = (TS_CN*)ts->data;
162: 
164:   *steps = -ts->steps;
165:   TSMonitor(ts,ts->steps,ts->ptime,sol);

167:   for (i=0; i<max_steps; i++) {
168:     ts->ptime += ts->time_step;
169:     if (ts->ptime > ts->max_time) break;
170:     VecCopy(sol,cn->update);
171:     SNESSolve(ts->snes,cn->update,&its);
172:     SNESGetNumberLinearIterations(ts->snes,&lits);
173:     ts->nonlinear_its += PetscAbsInt(its); ts->linear_its += lits;
174:     VecCopy(cn->update,sol);
175:     ts->steps++;
176:     TSMonitor(ts,ts->steps,ts->ptime,sol);
177:   }

179:   *steps += ts->steps;
180:   *ptime  = ts->ptime;
181:   return(0);
182: }

184: /*------------------------------------------------------------*/
185: static int TSDestroy_CN(TS ts)
186: {
187:   TS_CN *cn = (TS_CN*)ts->data;
188:   int       ierr;

191:   if (cn->update) {VecDestroy(cn->update);}
192:   if (cn->func) {VecDestroy(cn->func);}
193:   if (cn->rhs) {VecDestroy(cn->rhs);}
194:   if (ts->Ashell) {MatDestroy(ts->A);}
195:   PetscFree(cn);
196:   return(0);
197: }


200: /*------------------------------------------------------------*/
201: /*
202:     This matrix shell multiply where user provided Shell matrix
203: */

205: int TSCnMatMult(Mat mat,Vec x,Vec y)
206: {
207:   TS     ts;
208:   Scalar two = 2.0,neg_dt;
209:   int    ierr;

212:   MatShellGetContext(mat,(void **)&ts);
213:   neg_dt = -1.0*ts->time_step;

215:   /* apply user-provided function */
216:   MatMult(ts->Ashell,x,y);
217:   /* shift and scale by 2 - dt*F */
218:   VecAXPBY(&two,&neg_dt,x,y);
219:   return(0);
220: }

222: /* 
223:     This defines the nonlinear equation that is to be solved with SNES

225:               U^{n+1} - dt*F(U^{n+1}) - U^{n}
226: */
227: int TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
228: {
229:   TS     ts = (TS) ctx;
230:   Scalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
231:   int    ierr,i,n;

234:   /* apply user provided function */
235:   TSComputeRHSFunction(ts,ts->ptime,x,y);
236:   /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
237:   VecGetArray(ts->vec_sol,&un);
238:   VecGetArray(x,&unp1);
239:   VecGetArray(y,&Funp1);
240:   VecGetLocalSize(x,&n);

242:   for (i=0; i<n; i++) {
243:     Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
244:   }
245:   VecRestoreArray(ts->vec_sol,&un);
246:   VecRestoreArray(x,&unp1);
247:   VecRestoreArray(y,&Funp1);
248:   return(0);
249: }

251: /*
252:    This constructs the Jacobian needed for SNES 

254:              J = I/dt - J_{F}   where J_{F} is the given Jacobian of F.
255: */
256: int TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
257: {
258:   TS         ts = (TS) ctx;
259:   int        ierr;
260:   Scalar     mone = -1.0,mdt = 1.0/ts->time_step;
261:   PetscTruth isshell;

264:   /* construct user's Jacobian */
265:   TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str);

267:   /* shift and scale Jacobian, if not matrix-free */
268:   PetscTypeCompare((PetscObject)*AA,MATSHELL,&isshell);
269:   if (!isshell) {
270:     MatScale(&mone,*AA);
271:     MatShift(&mdt,*AA);
272:   }
273:   PetscTypeCompare((PetscObject)*BB,MATSHELL,&isshell);
274:   if (*BB != *AA && *str != SAME_PRECONDITIONER && !isshell) {
275:     MatScale(&mone,*BB);
276:     MatShift(&mdt,*BB);
277:   }

279:   return(0);
280: }

282: /* ------------------------------------------------------------*/
283: static int TSSetUp_CN_Linear_Constant_Matrix(TS ts)
284: {
285:   TS_CN   *cn = (TS_CN*)ts->data;
286:   int     ierr,M,m;
287:   Scalar  two = 2.0,neg_dt = -1.0*ts->time_step;

290:   VecDuplicate(ts->vec_sol,&cn->update);
291:   VecDuplicate(ts->vec_sol,&cn->rhs);
292: 
293:   /* build linear system to be solved */
294:   if (!ts->Ashell) {
295:     MatScale(&neg_dt,ts->A);
296:     MatShift(&two,ts->A);
297:   } else {
298:     /* construct new shell matrix */
299:     VecGetSize(ts->vec_sol,&M);
300:     VecGetLocalSize(ts->vec_sol,&m);
301:     MatCreateShell(ts->comm,m,M,M,M,ts,&ts->A);
302:     MatShellSetOperation(ts->A,MATOP_MULT,(void(*)())TSCnMatMult);
303:   }
304:   if (ts->A != ts->B && ts->Ashell != ts->B) {
305:     MatScale(&neg_dt,ts->B);
306:     MatShift(&two,ts->B);
307:   }
308:   SLESSetOperators(ts->sles,ts->A,ts->B,SAME_NONZERO_PATTERN);
309:   return(0);
310: }

312: static int TSSetUp_CN_Linear_Variable_Matrix(TS ts)
313: {
314:   TS_CN *cn = (TS_CN*)ts->data;
315:   int       ierr,M,m;

318:   VecDuplicate(ts->vec_sol,&cn->update);
319:   VecDuplicate(ts->vec_sol,&cn->rhs);
320:   if (ts->Ashell) { /* construct new shell matrix */
321:     VecGetSize(ts->vec_sol,&M);
322:     VecGetLocalSize(ts->vec_sol,&m);
323:     MatCreateShell(ts->comm,m,M,M,M,ts,&ts->A);
324:     MatShellSetOperation(ts->A,MATOP_MULT,(void(*)())TSCnMatMult);
325:   }
326:   return(0);
327: }

329: static int TSSetUp_CN_Nonlinear(TS ts)
330: {
331:   TS_CN *cn = (TS_CN*)ts->data;
332:   int       ierr,M,m;

335:   VecDuplicate(ts->vec_sol,&cn->update);
336:   VecDuplicate(ts->vec_sol,&cn->func);
337:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
338:   if (ts->Ashell) { /* construct new shell matrix */
339:     VecGetSize(ts->vec_sol,&M);
340:     VecGetLocalSize(ts->vec_sol,&m);
341:     MatCreateShell(ts->comm,m,M,M,M,ts,&ts->A);
342:     MatShellSetOperation(ts->A,MATOP_MULT,(void(*)())TSCnMatMult);
343:   }
344:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
345:   return(0);
346: }
347: /*------------------------------------------------------------*/

349: static int TSSetFromOptions_CN_Linear(TS ts)
350: {

354:   SLESSetFromOptions(ts->sles);
355: 
356:   return(0);
357: }

359: static int TSSetFromOptions_CN_Nonlinear(TS ts)
360: {

364:   SNESSetFromOptions(ts->snes);
365: 
366:   return(0);
367: }

369: static int TSView_CN(TS ts,PetscViewer viewer)
370: {
372:   return(0);
373: }

375: /* ------------------------------------------------------------ */
376: EXTERN_C_BEGIN
377: int TSCreate_CN(TS ts)
378: {
379:   TS_CN      *cn;
380:   int        ierr;
381:   KSP        ksp;
382:   PetscTruth isshell;

385:   ts->destroy         = TSDestroy_CN;
386:   ts->view            = TSView_CN;

388:   if (ts->problem_type == TS_LINEAR) {
389:     if (!ts->A) {
390:       SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
391:     }
392:     PetscTypeCompare((PetscObject)ts->A,MATSHELL,&isshell);
393:     if (!ts->rhsmatrix) {
394:       if (isshell) {
395:         ts->Ashell = ts->A;
396:       }
397:       ts->setup  = TSSetUp_CN_Linear_Constant_Matrix;
398:       ts->step   = TSStep_CN_Linear_Constant_Matrix;
399:     } else {
400:       if (isshell) {
401:         ts->Ashell = ts->A;
402:       }
403:       ts->setup  = TSSetUp_CN_Linear_Variable_Matrix;
404:       ts->step   = TSStep_CN_Linear_Variable_Matrix;
405:     }
406:     ts->setfromoptions  = TSSetFromOptions_CN_Linear;
407:     SLESCreate(ts->comm,&ts->sles);
408:     SLESGetKSP(ts->sles,&ksp);
409:     KSPSetInitialGuessNonzero(ksp);
410:   } else if (ts->problem_type == TS_NONLINEAR) {
411:     PetscTypeCompare((PetscObject)ts->A,MATSHELL,&isshell);
412:     if (isshell) {
413:       ts->Ashell = ts->A;
414:     }
415:     ts->setup           = TSSetUp_CN_Nonlinear;
416:     ts->step            = TSStep_CN_Nonlinear;
417:     ts->setfromoptions  = TSSetFromOptions_CN_Nonlinear;
418:     SNESCreate(ts->comm,SNES_NONLINEAR_EQUATIONS,&ts->snes);
419:   } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");

421:   PetscNew(TS_CN,&cn);
422:   PetscLogObjectMemory(ts,sizeof(TS_CN));
423:   ierr     = PetscMemzero(cn,sizeof(TS_CN));
424:   ts->data = (void*)cn;

426:   return(0);
427: }
428: EXTERN_C_END