Actual source code: cn.c

  1: #define PETSCTS_DLL

  3: /*
  4:        Code for Timestepping with implicit Crank-Nicholson method.
  5: */
 6:  #include private/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  rhsfunc, rhsfunc_old; /* work vectors to hold rhs function provided by user */
 12:   Vec  rhs;            /* work vector for RHS; vec_sol/dt */
 13:   TS   ts;             /* used by ShellMult_private() */
 14:   PetscScalar mdt;     /* 1/dt, used by ShellMult_private() */
 15:   PetscReal rhsfunc_time,rhsfunc_old_time; /* time at which rhsfunc holds the value */
 16: } TS_CN;

 18: /*------------------------------------------------------------------------------*/
 19: /* 
 20:    Scale ts->Alhs = 1/dt*Alhs, ts->Arhs = 0.5*Arhs
 21:    Set   ts->A    = Alhs - Arhs, used in KSPSolve()
 22: */
 25: PetscErrorCode TSSetKSPOperators_CN_Matrix(TS ts)
 26: {
 28:   PetscScalar    mdt = 1.0/ts->time_step;

 31:   /* scale Arhs = 0.5*Arhs, Alhs = 1/dt*Alhs - assume dt is constant! */
 32:   MatScale(ts->Arhs,0.5);
 33:   if (ts->Alhs){
 34:     MatScale(ts->Alhs,mdt);
 35:   }
 36:   if (ts->A){
 37:     MatDestroy(ts->A);
 38:   }
 39:   MatDuplicate(ts->Arhs,MAT_COPY_VALUES,&ts->A);
 40: 
 41:   if (ts->Alhs){
 42:     /* ts->A = - Arhs + Alhs */
 43:     MatAYPX(ts->A,-1.0,ts->Alhs,ts->matflg);
 44:   } else {
 45:     /* ts->A = 1/dt - Arhs */
 46:     MatScale(ts->A,-1.0);
 47:     MatShift(ts->A,mdt);
 48:   }
 49:   return(0);
 50: }

 52: /* 
 53:    Scale ts->Alhs = 1/dt*Alhs, ts->Arhs = 0.5*Arhs
 54:    Set   ts->A    = Alhs - Arhs, used in KSPSolve()
 55: */
 58: PetscErrorCode ShellMult_private(Mat mat,Vec x,Vec y)
 59: {
 60:   PetscErrorCode  ierr;
 61:   void            *ctx;
 62:   TS_CN           *cn;

 65:   MatShellGetContext(mat,(void **)&ctx);
 66:   cn   = (TS_CN*)ctx;
 67:   MatMult(cn->ts->Arhs,x,y); /* y = 0.5*Arhs*x */
 68:   VecScale(y,-1.0);          /* y = -0.5*Arhs*x */
 69:   if (cn->ts->Alhs){
 70:     MatMultAdd(cn->ts->Alhs,x,y,y); /* y = 1/dt*Alhs*x + y */
 71:   } else {
 72:     VecAXPY(y,cn->mdt,x); /* y = 1/dt*x + y */
 73:   }
 74:   return(0);
 75: }
 78: PetscErrorCode TSSetKSPOperators_CN_No_Matrix(TS ts)
 79: {
 81:   PetscScalar    mdt = 1.0/ts->time_step;
 82:   Mat            Arhs = ts->Arhs;
 83:   MPI_Comm       comm;
 84:   PetscInt       m,n,M,N;
 85:   TS_CN          *cn = (TS_CN*)ts->data;

 88:   /* scale Arhs = 0.5*Arhs, Alhs = 1/dt*Alhs - assume dt is constant! */
 89:   MatScale(ts->Arhs,0.5);
 90:   if (ts->Alhs){
 91:     MatScale(ts->Alhs,mdt);
 92:   }
 93: 
 94:   cn->ts  = ts;
 95:   cn->mdt = mdt;
 96:   if (ts->A) {
 97:     MatDestroy(ts->A);
 98:   }
 99:   MatGetSize(Arhs,&M,&N);
100:   MatGetLocalSize(Arhs,&m,&n);
101:   PetscObjectGetComm((PetscObject)Arhs,&comm);
102:   MatCreateShell(comm,m,n,M,N,cn,&ts->A);
103:   MatShellSetOperation(ts->A,MATOP_MULT,(void(*)(void))ShellMult_private);
104:   return(0);
105: }

107: /*
108:     Version for linear PDE where RHS does not depend on time. Has built a
109:   single matrix that is to be used for all timesteps.
110: */
113: static PetscErrorCode TSStep_CN_Linear_Constant_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    mdt = 1.0/ts->time_step;

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

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

