Actual source code: cn.c
1: /*
2: Code for Timestepping with implicit Crank-Nicholson method.
3: THIS IS NOT YET COMPLETE -- DO NOT USE!!
4: */
5: #include src/ts/tsimpl.h
7: typedef struct {
8: Vec update; /* work vector where new solution is formed */
9: Vec func; /* work vector where F(t[i],u[i]) is stored */
10: Vec rhs; /* work vector for RHS; vec_sol/dt */
11: } TS_CN;
13: /*------------------------------------------------------------------------------*/
16: /*
17: TSComputeRHSFunctionEuler - Evaluates the right-hand-side function.
19: Note: If the user did not provide a function but merely a matrix,
20: this routine applies the matrix.
21: */
22: PetscErrorCode TSComputeRHSFunctionEuler(TS ts,PetscReal t,Vec x,Vec y)
23: {
25: PetscScalar neg_two = -2.0,neg_mdt = -1.0/ts->time_step;
32: if (ts->ops->rhsfunction) {
33: PetscStackPush("TS user right-hand-side function");
34: (*ts->ops->rhsfunction)(ts,t,x,y,ts->funP);
35: PetscStackPop;
36: return(0);
37: }
39: if (ts->ops->rhsmatrix) { /* assemble matrix for this timestep */
40: MatStructure flg;
41: PetscStackPush("TS user right-hand-side matrix function");
42: (*ts->ops->rhsmatrix)(ts,t,&ts->A,&ts->B,&flg,ts->jacP);
43: PetscStackPop;
44: }
45: MatMult(ts->A,x,y);
46: /* shift: y = y -2*x */
47: VecAXPY(&neg_two,x,y);
48: /* scale: y = y -2*x */
49: VecScale(&neg_mdt,y);
51: /* apply user-provided boundary conditions (only needed if these are time dependent) */
52: TSComputeRHSBoundaryConditions(ts,t,y);
54: return(0);
55: }
57: /*
58: Version for linear PDE where RHS does not depend on time. Has built a
59: single matrix that is to be used for all timesteps.
60: */
63: static PetscErrorCode TSStep_CN_Linear_Constant_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
64: {
65: TS_CN *cn = (TS_CN*)ts->data;
66: Vec sol = ts->vec_sol,update = cn->update;
67: Vec rhs = cn->rhs;
69: PetscInt i,max_steps = ts->max_steps,its;
70: PetscScalar dt = ts->time_step,two = 2.0;
71: KSP ksp;
74: TSGetKSP(ts,&ksp);
75: *steps = -ts->steps;
76: TSMonitor(ts,ts->steps,ts->ptime,sol);
78: /* set initial guess to be previous solution */
79: VecCopy(sol,update);
81: for (i=0; i<max_steps; i++) {
82: ts->ptime += ts->time_step;
83: if (ts->ptime > ts->max_time) break;
85: /* phase 1 - explicit step */
86: TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
87: VecAXPBY(&dt,&two,update,sol);
89: /* phase 2 - implicit step */
90: VecCopy(sol,rhs);
91: /* apply user-provided boundary conditions (only needed if they are time dependent) */
92: TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);
94: KSPSolve(ts->ksp,rhs,update);
95: KSPGetIterationNumber(ksp,&its);
96: ts->linear_its += PetscAbsInt(its);
97: VecCopy(update,sol);
98: ts->steps++;
99: TSMonitor(ts,ts->steps,ts->ptime,sol);
100: }
102: *steps += ts->steps;
103: *ptime = ts->ptime;
104: return(0);
105: }
106: /*
107: Version where matrix depends on time
108: */
111: static PetscErrorCode TSStep_CN_Linear_Variable_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
112: {
113: TS_CN *cn = (TS_CN*)ts->data;
114: Vec sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
116: PetscInt i,max_steps = ts->max_steps,its;
117: PetscScalar dt = ts->time_step,two = 2.0;
118: MatStructure str;
119: KSP ksp;
122: TSGetKSP(ts,&ksp);
123: *steps = -ts->steps;
124: TSMonitor(ts,ts->steps,ts->ptime,sol);
126: /* set initial guess to be previous solution */
127: VecCopy(sol,update);
129: for (i=0; i<max_steps; i++) {
130: ts->ptime += ts->time_step;
131: if (ts->ptime > ts->max_time) break;
132: /*
133: evaluate matrix function
134: */
135: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
136: TSScaleShiftMatrices(ts,ts->A,ts->B,str);
138: /* phase 1 - explicit step */
139: TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
140: VecAXPBY(&dt,&two,update,sol);
142: /* phase 2 - implicit step */
143: VecCopy(sol,rhs);
145: /* apply user-provided boundary conditions (only needed if they are time dependent) */
146: TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);
148: KSPSetOperators(ts->ksp,ts->A,ts->B,str);
149: KSPSolve(ts->ksp,rhs,update);
150: KSPGetIterationNumber(ksp,&its);
151: ts->linear_its += PetscAbsInt(its);
152: VecCopy(update,sol);
153: ts->steps++;
154: TSMonitor(ts,ts->steps,ts->ptime,sol);
155: }
157: *steps += ts->steps;
158: *ptime = ts->ptime;
159: return(0);
160: }
161: /*
162: Version for nonlinear PDE.
163: */
166: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
167: {
168: Vec sol = ts->vec_sol;
170: PetscInt i,max_steps = ts->max_steps,its,lits;
171: TS_CN *cn = (TS_CN*)ts->data;
172:
174: *steps = -ts->steps;
175: TSMonitor(ts,ts->steps,ts->ptime,sol);
177: for (i=0; i<max_steps; i++) {
178: ts->ptime += ts->time_step;
179: if (ts->ptime > ts->max_time) break;
180: VecCopy(sol,cn->update);
181: SNESSolve(ts->snes,cn->update);
182: SNESGetIterationNumber(ts->snes,&its);
183: SNESGetNumberLinearIterations(ts->snes,&lits);
184: ts->nonlinear_its += its; ts->linear_its += lits;
185: VecCopy(cn->update,sol);
186: ts->steps++;
187: TSMonitor(ts,ts->steps,ts->ptime,sol);
188: }
190: *steps += ts->steps;
191: *ptime = ts->ptime;
192: return(0);
193: }
195: /*------------------------------------------------------------*/
198: static PetscErrorCode TSDestroy_CN(TS ts)
199: {
200: TS_CN *cn = (TS_CN*)ts->data;
204: if (cn->update) {VecDestroy(cn->update);}
205: if (cn->func) {VecDestroy(cn->func);}
206: if (cn->rhs) {VecDestroy(cn->rhs);}
207: PetscFree(cn);
208: return(0);
209: }
211: /*
212: This defines the nonlinear equation that is to be solved with SNES
214: U^{n+1} - dt*F(U^{n+1}) - U^{n}
215: */
218: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
219: {
220: TS ts = (TS) ctx;
221: PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
223: PetscInt i,n;
226: /* apply user provided function */
227: TSComputeRHSFunction(ts,ts->ptime,x,y);
228: /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
229: VecGetArray(ts->vec_sol,&un);
230: VecGetArray(x,&unp1);
231: VecGetArray(y,&Funp1);
232: VecGetLocalSize(x,&n);
234: for (i=0; i<n; i++) {
235: Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
236: }
237: VecRestoreArray(ts->vec_sol,&un);
238: VecRestoreArray(x,&unp1);
239: VecRestoreArray(y,&Funp1);
240: return(0);
241: }
243: /*
244: This constructs the Jacobian needed for SNES
246: J = I/dt - J_{F} where J_{F} is the given Jacobian of F.
247: */
250: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
251: {
252: TS ts = (TS) ctx;
256: /* construct user's Jacobian */
257: TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str);
259: /* shift and scale Jacobian */
260: TSScaleShiftMatrices(ts,*AA,*BB,*str);
261: return(0);
262: }
264: /* ------------------------------------------------------------*/
267: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
268: {
269: TS_CN *cn = (TS_CN*)ts->data;
273: VecDuplicate(ts->vec_sol,&cn->update);
274: VecDuplicate(ts->vec_sol,&cn->rhs);
275:
276: /* build linear system to be solved */
277: TSScaleShiftMatrices(ts,ts->A,ts->B,SAME_NONZERO_PATTERN);
278: KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
279: return(0);
280: }
284: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
285: {
286: TS_CN *cn = (TS_CN*)ts->data;
290: VecDuplicate(ts->vec_sol,&cn->update);
291: VecDuplicate(ts->vec_sol,&cn->rhs);
292: return(0);
293: }
297: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
298: {
299: TS_CN *cn = (TS_CN*)ts->data;
303: VecDuplicate(ts->vec_sol,&cn->update);
304: VecDuplicate(ts->vec_sol,&cn->func);
305: SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
306: SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
307: return(0);
308: }
309: /*------------------------------------------------------------*/
313: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
314: {
318: KSPSetFromOptions(ts->ksp);
319: return(0);
320: }
324: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
325: {
329: SNESSetFromOptions(ts->snes);
330: return(0);
331: }
335: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
336: {
338: return(0);
339: }
341: /* ------------------------------------------------------------ */
342: /*MC
343: TS_CN - ODE solver using the implicit Crank-Nicholson method
345: Level: beginner
347: .seealso: TSCreate(), TS, TSSetType()
349: M*/
353: PetscErrorCode TSCreate_CN(TS ts)
354: {
355: TS_CN *cn;
357: KSP ksp;
360: ts->ops->destroy = TSDestroy_CN;
361: ts->ops->view = TSView_CN;
363: if (ts->problem_type == TS_LINEAR) {
364: if (!ts->A) {
365: SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
366: }
367: if (!ts->ops->rhsmatrix) {
368: ts->ops->setup = TSSetUp_CN_Linear_Constant_Matrix;
369: ts->ops->step = TSStep_CN_Linear_Constant_Matrix;
370: } else {
371: ts->ops->setup = TSSetUp_CN_Linear_Variable_Matrix;
372: ts->ops->step = TSStep_CN_Linear_Variable_Matrix;
373: }
374: ts->ops->setfromoptions = TSSetFromOptions_CN_Linear;
375: KSPCreate(ts->comm,&ts->ksp);
376: TSGetKSP(ts,&ksp);
377: KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
378: } else if (ts->problem_type == TS_NONLINEAR) {
379: ts->ops->setup = TSSetUp_CN_Nonlinear;
380: ts->ops->step = TSStep_CN_Nonlinear;
381: ts->ops->setfromoptions = TSSetFromOptions_CN_Nonlinear;
382: SNESCreate(ts->comm,&ts->snes);
383: } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");
385: PetscNew(TS_CN,&cn);
386: PetscLogObjectMemory(ts,sizeof(TS_CN));
387: PetscMemzero(cn,sizeof(TS_CN));
388: ts->data = (void*)cn;
390: return(0);
391: }