Actual source code: cn.c

  1: /*$Id: cn.c,v 1.35 2001/09/11 16:34:19 bsmith 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: #undef __FUNCT__  
 17: /*
 18:    TSComputeRHSFunctionEuler - Evaluates the right-hand-side function. 

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


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

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

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

 54:   return(0);
 55: }

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

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

 78:   for (i=0; i<max_steps; i++) {
 79:     ts->ptime += ts->time_step;
 80:     if (ts->ptime > ts->max_time) break;

 82:     /* phase 1 - explicit step */
 83:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
 84:     VecAXPBY(&dt,&two,update,sol);

 86:     /* phase 2 - implicit step */
 87:     VecCopy(sol,rhs);
 88:     /* apply user-provided boundary conditions (only needed if they are time dependent) */
 89:     TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);

 91:     SLESSolve(ts->sles,rhs,update,&its);
 92:     ts->linear_its += PetscAbsInt(its);
 93:     VecCopy(update,sol);
 94:     ts->steps++;
 95:     TSMonitor(ts,ts->steps,ts->ptime,sol);
 96:   }

 98:   *steps += ts->steps;
 99:   *ptime  = ts->ptime;
100:   return(0);
101: }
102: /*
103:       Version where matrix depends on time 
104: */
105: #undef __FUNCT__  
107: static int TSStep_CN_Linear_Variable_Matrix(TS ts,int *steps,PetscReal *ptime)
108: {
109:   TS_CN        *cn = (TS_CN*)ts->data;
110:   Vec          sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
111:   int          ierr,i,max_steps = ts->max_steps,its;
112:   PetscScalar  dt = ts->time_step,two = 2.0,neg_dt = -1.0*ts->time_step;
113:   MatStructure str;

116:   *steps = -ts->steps;
117:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

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

136:     /* phase 1 - explicit step */
137:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
138:     VecAXPBY(&dt,&two,update,sol);

140:     /* phase 2 - implicit step */
141:     VecCopy(sol,rhs);

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

146:     SLESSetOperators(ts->sles,ts->A,ts->B,str);
147:     SLESSolve(ts->sles,rhs,update,&its);
148:     ts->linear_its += PetscAbsInt(its);
149:     VecCopy(update,sol);
150:     ts->steps++;
151:     TSMonitor(ts,ts->steps,ts->ptime,sol);
152:   }

154:   *steps += ts->steps;
155:   *ptime  = ts->ptime;
156:   return(0);
157: }
158: /*
159:     Version for nonlinear PDE.
160: */
161: #undef __FUNCT__  
163: static int TSStep_CN_Nonlinear(TS ts,int *steps,PetscReal *ptime)
164: {
165:   Vec   sol = ts->vec_sol;
166:   int   ierr,i,max_steps = ts->max_steps,its,lits;
167:   TS_CN *cn = (TS_CN*)ts->data;
168: 
170:   *steps = -ts->steps;
171:   TSMonitor(ts,ts->steps,ts->ptime,sol);

173:   for (i=0; i<max_steps; i++) {
174:     ts->ptime += ts->time_step;
175:     if (ts->ptime > ts->max_time) break;
176:     VecCopy(sol,cn->update);
177:     SNESSolve(ts->snes,cn->update,&its);
178:     SNESGetNumberLinearIterations(ts->snes,&lits);
179:     ts->nonlinear_its += PetscAbsInt(its); ts->linear_its += lits;
180:     VecCopy(cn->update,sol);
181:     ts->steps++;
182:     TSMonitor(ts,ts->steps,ts->ptime,sol);
183:   }

185:   *steps += ts->steps;
186:   *ptime  = ts->ptime;
187:   return(0);
188: }

190: /*------------------------------------------------------------*/
191: #undef __FUNCT__  
193: static int TSDestroy_CN(TS ts)
194: {
195:   TS_CN *cn = (TS_CN*)ts->data;
196:   int   ierr;

199:   if (cn->update) {VecDestroy(cn->update);}
200:   if (cn->func) {VecDestroy(cn->func);}
201:   if (cn->rhs) {VecDestroy(cn->rhs);}
202:   PetscFree(cn);
203:   return(0);
204: }

206: /* 
207:     This defines the nonlinear equation that is to be solved with SNES

209:               U^{n+1} - dt*F(U^{n+1}) - U^{n}
210: */
211: #undef __FUNCT__  
213: int TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
214: {
215:   TS          ts = (TS) ctx;
216:   PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
217:   int         ierr,i,n;

220:   /* apply user provided function */
221:   TSComputeRHSFunction(ts,ts->ptime,x,y);
222:   /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
223:   VecGetArray(ts->vec_sol,&un);
224:   VecGetArray(x,&unp1);
225:   VecGetArray(y,&Funp1);
226:   VecGetLocalSize(x,&n);

228:   for (i=0; i<n; i++) {
229:     Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
230:   }
231:   VecRestoreArray(ts->vec_sol,&un);
232:   VecRestoreArray(x,&unp1);
233:   VecRestoreArray(y,&Funp1);
234:   return(0);
235: }

