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: }