Actual source code: cn.c

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

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

 13: /*------------------------------------------------------------------------------*/
 16: /*
 17:    TSComputeRHSFunctionEuler - Evaluates the right-hand-side function. 

 19:    Note: If the user did not provide a function but merely a matrix,
 20:    this routine applies the matrix.
 21: */
 22: PetscErrorCode TSComputeRHSFunctionEuler(TS ts,PetscReal t,Vec x,Vec y)
 23: {
 25:   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: */
 63: static PetscErrorCode TSStep_CN_Linear_Constant_Matrix(TS ts,PetscInt *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;
 69:   PetscInt       i,max_steps = ts->max_steps,its;
 70:   PetscScalar    dt = ts->time_step,two = 2.0;
 71:   KSP            ksp;

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

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

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

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

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

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

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

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

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

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

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

142:     /* phase 2 - implicit step */
143:     VecCopy(sol,rhs);

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

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

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

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

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

195: /*------------------------------------------------------------*/
198: static PetscErrorCode TSDestroy_CN(TS ts)
199: {
200:   TS_CN          *cn = (TS_CN*)ts->data;

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

211: /* 
212:     This defines the nonlinear equation that is to be solved with SNES

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

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

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

243: /*
244:    This constructs the Jacobian needed for SNES 

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

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

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

264: /* ------------------------------------------------------------*/
267: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
268: {
269:   TS_CN          *cn = (TS_CN*)ts->data;

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

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

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

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

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

313: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
314: {

318:   KSPSetFromOptions(ts->ksp);
319:   return(0);
320: }

324: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
325: {

329:   SNESSetFromOptions(ts->snes);
330:   return(0);
331: }

335: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
336: {
338:   return(0);
339: }

341: /* ------------------------------------------------------------ */
342: /*MC
343:       TS_CN - ODE solver using the implicit Crank-Nicholson method

345:   Level: beginner

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

349: M*/
353: PetscErrorCode TSCreate_CN(TS ts)
354: {
355:   TS_CN          *cn;
357:   KSP            ksp;

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

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

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

390:   return(0);
391: }