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: return(0);
54: }
56: /*
57: Version for linear PDE where RHS does not depend on time. Has built a
58: single matrix that is to be used for all timesteps.
59: */
62: static PetscErrorCode TSStep_CN_Linear_Constant_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
63: {
64: TS_CN *cn = (TS_CN*)ts->data;
65: Vec sol = ts->vec_sol,update = cn->update;
66: Vec rhs = cn->rhs;
68: PetscInt i,max_steps = ts->max_steps,its;
69: PetscScalar dt = ts->time_step,two = 2.0;
70: KSP ksp;
73: TSGetKSP(ts,&ksp);
74: *steps = -ts->steps;
75: TSMonitor(ts,ts->steps,ts->ptime,sol);
77: /* set initial guess to be previous solution */
78: VecCopy(sol,update);
80: for (i=0; i<max_steps; i++) {
81: ts->ptime += ts->time_step;
82: if (ts->ptime > ts->max_time) break;
84: /* phase 1 - explicit step */
85: TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
86: VecAXPBY(sol,dt,two,update);
88: /* phase 2 - implicit step */
89: VecCopy(sol,rhs);
91: KSPSolve(ts->ksp,rhs,update);
92: KSPGetIterationNumber(ksp,&its);
93: ts->linear_its += PetscAbsInt(its);
94: VecCopy(update,sol);
95: ts->steps++;
96: TSMonitor(ts,ts->steps,ts->ptime,sol);
97: }
99: *steps += ts->steps;
100: *ptime = ts->ptime;
101: return(0);
102: }
103: /*
104: Version where matrix depends on time
105: */
108: static PetscErrorCode TSStep_CN_Linear_Variable_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
109: {
110: TS_CN *cn = (TS_CN*)ts->data;
111: Vec sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
113: PetscInt i,max_steps = ts->max_steps,its;
114: PetscScalar dt = ts->time_step,two = 2.0;
115: MatStructure str;
116: KSP ksp;
119: TSGetKSP(ts,&ksp);
120: *steps = -ts->steps;
121: TSMonitor(ts,ts->steps,ts->ptime,sol);
123: /* set initial guess to be previous solution */
124: VecCopy(sol,update);
126: for (i=0; i<max_steps; i++) {
127: ts->ptime += ts->time_step;
128: if (ts->ptime > ts->max_time) break;
129: /*
130: evaluate matrix function
131: */
132: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
133: TSScaleShiftMatrices(ts,ts->A,ts->B,str);
135: /* phase 1 - explicit step */
136: TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
137: VecAXPBY(sol,dt,two,update);
139: /* phase 2 - implicit step */
140: VecCopy(sol,rhs);
142: KSPSetOperators(ts->ksp,ts->A,ts->B,str);
143: KSPSolve(ts->ksp,rhs,update);
144: KSPGetIterationNumber(ksp,&its);
145: ts->linear_its += PetscAbsInt(its);
146: VecCopy(update,sol);
147: ts->steps++;
148: TSMonitor(ts,ts->steps,ts->ptime,sol);
149: }
151: *steps += ts->steps;
152: *ptime = ts->ptime;
153: return(0);
154: }
155: /*
156: Version for nonlinear PDE.
157: */
160: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
161: {
162: Vec sol = ts->vec_sol;
164: PetscInt i,max_steps = ts->max_steps,its,lits;
165: TS_CN *cn = (TS_CN*)ts->data;
166:
168: *steps = -ts->steps;
169: TSMonitor(ts,ts->steps,ts->ptime,sol);
171: for (i=0; i<max_steps; i++) {
172: ts->ptime += ts->time_step;
173: if (ts->ptime > ts->max_time) break;
174: VecCopy(sol,cn->update);
175: SNESSolve(ts->snes,PETSC_NULL,cn->update);
176: SNESGetIterationNumber(ts->snes,&its);
177: SNESGetNumberLinearIterations(ts->snes,&lits);
178: ts->nonlinear_its += its; ts->linear_its += lits;
179: VecCopy(cn->update,sol);
180: ts->steps++;
181: TSMonitor(ts,ts->steps,ts->ptime,sol);
182: }
184: *steps += ts->steps;
185: *ptime = ts->ptime;
186: return(0);
187: }
189: /*------------------------------------------------------------*/
192: static PetscErrorCode TSDestroy_CN(TS ts)
193: {
194: TS_CN *cn = (TS_CN*)ts->data;
198: if (cn->update) {VecDestroy(cn->update);}
199: if (cn->func) {VecDestroy(cn->func);}
200: if (cn->rhs) {VecDestroy(cn->rhs);}
201: PetscFree(cn);
202: return(0);
203: }
205: /*
206: This defines the nonlinear equation that is to be solved with SNES
208: U^{n+1} - dt*F(U^{n+1}) - U^{n}
209: */
212: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
213: {
214: TS ts = (TS) ctx;
215: PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
217: PetscInt i,n;
220: /* apply user provided function */
221: TSComputeRHSFunction(ts,ts->ptime,x,y);
222: /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
223: VecGetArray(ts->vec_sol,&un);
224: VecGetArray(x,&unp1);
225: VecGetArray(y,&Funp1);
226: VecGetLocalSize(x,&n);
228: for (i=0; i<n; i++) {
229: Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
230: }
231: VecRestoreArray(ts->vec_sol,&un);
232: VecRestoreArray(x,&unp1);
233: VecRestoreArray(y,&Funp1);
234: return(0);
235: }
237: /*
238: This constructs the Jacobian needed for SNES
240: J = I/dt - J_{F} where J_{F} is the given Jacobian of F.
241: */
244: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
245: {
246: TS ts = (TS) ctx;
250: /* construct user's Jacobian */
251: TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str);
253: /* shift and scale Jacobian */
254: TSScaleShiftMatrices(ts,*AA,*BB,*str);
255: return(0);
256: }
258: /* ------------------------------------------------------------*/
261: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
262: {
263: TS_CN *cn = (TS_CN*)ts->data;
267: VecDuplicate(ts->vec_sol,&cn->update);
268: VecDuplicate(ts->vec_sol,&cn->rhs);
269:
270: /* build linear system to be solved */
271: TSScaleShiftMatrices(ts,ts->A,ts->B,SAME_NONZERO_PATTERN);
272: KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
273: return(0);
274: }
278: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
279: {
280: TS_CN *cn = (TS_CN*)ts->data;
284: VecDuplicate(ts->vec_sol,&cn->update);
285: VecDuplicate(ts->vec_sol,&cn->rhs);
286: return(0);
287: }
291: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
292: {
293: TS_CN *cn = (TS_CN*)ts->data;
297: VecDuplicate(ts->vec_sol,&cn->update);
298: VecDuplicate(ts->vec_sol,&cn->func);
299: SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
300: SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
301: return(0);
302: }
303: /*------------------------------------------------------------*/
307: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
308: {
312: KSPSetFromOptions(ts->ksp);
313: return(0);
314: }
318: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
319: {
323: SNESSetFromOptions(ts->snes);
324: return(0);
325: }
329: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
330: {
332: return(0);
333: }
335: /* ------------------------------------------------------------ */
336: /*MC
337: TS_CN - ODE solver using the implicit Crank-Nicholson method
339: Level: beginner
341: .seealso: TSCreate(), TS, TSSetType()
343: M*/
347: PetscErrorCode TSCreate_CN(TS ts)
348: {
349: TS_CN *cn;
351: KSP ksp;
354: ts->ops->destroy = TSDestroy_CN;
355: ts->ops->view = TSView_CN;
357: if (ts->problem_type == TS_LINEAR) {
358: if (!ts->A) {
359: SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
360: }
361: if (!ts->ops->rhsmatrix) {
362: ts->ops->setup = TSSetUp_CN_Linear_Constant_Matrix;
363: ts->ops->step = TSStep_CN_Linear_Constant_Matrix;
364: } else {
365: ts->ops->setup = TSSetUp_CN_Linear_Variable_Matrix;
366: ts->ops->step = TSStep_CN_Linear_Variable_Matrix;
367: }
368: ts->ops->setfromoptions = TSSetFromOptions_CN_Linear;
369: KSPCreate(ts->comm,&ts->ksp);
370: TSGetKSP(ts,&ksp);
371: KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
372: } else if (ts->problem_type == TS_NONLINEAR) {
373: ts->ops->setup = TSSetUp_CN_Nonlinear;
374: ts->ops->step = TSStep_CN_Nonlinear;
375: ts->ops->setfromoptions = TSSetFromOptions_CN_Nonlinear;
376: SNESCreate(ts->comm,&ts->snes);
377: } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");
379: PetscNew(TS_CN,&cn);
380: PetscLogObjectMemory(ts,sizeof(TS_CN));
381: ts->data = (void*)cn;
383: SETERRQ(PETSC_ERR_SUP,"The code for Crank-Nicholson is not complete\n emai petsc-maint@mcs.anl.gov for more info");
384: return(0);
385: }