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:   return(0);
 54: }

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

 73:   TSGetKSP(ts,&ksp);
 74:   *steps = -ts->steps;
 75:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

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

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

 88:     /* phase 2 - implicit step */
 89:     VecCopy(sol,rhs);

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

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

119:   TSGetKSP(ts,&ksp);
120:   *steps = -ts->steps;
121:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

126:   for (i=0; i<max_steps; i++) {
127:     ts->ptime += ts->time_step;
128:     if (ts->ptime > ts->max_time) break;
129:     /*
130:         evaluate matrix function 
131:     */
132:     (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
133:     TSScaleShiftMatrices(ts,ts->A,ts->B,str);

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

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

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

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

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

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

189: /*------------------------------------------------------------*/
192: static PetscErrorCode TSDestroy_CN(TS ts)
193: {
194:   TS_CN          *cn = (TS_CN*)ts->data;

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

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

208:               U^{n+1} - dt*F(U^{n+1}) - U^{n}
209: */
212: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
213: {
214:   TS             ts = (TS) ctx;
215:   PetscScalar    mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
217:   PetscInt       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: */
244: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
245: {
246:   TS             ts = (TS) ctx;

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

253:   /* shift and scale Jacobian */
254:   TSScaleShiftMatrices(ts,*AA,*BB,*str);
255:   return(0);
256: }

258: /* ------------------------------------------------------------*/
261: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
262: {
263:   TS_CN          *cn = (TS_CN*)ts->data;

267:   VecDuplicate(ts->vec_sol,&cn->update);
268:   VecDuplicate(ts->vec_sol,&cn->rhs);
269: 
270:   /* build linear system to be solved */
271:   TSScaleShiftMatrices(ts,ts->A,ts->B,SAME_NONZERO_PATTERN);
272:   KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
273:   return(0);
274: }

278: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
279: {
280:   TS_CN          *cn = (TS_CN*)ts->data;

284:   VecDuplicate(ts->vec_sol,&cn->update);
285:   VecDuplicate(ts->vec_sol,&cn->rhs);
286:   return(0);
287: }

291: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
292: {
293:   TS_CN          *cn = (TS_CN*)ts->data;

297:   VecDuplicate(ts->vec_sol,&cn->update);
298:   VecDuplicate(ts->vec_sol,&cn->func);
299:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
300:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
301:   return(0);
302: }
303: /*------------------------------------------------------------*/

307: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
308: {

312:   KSPSetFromOptions(ts->ksp);
313:   return(0);
314: }

318: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
319: {

323:   SNESSetFromOptions(ts->snes);
324:   return(0);
325: }

329: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
330: {
332:   return(0);
333: }

335: /* ------------------------------------------------------------ */
336: /*MC
337:       TS_CN - ODE solver using the implicit Crank-Nicholson method

339:   Level: beginner

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

343: M*/
347: PetscErrorCode  TSCreate_CN(TS ts)
348: {
349:   TS_CN          *cn;
351:   KSP            ksp;

354:   ts->ops->destroy         = TSDestroy_CN;
355:   ts->ops->view            = TSView_CN;

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

379:   PetscNew(TS_CN,&cn);
380:   PetscLogObjectMemory(ts,sizeof(TS_CN));
381:   ts->data = (void*)cn;

383:   SETERRQ(PETSC_ERR_SUP,"The code for Crank-Nicholson is not complete\n         emai petsc-maint@mcs.anl.gov for more info");
384:   return(0);
385: }