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: /*------------------------------------------------------------------------------*/
 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;


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

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

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

 55:   return(0);
 56: }

 58: /*
 59:     Version for linear PDE where RHS does not depend on time. Has built a
 60:   single matrix that is to be used for all timesteps.
 61: */
 64: static int TSStep_CN_Linear_Constant_Matrix(TS ts,int *steps,PetscReal *ptime)
 65: {
 66:   TS_CN       *cn = (TS_CN*)ts->data;
 67:   Vec         sol = ts->vec_sol,update = cn->update;
 68:   Vec         rhs = cn->rhs;
 69:   int         ierr,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:     KSPSetRhs(ts->ksp,rhs);
 95:     KSPSetSolution(ts->ksp,update);
 96:     KSPSolve(ts->ksp);
 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 int TSStep_CN_Linear_Variable_Matrix(TS ts,int *steps,PetscReal *ptime)
114: {
115:   TS_CN        *cn = (TS_CN*)ts->data;
116:   Vec          sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
117:   int          ierr,i,max_steps = ts->max_steps,its;
118:   PetscScalar  dt = ts->time_step,two = 2.0,neg_dt = -1.0*ts->time_step;
119:   MatStructure str;
120:   KSP          ksp;

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

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

130:   for (i=0; i<max_steps; i++) {
131:     ts->ptime += ts->time_step;
132:     if (ts->ptime > ts->max_time) break;
133:     /*
134:         evaluate matrix function 
135:     */
136:     (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
137:     MatScale(&neg_dt,ts->A);
138:     MatShift(&two,ts->A);
139:     if (ts->B != ts->A && str != SAME_PRECONDITIONER) {
140:       MatScale(&neg_dt,ts->B);
141:       MatShift(&two,ts->B);
142:     }

144:     /* phase 1 - explicit step */
145:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
146:     VecAXPBY(&dt,&two,update,sol);

148:     /* phase 2 - implicit step */
149:     VecCopy(sol,rhs);

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

154:     KSPSetOperators(ts->ksp,ts->A,ts->B,str);
155:     KSPSetRhs(ts->ksp,rhs);
156:     KSPSetSolution(ts->ksp,update);
157:     KSPSolve(ts->ksp);
158:     KSPGetIterationNumber(ksp,&its);
159:     ts->linear_its += PetscAbsInt(its);
160:     VecCopy(update,sol);
161:     ts->steps++;
162:     TSMonitor(ts,ts->steps,ts->ptime,sol);
163:   }

165:   *steps += ts->steps;
166:   *ptime  = ts->ptime;
167:   return(0);
168: }
169: /*
170:     Version for nonlinear PDE.
171: */
174: static int TSStep_CN_Nonlinear(TS ts,int *steps,PetscReal *ptime)
175: {
176:   Vec   sol = ts->vec_sol;
177:   int   ierr,i,max_steps = ts->max_steps,its,lits;
178:   TS_CN *cn = (TS_CN*)ts->data;
179: 
181:   *steps = -ts->steps;
182:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

197:   *steps += ts->steps;
198:   *ptime  = ts->ptime;
199:   return(0);
200: }

202: /*------------------------------------------------------------*/
205: static int TSDestroy_CN(TS ts)
206: {
207:   TS_CN *cn = (TS_CN*)ts->data;
208:   int   ierr;

211:   if (cn->update) {VecDestroy(cn->update);}
212:   if (cn->func) {VecDestroy(cn->func);}
213:   if (cn->rhs) {VecDestroy(cn->rhs);}
214:   PetscFree(cn);
215:   return(0);
216: }

218: /* 
219:     This defines the nonlinear equation that is to be solved with SNES

221:               U^{n+1} - dt*F(U^{n+1}) - U^{n}
222: */
225: int TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
226: {
227:   TS          ts = (TS) ctx;
228:   PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
229:   int         ierr,i,n;

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

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

249: /*
250:    This constructs the Jacobian needed for SNES 

252:              J = I/dt - J_{F}   where J_{F} is the given Jacobian of F.
253: */
256: int TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
257: {
258:   TS           ts = (TS) ctx;
259:   int          ierr;
260:   PetscScalar  mone = -1.0,mdt = 1.0/ts->time_step;

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

266:   /* shift and scale Jacobian */
267:   MatScale(&mone,*AA);
268:   MatShift(&mdt,*AA);
269:   if (*BB != *AA && *str != SAME_PRECONDITIONER) {
270:     MatScale(&mone,*BB);
271:     MatShift(&mdt,*BB);
272:   }

274:   return(0);
275: }

277: /* ------------------------------------------------------------*/
280: static int TSSetUp_CN_Linear_Constant_Matrix(TS ts)
281: {
282:   TS_CN        *cn = (TS_CN*)ts->data;
283:   int          ierr;
284:   PetscScalar  two = 2.0,neg_dt = -1.0*ts->time_step;

287:   VecDuplicate(ts->vec_sol,&cn->update);
288:   VecDuplicate(ts->vec_sol,&cn->rhs);
289: 
290:   /* build linear system to be solved */
291:   MatScale(&neg_dt,ts->A);
292:   MatShift(&two,ts->A);
293:   if (ts->A != ts->B) {
294:     MatScale(&neg_dt,ts->B);
295:     MatShift(&two,ts->B);
296:   }
297:   KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
298:   return(0);
299: }

303: static int TSSetUp_CN_Linear_Variable_Matrix(TS ts)
304: {
305:   TS_CN *cn = (TS_CN*)ts->data;
306:   int   ierr;

309:   VecDuplicate(ts->vec_sol,&cn->update);
310:   VecDuplicate(ts->vec_sol,&cn->rhs);
311:   return(0);
312: }

316: static int TSSetUp_CN_Nonlinear(TS ts)
317: {
318:   TS_CN *cn = (TS_CN*)ts->data;
319:   int   ierr;

322:   VecDuplicate(ts->vec_sol,&cn->update);
323:   VecDuplicate(ts->vec_sol,&cn->func);
324:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
325:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
326:   return(0);
327: }
328: /*------------------------------------------------------------*/

332: static int TSSetFromOptions_CN_Linear(TS ts)
333: {

337:   KSPSetFromOptions(ts->ksp);
338:   return(0);
339: }

343: static int TSSetFromOptions_CN_Nonlinear(TS ts)
344: {

348:   SNESSetFromOptions(ts->snes);
349:   return(0);
350: }

354: static int TSView_CN(TS ts,PetscViewer viewer)
355: {
357:   return(0);
358: }

360: /* ------------------------------------------------------------ */
361: /*MC
362:       TS_CN - ODE solver using the implicit Crank-Nicholson method

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

366: M*/
367: EXTERN_C_BEGIN
370: int TSCreate_CN(TS ts)
371: {
372:   TS_CN      *cn;
373:   int        ierr;
374:   KSP        ksp;

377:   ts->ops->destroy         = TSDestroy_CN;
378:   ts->ops->view            = TSView_CN;

380:   if (ts->problem_type == TS_LINEAR) {
381:     if (!ts->A) {
382:       SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
383:     }
384:     if (!ts->ops->rhsmatrix) {
385:       ts->ops->setup  = TSSetUp_CN_Linear_Constant_Matrix;
386:       ts->ops->step   = TSStep_CN_Linear_Constant_Matrix;
387:     } else {
388:       ts->ops->setup  = TSSetUp_CN_Linear_Variable_Matrix;
389:       ts->ops->step   = TSStep_CN_Linear_Variable_Matrix;
390:     }
391:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Linear;
392:     KSPCreate(ts->comm,&ts->ksp);
393:     TSGetKSP(ts,&ksp);
394:     KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
395:   } else if (ts->problem_type == TS_NONLINEAR) {
396:     ts->ops->setup           = TSSetUp_CN_Nonlinear;
397:     ts->ops->step            = TSStep_CN_Nonlinear;
398:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Nonlinear;
399:     SNESCreate(ts->comm,&ts->snes);
400:   } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");

402:   PetscNew(TS_CN,&cn);
403:   PetscLogObjectMemory(ts,sizeof(TS_CN));
404:   PetscMemzero(cn,sizeof(TS_CN));
405:   ts->data = (void*)cn;

407:   return(0);
408: }
409: EXTERN_C_END