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:     TSPreStep(ts);
131:     /* set rhs = (1/dt*Alhs + 0.5*Arhs)*sol */
132:     MatMult(ts->Arhs,sol,rhs); /* rhs = 0.5*Arhs*sol */
133:     if (ts->Alhs){
134:       MatMultAdd(ts->Alhs,sol,rhs,rhs);      /* rhs = rhs + 1/dt*Alhs*sol */
135:     } else {
136:       VecAXPY(rhs,mdt,sol);    /* rhs = rhs + 1/dt*sol */
137:     }

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

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

170:   *steps = -ts->steps;
171:   TSMonitor(ts,ts->steps,ts->ptime,sol);

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

176:   for (i=0; i<max_steps; i++) {
177:     if (ts->ptime + ts->time_step > ts->max_time) break;
178:     TSPreStep(ts);

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

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

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

204:     KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
205:     KSPSolve(ts->ksp,rhs,update);
206:     KSPGetIterationNumber(ts->ksp,&its);
207:     ts->linear_its += PetscAbsInt(its);
208:     VecCopy(update,sol);
209:     ts->steps++;
210:     TSPostStep(ts);
211:     TSMonitor(ts,ts->steps,ts->ptime,sol);
212:   }

214:   *steps += ts->steps;
215:   *ptime  = ts->ptime;
216:   return(0);
217: }
218: /*
219:     Version for nonlinear PDE.
220: */
223: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
224: {
225:   Vec            sol = ts->vec_sol;
227:   PetscInt       i,max_steps = ts->max_steps,its,lits;
228:   TS_CN          *cn = (TS_CN*)ts->data;
229: 
231:   *steps = -ts->steps;
232:   TSMonitor(ts,ts->steps,ts->ptime,sol);

234:   for (i=0; i<max_steps; i++) {
235:     if (ts->ptime + ts->time_step > ts->max_time) break;
236:     TSPreStep(ts);
237:     ts->ptime += ts->time_step;
238: 
239:     VecCopy(sol,cn->update);
240:     SNESSolve(ts->snes,PETSC_NULL,cn->update);
241:     SNESGetIterationNumber(ts->snes,&its);
242:     SNESGetLinearSolveIterations(ts->snes,&lits);
243:     ts->nonlinear_its += its; ts->linear_its += lits;
244:     VecCopy(cn->update,sol);
245:     ts->steps++;
246:     TSPostStep(ts);
247:     TSMonitor(ts,ts->steps,ts->ptime,sol);
248:   }

250:   *steps += ts->steps;
251:   *ptime  = ts->ptime;
252:   return(0);
253: }

255: /*------------------------------------------------------------*/
258: static PetscErrorCode TSDestroy_CN(TS ts)
259: {
260:   TS_CN          *cn = (TS_CN*)ts->data;

264:   if (cn->update) {VecDestroy(cn->update);}
265:   if (cn->func) {VecDestroy(cn->func);}
266:   if (cn->rhsfunc) {VecDestroy(cn->rhsfunc);}
267:   if (cn->rhsfunc_old) {VecDestroy(cn->rhsfunc_old);}
268:   if (cn->rhs) {VecDestroy(cn->rhs);}
269:   PetscFree(cn);
270:   return(0);
271: }

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

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

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

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

341:   PetscTypeCompare((PetscObject)B,MATMFFD,&flg);
342:   if (!flg) {
343:     MatScale(B,-0.5);
344:     if (ts->Alhs){
345:       MatAXPY(B,mdt,ts->Alhs,DIFFERENT_NONZERO_PATTERN); /* DIFFERENT_NONZERO_PATTERN? */
346:     } else {
347:       MatShift(B,mdt);
348:     }
349:   } else {
350:     SETERRQ(PETSC_ERR_SUP,"Matrix type MATMFFD is not supported yet"); /* ref TSScaleShiftMatrices() */
351:   }
352:   return(0);
353: }

355: /*
356:    This constructs the Jacobian needed for SNES 
357:      J = I/dt - 0.5*J_{F}   where J_{F} is the given Jacobian of F.
358:      x  - input vector
359:      AA - Jacobian matrix 
360:      BB - preconditioner matrix, usually the same as AA
361: */
364: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
365: {
366:   TS             ts = (TS) ctx;

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

373:   /* shift and scale Jacobian */
374:   TSScaleShiftMatrices_CN(ts,*AA,*BB,*str); /* Set BB = 1/dt*Alhs - 0.5*BB */
375:   return(0);
376: }

378: /* ------------------------------------------------------------*/
381: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
382: {
383:   TS_CN          *cn = (TS_CN*)ts->data;
385:   PetscTruth shelltype;

388:   VecDuplicate(ts->vec_sol,&cn->update);
389:   VecDuplicate(ts->vec_sol,&cn->rhs);
390: 
391:   /* build linear system to be solved */
392:   /* scale ts->Alhs = 1/dt*Alhs, ts->Arhs = 0.5*Arhs; set ts->A = Alhs - Arhs */
393:   PetscTypeCompare((PetscObject)ts->Arhs,MATSHELL,&shelltype);
394:   if (shelltype){
395:     TSSetKSPOperators_CN_No_Matrix(ts);
396:   } else {
397:     TSSetKSPOperators_CN_Matrix(ts);
398:   }
399:   KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
400:   return(0);
401: }

405: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
406: {
407:   TS_CN          *cn = (TS_CN*)ts->data;

411:   VecDuplicate(ts->vec_sol,&cn->update);
412:   VecDuplicate(ts->vec_sol,&cn->rhs);
413:   return(0);
414: }

418: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
419: {
420:   TS_CN          *cn = (TS_CN*)ts->data;

424:   VecDuplicate(ts->vec_sol,&cn->update);
425:   VecDuplicate(ts->vec_sol,&cn->func);
426:   VecDuplicate(ts->vec_sol,&cn->rhsfunc);
427:   VecDuplicate(ts->vec_sol,&cn->rhsfunc_old);
428:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
429:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
430:   cn->rhsfunc_time     = -100.0; /* cn->rhsfunc is not evaluated yet */
431:   cn->rhsfunc_old_time = -100.0;
432:   return(0);
433: }
434: /*------------------------------------------------------------*/

438: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
439: {
441:   return(0);
442: }

446: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
447: {
449:   return(0);
450: }

454: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
455: {
457:   return(0);
458: }

460: /* ------------------------------------------------------------ */
461: /*MC
462:       TSCN - ODE solver using the implicit Crank-Nicholson method

464:   Level: beginner

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

468: M*/
472: PetscErrorCode  TSCreate_CN(TS ts)
473: {
474:   TS_CN          *cn;

478:   ts->ops->destroy = TSDestroy_CN;
479:   ts->ops->view    = TSView_CN;

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

504:   PetscNewLog(ts,TS_CN,&cn);
505:   ts->data = (void*)cn;
506:   return(0);
507: }