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: /*
 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,PetscReal t,Vec x,Vec y)
 22: {
 23:   int         ierr;
 24:   PetscScalar neg_two = -2.0,neg_mdt = -1.0/ts->time_step;


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

 37:   if (ts->ops->rhsmatrix) { /* assemble matrix for this timestep */
 38:     MatStructure flg;
 39:     PetscStackPush("TS user right-hand-side matrix function");
 40:     (*ts->ops->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,PetscReal *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:   PetscScalar 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,PetscReal *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:   PetscScalar  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->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
123:     MatScale(&neg_dt,ts->A);
124:     MatShift(&two,ts->A);
125:     if (ts->B != ts->A && str != SAME_PRECONDITIONER) {
126:       MatScale(&neg_dt,ts->B);
127:       MatShift(&two,ts->B);
128:     }

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

134:     /* phase 2 - implicit step */
135:     VecCopy(sol,rhs);

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

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

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

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

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

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

189:   if (cn->update) {VecDestroy(cn->update);}
190:   if (cn->func) {VecDestroy(cn->func);}
191:   if (cn->rhs) {VecDestroy(cn->rhs);}
192:   PetscFree(cn);
193:   return(0);
194: }

196: /* 
197:     This defines the nonlinear equation that is to be solved with SNES

199:               U^{n+1} - dt*F(U^{n+1}) - U^{n}
200: */
201: int TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
202: {
203:   TS          ts = (TS) ctx;
204:   PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
205:   int         ierr,i,n;

208:   /* apply user provided function */
209:   TSComputeRHSFunction(ts,ts->ptime,x,y);
210:   /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
211:   VecGetArray(ts->vec_sol,&un);
212:   VecGetArray(x,&unp1);
213:   VecGetArray(y,&Funp1);
214:   VecGetLocalSize(x,&n);

216:   for (i=0; i<n; i++) {
217:     Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
218:   }
219:   VecRestoreArray(ts->vec_sol,&un);
220:   VecRestoreArray(x,&unp1);
221:   VecRestoreArray(y,&Funp1);
222:   return(0);
223: }

225: /*
226:    This constructs the Jacobian needed for SNES 

228:              J = I/dt - J_{F}   where J_{F} is the given Jacobian of F.
229: */
230: int TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
231: {
232:   TS           ts = (TS) ctx;
233:   int          ierr;
234:   PetscScalar  mone = -1.0,mdt = 1.0/ts->time_step;

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

240:   /* shift and scale Jacobian */
241:   MatScale(&mone,*AA);
242:   MatShift(&mdt,*AA);
243:   if (*BB != *AA && *str != SAME_PRECONDITIONER) {
244:     MatScale(&mone,*BB);
245:     MatShift(&mdt,*BB);
246:   }

248:   return(0);
249: }

251: /* ------------------------------------------------------------*/
252: static int TSSetUp_CN_Linear_Constant_Matrix(TS ts)
253: {
254:   TS_CN        *cn = (TS_CN*)ts->data;
255:   int          ierr;
256:   PetscScalar  two = 2.0,neg_dt = -1.0*ts->time_step;

259:   VecDuplicate(ts->vec_sol,&cn->update);
260:   VecDuplicate(ts->vec_sol,&cn->rhs);
261: 
262:   /* build linear system to be solved */
263:   MatScale(&neg_dt,ts->A);
264:   MatShift(&two,ts->A);
265:   if (ts->A != ts->B) {
266:     MatScale(&neg_dt,ts->B);
267:     MatShift(&two,ts->B);
268:   }
269:   SLESSetOperators(ts->sles,ts->A,ts->B,SAME_NONZERO_PATTERN);
270:   return(0);
271: }

273: static int TSSetUp_CN_Linear_Variable_Matrix(TS ts)
274: {
275:   TS_CN *cn = (TS_CN*)ts->data;
276:   int   ierr;

279:   VecDuplicate(ts->vec_sol,&cn->update);
280:   VecDuplicate(ts->vec_sol,&cn->rhs);
281:   return(0);
282: }

284: static int TSSetUp_CN_Nonlinear(TS ts)
285: {
286:   TS_CN *cn = (TS_CN*)ts->data;
287:   int   ierr;

290:   VecDuplicate(ts->vec_sol,&cn->update);
291:   VecDuplicate(ts->vec_sol,&cn->func);
292:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
293:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
294:   return(0);
295: }
296: /*------------------------------------------------------------*/

298: static int TSSetFromOptions_CN_Linear(TS ts)
299: {

303:   SLESSetFromOptions(ts->sles);
304:   return(0);
305: }

307: static int TSSetFromOptions_CN_Nonlinear(TS ts)
308: {

312:   SNESSetFromOptions(ts->snes);
313:   return(0);
314: }

316: static int TSView_CN(TS ts,PetscViewer viewer)
317: {
319:   return(0);
320: }

322: /* ------------------------------------------------------------ */
323: EXTERN_C_BEGIN
324: int TSCreate_CN(TS ts)
325: {
326:   TS_CN      *cn;
327:   int        ierr;
328:   KSP        ksp;

331:   ts->ops->destroy         = TSDestroy_CN;
332:   ts->ops->view            = TSView_CN;

334:   if (ts->problem_type == TS_LINEAR) {
335:     if (!ts->A) {
336:       SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
337:     }
338:     if (!ts->ops->rhsmatrix) {
339:       ts->ops->setup  = TSSetUp_CN_Linear_Constant_Matrix;
340:       ts->ops->step   = TSStep_CN_Linear_Constant_Matrix;
341:     } else {
342:       ts->ops->setup  = TSSetUp_CN_Linear_Variable_Matrix;
343:       ts->ops->step   = TSStep_CN_Linear_Variable_Matrix;
344:     }
345:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Linear;
346:     SLESCreate(ts->comm,&ts->sles);
347:     SLESGetKSP(ts->sles,&ksp);
348:     KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
349:   } else if (ts->problem_type == TS_NONLINEAR) {
350:     ts->ops->setup           = TSSetUp_CN_Nonlinear;
351:     ts->ops->step            = TSStep_CN_Nonlinear;
352:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Nonlinear;
353:     SNESCreate(ts->comm,&ts->snes);
354:   } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");

356:   PetscNew(TS_CN,&cn);
357:   PetscLogObjectMemory(ts,sizeof(TS_CN));
358:   ierr     = PetscMemzero(cn,sizeof(TS_CN));
359:   ts->data = (void*)cn;

361:   return(0);
362: }
363: EXTERN_C_END