Actual source code: cn.c

  1: #define PETSCTS_DLL

  3: /*
  4:        Code for Timestepping with implicit Crank-Nicholson method.
  5:     THIS IS NOT YET COMPLETE -- DO NOT USE!!
  6: */
 7:  #include src/ts/tsimpl.h

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

 15: /*------------------------------------------------------------------------------*/
 18: /*
 19:    TSComputeRHSFunctionEuler - Evaluates the right-hand-side function. 

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


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

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

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

 56:   return(0);
 57: }

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

 76:   TSGetKSP(ts,&ksp);
 77:   *steps = -ts->steps;
 78:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

 83:   for (i=0; i<max_steps; i++) {
 84:     ts->ptime += ts->time_step;
 85:     if (ts->ptime > ts->max_time) break;

 87:     /* phase 1 - explicit step */
 88:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
 89:     VecAXPBY(sol,dt,two,update);

 91:     /* phase 2 - implicit step */
 92:     VecCopy(sol,rhs);
 93:     /* apply user-provided boundary conditions (only needed if they are time dependent) */
 94:     TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);

 96:     KSPSolve(ts->ksp,rhs,update);
 97:     KSPGetIterationNumber(ksp,&its);
 98:     ts->linear_its += PetscAbsInt(its);
 99:     VecCopy(update,sol);
100:     ts->steps++;
101:     TSMonitor(ts,ts->steps,ts->ptime,sol);
102:   }