128:   for (i=0; i<max_steps; i++) {
129:     if (ts->ptime + ts->time_step > ts->max_time) break;
130:     /* set rhs = (1/dt*Alhs + 0.5*Arhs)*sol */
131:     MatMult(ts->Arhs,sol,rhs); /* rhs = 0.5*Arhs*sol */
132:     if (ts->Alhs){
133:       MatMultAdd(ts->Alhs,sol,rhs,rhs);      /* rhs = rhs + 1/dt*Alhs*sol */
134:     } else {
135:       VecAXPY(rhs,mdt,sol);    /* rhs = rhs + 1/dt*sol */
136:     }

138:     ts->ptime += ts->time_step;

140:     /* solve (1/dt*Alhs - 0.5*Arhs)*update = rhs */
141:     KSPSolve(ts->ksp,rhs,update);
142:     KSPGetIterationNumber(ts->ksp,&its);
143:     ts->linear_its += PetscAbsInt(its);
144:     VecCopy(update,sol);
145:     ts->steps++;
146:     TSMonitor(ts,ts->steps,ts->ptime,sol);
147:   }  *steps += ts->steps;
148:   *ptime  = ts->ptime;
149:   return(0);
150: }
151: /*
152:       Version where matrix depends on time 
153: */
156: static PetscErrorCode TSStep_CN_Linear_Variable_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
157: {
158:   TS_CN          *cn = (TS_CN*)ts->data;
159:   Vec            sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
161:   PetscInt       i,max_steps = ts->max_steps,its;
162:   PetscScalar    mdt = 1.0/ts->time_step;
163:   PetscReal      t_mid;
164:   MatStructure   str;

167:   *steps = -ts->steps;
168:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

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

176:     /* set rhs = (1/dt*Alhs(t_mid) + 0.5*Arhs(t_n)) * sol */
177:     if (i==0){
178:       /* evaluate 0.5*Arhs(t_0) */
179:       (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->Arhs,PETSC_NULL,&str,ts->jacP);
180:       MatScale(ts->Arhs,0.5);
181:     }
182:     if (ts->Alhs){
183:       /* evaluate Alhs(t_mid) */
184:       t_mid = ts->ptime+ts->time_step/2.0;
185:       (*ts->ops->lhsmatrix)(ts,t_mid,&ts->Alhs,PETSC_NULL,&str,ts->jacPlhs);
186:       MatMult(ts->Alhs,sol,rhs); /* rhs = Alhs_mid*sol */
187:       VecScale(rhs,mdt);         /* rhs = 1/dt*Alhs_mid*sol */
188:       MatMultAdd(ts->Arhs,sol,rhs,rhs);        /* rhs = rhs + 0.5*Arhs_mid*sol */
189:     } else {
190:       MatMult(ts->Arhs,sol,rhs); /* rhs = 0.5*Arhs_n*sol */
191:       VecAXPY(rhs,mdt,sol);      /* rhs = rhs + 1/dt*sol */
192:     }

194:     ts->ptime += ts->time_step;

196:     /* evaluate Arhs at current ptime t_{n+1} */
197:     (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->Arhs,PETSC_NULL,&str,ts->jacP);
198:     TSSetKSPOperators_CN_Matrix(ts);

200:     KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
201:     KSPSolve(ts->ksp,rhs,update);
202:     KSPGetIterationNumber(ts->ksp,&its);
203:     ts->linear_its += PetscAbsInt(its);
204:     VecCopy(update,sol);
205:     ts->steps++;
206:     TSMonitor(ts,ts->steps,ts->ptime,sol);
207:   }

209:   *steps += ts->steps;
210:   *ptime  = ts->ptime;
211:   return(0);
212: }
213: /*
214:     Version for nonlinear PDE.
215: */
218: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
219: {
220:   Vec            sol = ts->vec_sol;
222:   PetscInt       i,max_steps = ts->max_steps,its,lits;
223:   TS_CN          *cn = (TS_CN*)ts->data;
224: 
226:   *steps = -ts->steps;
227:   TSMonitor(ts,ts->steps,ts->ptime,sol);

229:   for (i=0; i<max_steps; i++) {
230:     if (ts->ptime + ts->time_step > ts->max_time) break;
231:     ts->ptime += ts->time_step;
232: 
233:     VecCopy(sol,cn->update);
234:     SNESSolve(ts->snes,PETSC_NULL,cn->update);
235:     SNESGetIterationNumber(ts->snes,&its);
236:     SNESGetLinearSolveIterations(ts->snes,&lits);
237:     ts->nonlinear_its += its; ts->linear_its += lits;
238:     VecCopy(cn->update,sol);
239:     ts->steps++;
240:     TSMonitor(ts,ts->steps,ts->ptime,sol);
241:   }

243:   *steps += ts->steps;
244:   *ptime  = ts->ptime;
245:   return(0);
246: }