237: /*
238:    This constructs the Jacobian needed for SNES 

240:              J = I/dt - J_{F}   where J_{F} is the given Jacobian of F.
241: */
242: #undef __FUNCT__  
244: int TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
245: {
246:   TS           ts = (TS) ctx;
247:   int          ierr;
248:   PetscScalar  mone = -1.0,mdt = 1.0/ts->time_step;

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

254:   /* shift and scale Jacobian */
255:   MatScale(&mone,*AA);
256:   MatShift(&mdt,*AA);
257:   if (*BB != *AA && *str != SAME_PRECONDITIONER) {
258:     MatScale(&mone,*BB);
259:     MatShift(&mdt,*BB);
260:   }

262:   return(0);
263: }

265: /* ------------------------------------------------------------*/
266: #undef __FUNCT__  
268: static int TSSetUp_CN_Linear_Constant_Matrix(TS ts)
269: {
270:   TS_CN        *cn = (TS_CN*)ts->data;
271:   int          ierr;
272:   PetscScalar  two = 2.0,neg_dt = -1.0*ts->time_step;

275:   VecDuplicate(ts->vec_sol,&cn->update);
276:   VecDuplicate(ts->vec_sol,&cn->rhs);
277: 
278:   /* build linear system to be solved */
279:   MatScale(&neg_dt,ts->A);
280:   MatShift(&two,ts->A);
281:   if (ts->A != ts->B) {
282:     MatScale(&neg_dt,ts->B);
283:     MatShift(&two,ts->B);
284:   }
285:   SLESSetOperators(ts->sles,ts->A,ts->B,SAME_NONZERO_PATTERN);
286:   return(0);
287: }

289: #undef __FUNCT__  
291: static int TSSetUp_CN_Linear_Variable_Matrix(TS ts)
292: {
293:   TS_CN *cn = (TS_CN*)ts->data;
294:   int   ierr;

297:   VecDuplicate(ts->vec_sol,&cn->update);
298:   VecDuplicate(ts->vec_sol,&cn->rhs);
299:   return(0);
300: }

302: #undef __FUNCT__  
304: static int TSSetUp_CN_Nonlinear(TS ts)
305: {
306:   TS_CN *cn = (TS_CN*)ts->data;
307:   int   ierr;

310:   VecDuplicate(ts->vec_sol,&cn->update);
311:   VecDuplicate(ts->vec_sol,&cn->func);
312:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
313:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
314:   return(0);
315: }
316: /*------------------------------------------------------------*/

318: #undef __FUNCT__  
320: static int TSSetFromOptions_CN_Linear(TS ts)
321: {

325:   SLESSetFromOptions(ts->sles);
326:   return(0);
327: }

329: #undef __FUNCT__  
331: static int TSSetFromOptions_CN_Nonlinear(TS ts)
332: {

336:   SNESSetFromOptions(ts->snes);
337:   return(0);
338: }

340: #undef __FUNCT__  
342: static int TSView_CN(TS ts,PetscViewer viewer)
343: {
345:   return(0);
346: }

348: /* ------------------------------------------------------------ */
349: EXTERN_C_BEGIN
350: #undef __FUNCT__  
352: int TSCreate_CN(TS ts)
353: {
354:   TS_CN      *cn;
355:   int        ierr;
356:   KSP        ksp;

359:   ts->ops->destroy         = TSDestroy_CN;
360:   ts->ops->view            = TSView_CN;

362:   if (ts->problem_type == TS_LINEAR) {
363:     if (!ts->A) {
364:       SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
365:     }
366:     if (!ts->ops->rhsmatrix) {
367:       ts->ops->setup  = TSSetUp_CN_Linear_Constant_Matrix;
368:       ts->ops->step   = TSStep_CN_Linear_Constant_Matrix;
369:     } else {
370:       ts->ops->setup  = TSSetUp_CN_Linear_Variable_Matrix;
371:       ts->ops->step   = TSStep_CN_Linear_Variable_Matrix;
372:     }
373:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Linear;
374:     SLESCreate(ts->comm,&ts->sles);
375:     SLESGetKSP(ts->sles,&ksp);
376:     KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
377:   } else if (ts->problem_type == TS_NONLINEAR) {
378:     ts->ops->setup           = TSSetUp_CN_Nonlinear;
379:     ts->ops->step            = TSStep_CN_Nonlinear;
380:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Nonlinear;
381:     SNESCreate(ts->comm,&ts->snes);
382:   } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");

384:   PetscNew(TS_CN,&cn);
385:   PetscLogObjectMemory(ts,sizeof(TS_CN));
386:   ierr     = PetscMemzero(cn,sizeof(TS_CN));
387:   ts->data = (void*)cn;

389:   return(0);
390: }
391: EXTERN_C_END