Actual source code: cn.c
1: #define PETSCTS_DLL
3: /*
4: Code for Timestepping with implicit Crank-Nicholson method.
5: THIS IS NOT YET COMPLETE -- DO NOT USE!!
6: */
7: #include src/ts/tsimpl.h
9: typedef struct {
10: Vec update; /* work vector where new solution is formed */
11: Vec func; /* work vector where F(t[i],u[i]) is stored */
12: Vec rhs; /* work vector for RHS; vec_sol/dt */
13: } TS_CN;
15: /*------------------------------------------------------------------------------*/
18: /*
19: TSComputeRHSFunctionEuler - Evaluates the right-hand-side function.
21: Note: If the user did not provide a function but merely a matrix,
22: this routine applies the matrix.
23: */
24: PetscErrorCode TSComputeRHSFunctionEuler(TS ts,PetscReal t,Vec x,Vec y)
25: {
27: PetscScalar neg_two = -2.0,neg_mdt = -1.0/ts->time_step;
34: if (ts->ops->rhsfunction) {
35: PetscStackPush("TS user right-hand-side function");
36: (*ts->ops->rhsfunction)(ts,t,x,y,ts->funP);
37: PetscStackPop;
38: return(0);
39: }
41: if (ts->ops->rhsmatrix) { /* assemble matrix for this timestep */
42: MatStructure flg;
43: PetscStackPush("TS user right-hand-side matrix function");
44: (*ts->ops->rhsmatrix)(ts,t,&ts->A,&ts->B,&flg,ts->jacP);
45: PetscStackPop;
46: }
47: MatMult(ts->A,x,y);
48: /* shift: y = y -2*x */
49: VecAXPY(y,neg_two,x);
50: /* scale: y = y -2*x */
51: VecScale(y,neg_mdt);
53: /* apply user-provided boundary conditions (only needed if these are time dependent) */
54: TSComputeRHSBoundaryConditions(ts,t,y);
56: return(0);
57: }
59: /*
60: Version for linear PDE where RHS does not depend on time. Has built a
61: single matrix that is to be used for all timesteps.
62: */
65: static PetscErrorCode TSStep_CN_Linear_Constant_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
66: {
67: TS_CN *cn = (TS_CN*)ts->data;
68: Vec sol = ts->vec_sol,update = cn->update;
69: Vec rhs = cn->rhs;
71: PetscInt i,max_steps = ts->max_steps,its;
72: PetscScalar dt = ts->time_step,two = 2.0;
73: KSP ksp;
76: TSGetKSP(ts,&ksp);
77: *steps = -ts->steps;
78: TSMonitor(ts,ts->steps,ts->ptime,sol);
80: /* set initial guess to be previous solution */
81: VecCopy(sol,update);
83: for (i=0; i<max_steps; i++) {
84: ts->ptime += ts->time_step;
85: if (ts->ptime > ts->max_time) break;
87: /* phase 1 - explicit step */
88: TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
89: VecAXPBY(sol,dt,two,update);
91: /* phase 2 - implicit step */
92: VecCopy(sol,rhs);
93: /* apply user-provided boundary conditions (only needed if they are time dependent) */
94: TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);
96: KSPSolve(ts->ksp,rhs,update);
97: KSPGetIterationNumber(ksp,&its);
98: ts->linear_its += PetscAbsInt(its);
99: VecCopy(update,sol);
100: ts->steps++;
101: TSMonitor(ts,ts->steps,ts->ptime,sol);
102: }
104: *steps += ts->steps;
105: *ptime = ts->ptime;
106: return(0);
107: }
108: /*
109: Version where matrix depends on time
110: */
113: static PetscErrorCode TSStep_CN_Linear_Variable_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 dt = ts->time_step,two = 2.0;
120: MatStructure str;
121: KSP ksp;
124: TSGetKSP(ts,&ksp);
125: *steps = -ts->steps;
126: TSMonitor(ts,ts->steps,ts->ptime,sol);
128: /* set initial guess to be previous solution */
129: VecCopy(sol,update);
131: for (i=0; i<max_steps; i++) {
132: ts->ptime += ts->time_step;
133: if (ts->ptime > ts->max_time) break;
134: /*
135: evaluate matrix function
136: */
137: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
138: TSScaleShiftMatrices(ts,ts->A,ts->B,str);
140: /* phase 1 - explicit step */
141: TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
142: VecAXPBY(sol,dt,two,update);
144: /* phase 2 - implicit step */
145: VecCopy(sol,rhs);
147: /* apply user-provided boundary conditions (only needed if they are time dependent) */
148: TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);
150: KSPSetOperators(ts->ksp,ts->A,ts->B,str);
151: KSPSolve(ts->ksp,rhs,update);
152: KSPGetIterationNumber(ksp,&its);
153: ts->linear_its += PetscAbsInt(its);
154: VecCopy(update,sol);
155: ts->steps++;
156: TSMonitor(ts,ts->steps,ts->ptime,sol);
157: }
159: *steps += ts->steps;
160: *ptime = ts->ptime;
161: return(0);
162: }
163: /*
164: Version for nonlinear PDE.
165: */
168: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
169: {
170: Vec sol = ts->vec_sol;
172: PetscInt i,max_steps = ts->max_steps,its,lits;
173: TS_CN *cn = (TS_CN*)ts->data;
174:
176: *steps = -ts->steps;
177: TSMonitor(ts,ts->steps,ts->ptime,sol);
179: for (i=0; i<max_steps; i++) {
180: ts->ptime += ts->time_step;
181: if (ts->ptime > ts->max_time) break;
182: VecCopy(sol,cn->update);
183: SNESSolve(ts->snes,PETSC_NULL,cn->update);
184: SNESGetIterationNumber(ts->snes,&its);
185: SNESGetNumberLinearIterations(ts->snes,&lits);
186: ts->nonlinear_its += its; ts->linear_its += lits;
187: VecCopy(cn->update,sol);
188: ts->steps++;
189: TSMonitor(ts,ts->steps,ts->ptime,sol);
190: }
192: *steps += ts->steps;
193: *ptime = ts->ptime;
194: return(0);
195: }
197: /*------------------------------------------------------------*/
200: static PetscErrorCode TSDestroy_CN(TS ts)
201: {
202: TS_CN *cn = (TS_CN*)ts->data;
206: if (cn->update) {VecDestroy(cn->update);}
207: if (cn->func) {VecDestroy(cn->func);}
208: if (cn->rhs) {VecDestroy(cn->rhs);}
209: PetscFree(cn);
210: return(0);
211: }
213: /*
214: This defines the nonlinear equation that is to be solved with SNES
216: U^{n+1} - dt*F(U^{n+1}) - U^{n}
217: */
220: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
221: {
222: TS ts = (TS) ctx;
223: PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
225: PetscInt i,n;
228: /* apply user provided function */
229: TSComputeRHSFunction(ts,ts->ptime,x,y);
230: /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
231: VecGetArray(ts->vec_sol,&un);
232: VecGetArray(x,&unp1);
233: VecGetArray(y,&Funp1);
234: VecGetLocalSize(x,&n);
236: for (i=0; i<n; i++) {
237: Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
238: }
239: VecRestoreArray(ts->vec_sol,&un);
240: VecRestoreArray(x,&unp1);
241: VecRestoreArray(y,&Funp1);
242: return(0);
243: }
245: /*
246: This constructs the Jacobian needed for SNES
248: J = I/dt - J_{F} where J_{F} is the given Jacobian of F.
249: */
252: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
253: {
254: TS ts = (TS) ctx;
258: /* construct user's Jacobian */
259: TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str);
261: /* shift and scale Jacobian */
262: TSScaleShiftMatrices(ts,*AA,*BB,*str);
263: return(0);
264: }
266: /* ------------------------------------------------------------*/
269: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
270: {
271: TS_CN *cn = (TS_CN*)ts->data;
275: VecDuplicate(ts->vec_sol,&cn->update);
276: VecDuplicate(ts->vec_sol,&cn->rhs);
277:
278: /* build linear system to be solved */
279: TSScaleShiftMatrices(ts,ts->A,ts->B,SAME_NONZERO_PATTERN);
280: KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
281: return(0);
282: }
286: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
287: {
288: TS_CN *cn = (TS_CN*)ts->data;
292: VecDuplicate(ts->vec_sol,&cn->update);
293: VecDuplicate(ts->vec_sol,&cn->rhs);
294: return(0);
295: }
299: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
300: {
301: TS_CN *cn = (TS_CN*)ts->data;
305: VecDuplicate(ts->vec_sol,&cn->update);
306: VecDuplicate(ts->vec_sol,&cn->func);
307: SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
308: SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
309: return(0);
310: }
311: /*------------------------------------------------------------*/
315: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
316: {
320: KSPSetFromOptions(ts->ksp);
321: return(0);
322: }
326: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
327: {
331: SNESSetFromOptions(ts->snes);
332: return(0);
333: }
337: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
338: {
340: return(0);
341: }
343: /* ------------------------------------------------------------ */
344: /*MC
345: TS_CN - ODE solver using the implicit Crank-Nicholson method
347: Level: beginner
349: .seealso: TSCreate(), TS, TSSetType()
351: M*/
355: PetscErrorCode PETSCTS_DLLEXPORT TSCreate_CN(TS ts)
356: {
357: TS_CN *cn;
359: KSP ksp;
362: ts->ops->destroy = TSDestroy_CN;
363: ts->ops->view = TSView_CN;
365: if (ts->problem_type == TS_LINEAR) {
366: if (!ts->A) {
367: SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
368: }
369: if (!ts->ops->rhsmatrix) {
370: ts->ops->setup = TSSetUp_CN_Linear_Constant_Matrix;
371: ts->ops->step = TSStep_CN_Linear_Constant_Matrix;
372: } else {
373: ts->ops->setup = TSSetUp_CN_Linear_Variable_Matrix;
374: ts->ops->step = TSStep_CN_Linear_Variable_Matrix;
375: }
376: ts->ops->setfromoptions = TSSetFromOptions_CN_Linear;
377: KSPCreate(ts->comm,&ts->ksp);
378: TSGetKSP(ts,&ksp);
379: KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
380: } else if (ts->problem_type == TS_NONLINEAR) {
381: ts->ops->setup = TSSetUp_CN_Nonlinear;
382: ts->ops->step = TSStep_CN_Nonlinear;
383: ts->ops->setfromoptions = TSSetFromOptions_CN_Nonlinear;
384: SNESCreate(ts->comm,&ts->snes);
385: } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");
387: PetscNew(TS_CN,&cn);
388: PetscLogObjectMemory(ts,sizeof(TS_CN));
389: ts->data = (void*)cn;
391: return(0);
392: }