248: /*------------------------------------------------------------*/
251: static PetscErrorCode TSDestroy_CN(TS ts)
252: {
253:   TS_CN          *cn = (TS_CN*)ts->data;

257:   if (cn->update) {VecDestroy(cn->update);}
258:   if (cn->func) {VecDestroy(cn->func);}
259:   if (cn->rhsfunc) {VecDestroy(cn->rhsfunc);}
260:   if (cn->rhsfunc_old) {VecDestroy(cn->rhsfunc_old);}
261:   if (cn->rhs) {VecDestroy(cn->rhs);}
262:   PetscFree(cn);
263:   return(0);
264: }

266: /* 
267:     This defines the nonlinear equation that is to be solved with SNES
268:        1/dt*Alhs*(U^{n+1} - U^{n}) - 0.5*(F(U^{n+1}) + F(U^{n}))
269: */
272: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
273: {
274:   TS             ts = (TS) ctx;
275:   PetscScalar    mdt = 1.0/ts->time_step,*unp1,*un,*Funp1,*Fun,*yarray;
277:   PetscInt       i,n;
278:   TS_CN          *cn = (TS_CN*)ts->data;

281:   /* apply user provided function */
282:   if (cn->rhsfunc_time == (ts->ptime - ts->time_step)){
283:     /* printf("   copy rhsfunc to rhsfunc_old, then eval rhsfunc\n"); */
284:     VecCopy(cn->rhsfunc,cn->rhsfunc_old);
285:     cn->rhsfunc_old_time = cn->rhsfunc_time;
286:   } else if (cn->rhsfunc_time != ts->ptime && cn->rhsfunc_old_time != ts->ptime-ts->time_step){
287:     /* printf("   eval both rhsfunc_old and rhsfunc\n"); */
288:     TSComputeRHSFunction(ts,ts->ptime-ts->time_step,ts->vec_sol,cn->rhsfunc_old); /* rhsfunc_old=F(U^{n}) */
289:     cn->rhsfunc_old_time = ts->ptime - ts->time_step;
290:   }
291: 
292:   if (ts->Alhs){
293:     /* compute y=Alhs*(U^{n+1} - U^{n}) with cn->rhsfunc as workspce */
294:     VecWAXPY(cn->rhsfunc,-1.0,ts->vec_sol,x);
295:     MatMult(ts->Alhs,cn->rhsfunc,y);
296:   }

298:   TSComputeRHSFunction(ts,ts->ptime,x,cn->rhsfunc); /* rhsfunc = F(U^{n+1}) */
299:   cn->rhsfunc_time = ts->ptime;
300: 
301:   VecGetArray(ts->vec_sol,&un); /* U^{n} */
302:   VecGetArray(x,&unp1);         /* U^{n+1} */
303:   VecGetArray(cn->rhsfunc,&Funp1);
304:   VecGetArray(cn->rhsfunc_old,&Fun);
305:   VecGetArray(y,&yarray);
306:   VecGetLocalSize(x,&n);
307:   if (ts->Alhs){
308:     for (i=0; i<n; i++) {
309:       yarray[i] = mdt*yarray[i] - 0.5*(Funp1[i] + Fun[i]);
310:     }
311:   } else {
312:     for (i=0; i<n; i++) {
313:       yarray[i] = mdt*(unp1[i] - un[i]) - 0.5*(Funp1[i] + Fun[i]);
314:     }
315:   }
316:   VecRestoreArray(ts->vec_sol,&un);
317:   VecRestoreArray(x,&unp1);
318:   VecRestoreArray(cn->rhsfunc,&Funp1);
319:   VecRestoreArray(cn->rhsfunc_old,&Fun);
320:   VecRestoreArray(y,&yarray);
321:   return(0);
322: }

324: /* Set B = 1/dt*Alh - 0.5*B */
327: PetscErrorCode TSScaleShiftMatrices_CN(TS ts,Mat A,Mat B,MatStructure str)
328: {
329:   PetscTruth     flg;
331:   PetscScalar    mdt = 1.0/ts->time_step;

334:   PetscTypeCompare((PetscObject)B,MATMFFD,&flg);
335:   if (!flg) {
336:     MatScale(B,-0.5);
337:     if (ts->Alhs){
338:       MatAXPY(B,mdt,ts->Alhs,DIFFERENT_NONZERO_PATTERN); /* DIFFERENT_NONZERO_PATTERN? */
339:     } else {
340:       MatShift(B,mdt);
341:     }
342:   } else {
343:     SETERRQ(PETSC_ERR_SUP,"Matrix type MATMFFD is not supported yet"); /* ref TSScaleShiftMatrices() */
344:   }
345:   return(0);
346: }

