Actual source code: cn.c

  1: #define PETSCTS_DLL

  3: /*
  4:        Code for Timestepping with implicit Crank-Nicholson method.
  5: */
 6:  #include 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 A = B = 1/dt*A - 0.5*A */
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:   /* this function requires additional work! */
335:   PetscTypeCompare((PetscObject)A,MATMFFD,&flg);
336:   if (!flg) {
337:     MatScale(A,-0.5);
338:     if (ts->Alhs){
339:       MatAXPY(A,mdt,ts->Alhs,DIFFERENT_NONZERO_PATTERN); /* DIFFERENT_NONZERO_PATTERN? */
340:     } else {
341:       MatShift(A,mdt);
342:     }
343:   } else {
344:     SETERRQ(PETSC_ERR_SUP,"Matrix type MATMFFD is not supported yet"); /* ref TSScaleShiftMatrices() */
345:   }
346:   if (B != A && str != SAME_PRECONDITIONER) {
347:     SETERRQ(PETSC_ERR_SUP,"not supported yet");
348:   }
349:   return(0);
350: }

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

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

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

375: /* ------------------------------------------------------------*/
378: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
379: {
380:   TS_CN          *cn = (TS_CN*)ts->data;
382:   PetscTruth shelltype;

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

402: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
403: {
404:   TS_CN          *cn = (TS_CN*)ts->data;

408:   VecDuplicate(ts->vec_sol,&cn->update);
409:   VecDuplicate(ts->vec_sol,&cn->rhs);
410:   return(0);
411: }

415: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
416: {
417:   TS_CN          *cn = (TS_CN*)ts->data;

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

435: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
436: {
438:   return(0);
439: }

443: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
444: {
446:   return(0);
447: }

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

457: /* ------------------------------------------------------------ */
458: /*MC
459:       TS_CN - ODE solver using the implicit Crank-Nicholson method

461:   Level: beginner

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

465: M*/
469: PetscErrorCode  TSCreate_CN(TS ts)
470: {
471:   TS_CN          *cn;

475:   ts->ops->destroy = TSDestroy_CN;
476:   ts->ops->view    = TSView_CN;

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

499:   PetscNew(TS_CN,&cn);
500:   PetscLogObjectMemory(ts,sizeof(TS_CN));
501:   ts->data = (void*)cn;
502:   return(0);
503: }