104:   *steps += ts->steps;
105:   *ptime  = ts->ptime;
106:   return(0);
107: }
108: /*
109:       Version where matrix depends on time 
110: */
113: static PetscErrorCode TSStep_CN_Linear_Variable_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
114: {
115:   TS_CN          *cn = (TS_CN*)ts->data;
116:   Vec            sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
118:   PetscInt       i,max_steps = ts->max_steps,its;
119:   PetscScalar    dt = ts->time_step,two = 2.0;
120:   MatStructure   str;
121:   KSP            ksp;

124:   TSGetKSP(ts,&ksp);
125:   *steps = -ts->steps;
126:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

131:   for (i=0; i<max_steps; i++) {
132:     ts->ptime += ts->time_step;
133:     if (ts->ptime > ts->max_time) break;
134:     /*
135:         evaluate matrix function 
136:     */
137:     (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
138:     TSScaleShiftMatrices(ts,ts->A,ts->B,str);

140:     /* phase 1 - explicit step */
141:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
142:     VecAXPBY(sol,dt,two,update);

144:     /* phase 2 - implicit step */
145:     VecCopy(sol,rhs);

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

150:     KSPSetOperators(ts->ksp,ts->A,ts->B,str);
151:     KSPSolve(ts->ksp,rhs,update);
152:     KSPGetIterationNumber(ksp,&its);
153:     ts->linear_its += PetscAbsInt(its);
154:     VecCopy(update,sol);
155:     ts->steps++;
156:     TSMonitor(ts,ts->steps,ts->ptime,sol);
157:   }

159:   *steps += ts->steps;
160:   *ptime  = ts->ptime;
161:   return(0);
162: }
163: /*
164:     Version for nonlinear PDE.
165: */
168: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
169: {
170:   Vec            sol = ts->vec_sol;
172:   PetscInt       i,max_steps = ts->max_steps,its,lits;
173:   TS_CN          *cn = (TS_CN*)ts->data;
174: 
176:   *steps = -ts->steps;
177:   TSMonitor(ts,ts->steps,ts->ptime,sol);

179:   for (i=0; i<max_steps; i++) {
180:     ts->ptime += ts->time_step;
181:     if (ts->ptime > ts->max_time) break;
182:     VecCopy(sol,cn->update);
183:     SNESSolve(ts->snes,PETSC_NULL,cn->update);
184:     SNESGetIterationNumber(ts->snes,&its);
185:     SNESGetNumberLinearIterations(ts->snes,&lits);
186:     ts->nonlinear_its += its; ts->linear_its += lits;
187:     VecCopy(cn->update,sol);
188:     ts->steps++;
189:     TSMonitor(ts,ts->steps,ts->ptime,sol);
190:   }

192:   *steps += ts->steps;
193:   *ptime  = ts->ptime;
194:   return(0);
195: }

197: /*------------------------------------------------------------*/
200: static PetscErrorCode TSDestroy_CN(TS ts)
201: {
202:   TS_CN          *cn = (TS_CN*)ts->data;

206:   if (cn->update) {VecDestroy(cn->update);}
207:   if (cn->func) {VecDestroy(cn->func);}
208:   if (cn->rhs) {VecDestroy(cn->rhs);}
209:   PetscFree(cn);
210:   return(0);
211: }

213: /* 
214:     This defines the nonlinear equation that is to be solved with SNES

216:               U^{n+1} - dt*F(U^{n+1}) - U^{n}
217: */
220: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
221: {
222:   TS             ts = (TS) ctx;
223:   PetscScalar    mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
225:   PetscInt       i,n;

228:   /* apply user provided function */
229:   TSComputeRHSFunction(ts,ts->ptime,x,y);
230:   /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
231:   VecGetArray(ts->vec_sol,&un);
232:   VecGetArray(x,&unp1);
233:   VecGetArray(y,&Funp1);
234:   VecGetLocalSize(x,&n);

236:   for (i=0; i<n; i++) {
237:     Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
238:   }
239:   VecRestoreArray(ts->vec_sol,&un);
240:   VecRestoreArray(x,&unp1);
241:   VecRestoreArray(y,&Funp1);
242:   return(0);
243: }

245: /*
246:    This constructs the Jacobian needed for SNES 

248:              J = I/dt - J_{F}   where J_{F} is the given Jacobian of F.
249: */
252: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
253: {
254:   TS             ts = (TS) ctx;

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

261:   /* shift and scale Jacobian */
262:   TSScaleShiftMatrices(ts,*AA,*BB,*str);
263:   return(0);
264: }

266: /* ------------------------------------------------------------*/
269: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
270: {
271:   TS_CN          *cn = (TS_CN*)ts->data;

275:   VecDuplicate(ts->vec_sol,&cn->update);
276:   VecDuplicate(ts->vec_sol,&cn->rhs);
277: 
278:   /* build linear system to be solved */
279:   TSScaleShiftMatrices(ts,ts->A,ts->B,SAME_NONZERO_PATTERN);
280:   KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
281:   return(0);
282: }

286: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
287: {
288:   TS_CN          *cn = (TS_CN*)ts->data;

292:   VecDuplicate(ts->vec_sol,&cn->update);
293:   VecDuplicate(ts->vec_sol,&cn->rhs);
294:   return(0);
295: }

299: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
300: {
301:   TS_CN          *cn = (TS_CN*)ts->data;

305:   VecDuplicate(ts->vec_sol,&cn->update);
306:   VecDuplicate(ts->vec_sol,&cn->func);
307:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
308:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
309:   return(0);
310: }
311: /*------------------------------------------------------------*/

315: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
316: {

320:   KSPSetFromOptions(ts->ksp);
321:   return(0);
322: }

326: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
327: {

331:   SNESSetFromOptions(ts->snes);
332:   return(0);
333: }

337: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
338: {
340:   return(0);
341: }

343: /* ------------------------------------------------------------ */
344: /*MC
345:       TS_CN - ODE solver using the implicit Crank-Nicholson method

347:   Level: beginner

349: .seealso:  TSCreate(), TS, TSSetType()

351: M*/
355: PetscErrorCode PETSCTS_DLLEXPORT TSCreate_CN(TS ts)
356: {
357:   TS_CN          *cn;
359:   KSP            ksp;

362:   ts->ops->destroy         = TSDestroy_CN;
363:   ts->ops->view            = TSView_CN;

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

387:   PetscNew(TS_CN,&cn);
388:   PetscLogObjectMemory(ts,sizeof(TS_CN));
389:   ts->data = (void*)cn;

391:   return(0);
392: }