348: /*
349:    This constructs the Jacobian needed for SNES 
350:      J = I/dt - 0.5*J_{F}   where J_{F} is the given Jacobian of F.
351:      x  - input vector
352:      AA - Jacobian matrix 
353:      BB - preconditioner matrix, usually the same as AA
354: */
357: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
358: {
359:   TS             ts = (TS) ctx;

363:   /* construct user's Jacobian */
364:   TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str); /* BB = J_{F} */

366:   /* shift and scale Jacobian */
367:   TSScaleShiftMatrices_CN(ts,*AA,*BB,*str); /* Set BB = 1/dt*Alhs - 0.5*BB */
368:   return(0);
369: }

371: /* ------------------------------------------------------------*/
374: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
375: {
376:   TS_CN          *cn = (TS_CN*)ts->data;
378:   PetscTruth shelltype;

381:   VecDuplicate(ts->vec_sol,&cn->update);
382:   VecDuplicate(ts->vec_sol,&cn->rhs);
383: 
384:   /* build linear system to be solved */
385:   /* scale ts->Alhs = 1/dt*Alhs, ts->Arhs = 0.5*Arhs; set ts->A = Alhs - Arhs */
386:   PetscTypeCompare((PetscObject)ts->Arhs,MATSHELL,&shelltype);
387:   if (shelltype){
388:     TSSetKSPOperators_CN_No_Matrix(ts);
389:   } else {
390:     TSSetKSPOperators_CN_Matrix(ts);
391:   }
392:   KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
393:   return(0);
394: }

398: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
399: {
400:   TS_CN          *cn = (TS_CN*)ts->data;

404:   VecDuplicate(ts->vec_sol,&cn->update);
405:   VecDuplicate(ts->vec_sol,&cn->rhs);
406:   return(0);
407: }

411: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
412: {
413:   TS_CN          *cn = (TS_CN*)ts->data;

417:   VecDuplicate(ts->vec_sol,&cn->update);
418:   VecDuplicate(ts->vec_sol,&cn->func);
419:   VecDuplicate(ts->vec_sol,&cn->rhsfunc);
420:   VecDuplicate(ts->vec_sol,&cn->rhsfunc_old);
421:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
422:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
423:   cn->rhsfunc_time     = -100.0; /* cn->rhsfunc is not evaluated yet */
424:   cn->rhsfunc_old_time = -100.0;
425:   return(0);
426: }
427: /*------------------------------------------------------------*/

431: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
432: {
434:   return(0);
435: }

439: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
440: {
442:   return(0);
443: }

447: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
448: {
450:   return(0);
451: }

453: /* ------------------------------------------------------------ */
454: /*MC
455:       TS_CN - ODE solver using the implicit Crank-Nicholson method

457:   Level: beginner

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

461: M*/
465: PetscErrorCode  TSCreate_CN(TS ts)
466: {
467:   TS_CN          *cn;

471:   ts->ops->destroy = TSDestroy_CN;
472:   ts->ops->view    = TSView_CN;

474:   if (ts->problem_type == TS_LINEAR) {
475:     if (!ts->Arhs) {
476:       SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
477:     }
478:     if (!ts->ops->rhsmatrix) {
479:       ts->ops->setup = TSSetUp_CN_Linear_Constant_Matrix;
480:       ts->ops->step  = TSStep_CN_Linear_Constant_Matrix;
481:     } else {
482:       ts->ops->setup = TSSetUp_CN_Linear_Variable_Matrix;
483:       ts->ops->step  = TSStep_CN_Linear_Variable_Matrix;
484:     }
485:     ts->ops->setfromoptions = TSSetFromOptions_CN_Linear;
486:     KSPCreate(((PetscObject)ts)->comm,&ts->ksp);
487:     PetscObjectIncrementTabLevel((PetscObject)ts->ksp,(PetscObject)ts,1);
488:     KSPSetInitialGuessNonzero(ts->ksp,PETSC_TRUE);
489:   } else if (ts->problem_type == TS_NONLINEAR) {
490:     ts->ops->setup          = TSSetUp_CN_Nonlinear;
491:     ts->ops->step           = TSStep_CN_Nonlinear;
492:     ts->ops->setfromoptions = TSSetFromOptions_CN_Nonlinear;
493:     SNESCreate(((PetscObject)ts)->comm,&ts->snes);
494:     PetscObjectIncrementTabLevel((PetscObject)ts->snes,(PetscObject)ts,1);
495:   } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");

497:   PetscNewLog(ts,TS_CN,&cn);
498:   ts->data = (void*)cn;
499:   return(0);
500: }