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