Actual source code: cn.c
1: /*$Id: cn.c,v 1.35 2001/09/11 16:34:19 bsmith Exp $*/
2: /*
3: Code for Timestepping with implicit Crank-Nicholson method.
4: THIS IS NOT YET COMPLETE -- DO NOT USE!!
5: */
6: #include src/ts/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 rhs; /* work vector for RHS; vec_sol/dt */
12: } TS_CN;
14: /*------------------------------------------------------------------------------*/
17: /*
18: TSComputeRHSFunctionEuler - Evaluates the right-hand-side function.
20: Note: If the user did not provide a function but merely a matrix,
21: this routine applies the matrix.
22: */
23: int TSComputeRHSFunctionEuler(TS ts,PetscReal t,Vec x,Vec y)
24: {
25: int ierr;
26: PetscScalar neg_two = -2.0,neg_mdt = -1.0/ts->time_step;
33: if (ts->ops->rhsfunction) {
34: PetscStackPush("TS user right-hand-side function");
35: (*ts->ops->rhsfunction)(ts,t,x,y,ts->funP);
36: PetscStackPop;
37: return(0);
38: }
40: if (ts->ops->rhsmatrix) { /* assemble matrix for this timestep */
41: MatStructure flg;
42: PetscStackPush("TS user right-hand-side matrix function");
43: (*ts->ops->rhsmatrix)(ts,t,&ts->A,&ts->B,&flg,ts->jacP);
44: PetscStackPop;
45: }
46: MatMult(ts->A,x,y);
47: /* shift: y = y -2*x */
48: VecAXPY(&neg_two,x,y);
49: /* scale: y = y -2*x */
50: VecScale(&neg_mdt,y);
52: /* apply user-provided boundary conditions (only needed if these are time dependent) */
53: TSComputeRHSBoundaryConditions(ts,t,y);
55: return(0);
56: }
58: /*
59: Version for linear PDE where RHS does not depend on time. Has built a
60: single matrix that is to be used for all timesteps.
61: */
64: static int TSStep_CN_Linear_Constant_Matrix(TS ts,int *steps,PetscReal *ptime)
65: {
66: TS_CN *cn = (TS_CN*)ts->data;
67: Vec sol = ts->vec_sol,update = cn->update;
68: Vec rhs = cn->rhs;
69: int ierr,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: KSPSetRhs(ts->ksp,rhs);
95: KSPSetSolution(ts->ksp,update);
96: KSPSolve(ts->ksp);
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 int TSStep_CN_Linear_Variable_Matrix(TS ts,int *steps,PetscReal *ptime)
114: {
115: TS_CN *cn = (TS_CN*)ts->data;
116: Vec sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
117: int ierr,i,max_steps = ts->max_steps,its;
118: PetscScalar dt = ts->time_step,two = 2.0,neg_dt = -1.0*ts->time_step;
119: MatStructure str;
120: KSP ksp;
123: TSGetKSP(ts,&ksp);
124: *steps = -ts->steps;
125: TSMonitor(ts,ts->steps,ts->ptime,sol);
127: /* set initial guess to be previous solution */
128: VecCopy(sol,update);
130: for (i=0; i<max_steps; i++) {
131: ts->ptime += ts->time_step;
132: if (ts->ptime > ts->max_time) break;
133: /*
134: evaluate matrix function
135: */
136: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
137: MatScale(&neg_dt,ts->A);
138: MatShift(&two,ts->A);
139: if (ts->B != ts->A && str != SAME_PRECONDITIONER) {
140: MatScale(&neg_dt,ts->B);
141: MatShift(&two,ts->B);
142: }
144: /* phase 1 - explicit step */
145: TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
146: VecAXPBY(&dt,&two,update,sol);
148: /* phase 2 - implicit step */
149: VecCopy(sol,rhs);
151: /* apply user-provided boundary conditions (only needed if they are time dependent) */
152: TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);
154: KSPSetOperators(ts->ksp,ts->A,ts->B,str);
155: KSPSetRhs(ts->ksp,rhs);
156: KSPSetSolution(ts->ksp,update);
157: KSPSolve(ts->ksp);
158: KSPGetIterationNumber(ksp,&its);
159: ts->linear_its += PetscAbsInt(its);
160: VecCopy(update,sol);
161: ts->steps++;
162: TSMonitor(ts,ts->steps,ts->ptime,sol);
163: }
165: *steps += ts->steps;
166: *ptime = ts->ptime;
167: return(0);
168: }
169: /*
170: Version for nonlinear PDE.
171: */
174: static int TSStep_CN_Nonlinear(TS ts,int *steps,PetscReal *ptime)
175: {
176: Vec sol = ts->vec_sol;
177: int ierr,i,max_steps = ts->max_steps,its,lits;
178: TS_CN *cn = (TS_CN*)ts->data;
179:
181: *steps = -ts->steps;
182: TSMonitor(ts,ts->steps,ts->ptime,sol);
184: for (i=0; i<max_steps; i++) {
185: ts->ptime += ts->time_step;
186: if (ts->ptime > ts->max_time) break;
187: VecCopy(sol,cn->update);
188: SNESSolve(ts->snes,cn->update);
189: SNESGetIterationNumber(ts->snes,&its);
190: SNESGetNumberLinearIterations(ts->snes,&lits);
191: ts->nonlinear_its += its; ts->linear_its += lits;
192: VecCopy(cn->update,sol);
193: ts->steps++;
194: TSMonitor(ts,ts->steps,ts->ptime,sol);
195: }
197: *steps += ts->steps;
198: *ptime = ts->ptime;
199: return(0);
200: }
202: /*------------------------------------------------------------*/
205: static int TSDestroy_CN(TS ts)
206: {
207: TS_CN *cn = (TS_CN*)ts->data;
208: int ierr;
211: if (cn->update) {VecDestroy(cn->update);}
212: if (cn->func) {VecDestroy(cn->func);}
213: if (cn->rhs) {VecDestroy(cn->rhs);}
214: PetscFree(cn);
215: return(0);
216: }
218: /*
219: This defines the nonlinear equation that is to be solved with SNES
221: U^{n+1} - dt*F(U^{n+1}) - U^{n}
222: */
225: int TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
226: {
227: TS ts = (TS) ctx;
228: PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
229: int ierr,i,n;
232: /* apply user provided function */
233: TSComputeRHSFunction(ts,ts->ptime,x,y);
234: /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
235: VecGetArray(ts->vec_sol,&un);
236: VecGetArray(x,&unp1);
237: VecGetArray(y,&Funp1);
238: VecGetLocalSize(x,&n);
240: for (i=0; i<n; i++) {
241: Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
242: }
243: VecRestoreArray(ts->vec_sol,&un);
244: VecRestoreArray(x,&unp1);
245: VecRestoreArray(y,&Funp1);
246: return(0);
247: }
249: /*
250: This constructs the Jacobian needed for SNES
252: J = I/dt - J_{F} where J_{F} is the given Jacobian of F.
253: */
256: int TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
257: {
258: TS ts = (TS) ctx;
259: int ierr;
260: PetscScalar mone = -1.0,mdt = 1.0/ts->time_step;
263: /* construct user's Jacobian */
264: TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str);
266: /* shift and scale Jacobian */
267: MatScale(&mone,*AA);
268: MatShift(&mdt,*AA);
269: if (*BB != *AA && *str != SAME_PRECONDITIONER) {
270: MatScale(&mone,*BB);
271: MatShift(&mdt,*BB);
272: }
274: return(0);
275: }
277: /* ------------------------------------------------------------*/
280: static int TSSetUp_CN_Linear_Constant_Matrix(TS ts)
281: {
282: TS_CN *cn = (TS_CN*)ts->data;
283: int ierr;
284: PetscScalar two = 2.0,neg_dt = -1.0*ts->time_step;
287: VecDuplicate(ts->vec_sol,&cn->update);
288: VecDuplicate(ts->vec_sol,&cn->rhs);
289:
290: /* build linear system to be solved */
291: MatScale(&neg_dt,ts->A);
292: MatShift(&two,ts->A);
293: if (ts->A != ts->B) {
294: MatScale(&neg_dt,ts->B);
295: MatShift(&two,ts->B);
296: }
297: KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
298: return(0);
299: }
303: static int TSSetUp_CN_Linear_Variable_Matrix(TS ts)
304: {
305: TS_CN *cn = (TS_CN*)ts->data;
306: int ierr;
309: VecDuplicate(ts->vec_sol,&cn->update);
310: VecDuplicate(ts->vec_sol,&cn->rhs);
311: return(0);
312: }
316: static int TSSetUp_CN_Nonlinear(TS ts)
317: {
318: TS_CN *cn = (TS_CN*)ts->data;
319: int ierr;
322: VecDuplicate(ts->vec_sol,&cn->update);
323: VecDuplicate(ts->vec_sol,&cn->func);
324: SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
325: SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
326: return(0);
327: }
328: /*------------------------------------------------------------*/
332: static int TSSetFromOptions_CN_Linear(TS ts)
333: {
337: KSPSetFromOptions(ts->ksp);
338: return(0);
339: }
343: static int TSSetFromOptions_CN_Nonlinear(TS ts)
344: {
348: SNESSetFromOptions(ts->snes);
349: return(0);
350: }
354: static int TSView_CN(TS ts,PetscViewer viewer)
355: {
357: return(0);
358: }
360: /* ------------------------------------------------------------ */
361: /*MC
362: TS_CN - ODE solver using the implicit Crank-Nicholson method
364: .seealso: TSCreate(), TS, TSSetType()
366: M*/
367: EXTERN_C_BEGIN
370: int TSCreate_CN(TS ts)
371: {
372: TS_CN *cn;
373: int ierr;
374: KSP ksp;
377: ts->ops->destroy = TSDestroy_CN;
378: ts->ops->view = TSView_CN;
380: if (ts->problem_type == TS_LINEAR) {
381: if (!ts->A) {
382: SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
383: }
384: if (!ts->ops->rhsmatrix) {
385: ts->ops->setup = TSSetUp_CN_Linear_Constant_Matrix;
386: ts->ops->step = TSStep_CN_Linear_Constant_Matrix;
387: } else {
388: ts->ops->setup = TSSetUp_CN_Linear_Variable_Matrix;
389: ts->ops->step = TSStep_CN_Linear_Variable_Matrix;
390: }
391: ts->ops->setfromoptions = TSSetFromOptions_CN_Linear;
392: KSPCreate(ts->comm,&ts->ksp);
393: TSGetKSP(ts,&ksp);
394: KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
395: } else if (ts->problem_type == TS_NONLINEAR) {
396: ts->ops->setup = TSSetUp_CN_Nonlinear;
397: ts->ops->step = TSStep_CN_Nonlinear;
398: ts->ops->setfromoptions = TSSetFromOptions_CN_Nonlinear;
399: SNESCreate(ts->comm,&ts->snes);
400: } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");
402: PetscNew(TS_CN,&cn);
403: PetscLogObjectMemory(ts,sizeof(TS_CN));
404: PetscMemzero(cn,sizeof(TS_CN));
405: ts->data = (void*)cn;
407: return(0);
408: }
409: EXTERN_C_END