Actual source code: cn.c
1: #define PETSCTS_DLL
3: /*
4: Code for Timestepping with implicit Crank-Nicholson method.
5: */
6: #include 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: /* set rhs = (1/dt*Alhs + 0.5*Arhs)*sol */
130: MatMult(ts->Arhs,sol,rhs); /* rhs = 0.5*Arhs*sol */
131: if (ts->Alhs){
132: MatMultAdd(ts->Alhs,sol,rhs,rhs); /* rhs = rhs + 1/dt*Alhs*sol */
133: } else {
134: VecAXPY(rhs,mdt,sol); /* rhs = rhs + 1/dt*sol */
135: }
137: ts->ptime += ts->time_step;
138: if (ts->ptime > ts->max_time) break;
140: /* solve (1/dt*Alhs - 0.5*Arhs)*update = rhs */
141: KSPSolve(ts->ksp,rhs,update);
142: KSPGetIterationNumber(ts->ksp,&its);
143: ts->linear_its += PetscAbsInt(its);
144: VecCopy(update,sol);
145: ts->steps++;
146: TSMonitor(ts,ts->steps,ts->ptime,sol);
147: } *steps += ts->steps;
148: *ptime = ts->ptime;
149: return(0);
150: }
151: /*
152: Version where matrix depends on time
153: */
156: static PetscErrorCode TSStep_CN_Linear_Variable_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
157: {
158: TS_CN *cn = (TS_CN*)ts->data;
159: Vec sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
161: PetscInt i,max_steps = ts->max_steps,its;
162: PetscScalar mdt = 1.0/ts->time_step;
163: PetscReal t_mid;
164: MatStructure str;
167: *steps = -ts->steps;
168: TSMonitor(ts,ts->steps,ts->ptime,sol);
170: /* set initial guess to be previous solution */
171: VecCopy(sol,update);
173: for (i=0; i<max_steps; i++) {
174: /* set rhs = (1/dt*Alhs(t_mid) + 0.5*Arhs(t_n)) * sol */
175: if (i==0){
176: /* evaluate 0.5*Arhs(t_0) */
177: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->Arhs,PETSC_NULL,&str,ts->jacP);
178: MatScale(ts->Arhs,0.5);
179: }
180: if (ts->Alhs){
181: /* evaluate Alhs(t_mid) */
182: t_mid = ts->ptime+ts->time_step/2.0;
183: (*ts->ops->lhsmatrix)(ts,t_mid,&ts->Alhs,PETSC_NULL,&str,ts->jacPlhs);
184: MatMult(ts->Alhs,sol,rhs); /* rhs = Alhs_mid*sol */
185: VecScale(rhs,mdt); /* rhs = 1/dt*Alhs_mid*sol */
186: MatMultAdd(ts->Arhs,sol,rhs,rhs); /* rhs = rhs + 0.5*Arhs_mid*sol */
187: } else {
188: MatMult(ts->Arhs,sol,rhs); /* rhs = 0.5*Arhs_n*sol */
189: VecAXPY(rhs,mdt,sol); /* rhs = rhs + 1/dt*sol */
190: }
192: ts->ptime += ts->time_step;
193: if (ts->ptime > ts->max_time) break;
195: /* evaluate Arhs at current ptime t_{n+1} */
196: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->Arhs,PETSC_NULL,&str,ts->jacP);
197: TSSetKSPOperators_CN_Matrix(ts);
199: KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
200: KSPSolve(ts->ksp,rhs,update);
201: KSPGetIterationNumber(ts->ksp,&its);
202: ts->linear_its += PetscAbsInt(its);
203: VecCopy(update,sol);
204: ts->steps++;
205: TSMonitor(ts,ts->steps,ts->ptime,sol);
206: }
208: *steps += ts->steps;
209: *ptime = ts->ptime;
210: return(0);
211: }
212: /*
213: Version for nonlinear PDE.
214: */
217: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
218: {
219: Vec sol = ts->vec_sol;
221: PetscInt i,max_steps = ts->max_steps,its,lits;
222: TS_CN *cn = (TS_CN*)ts->data;
223:
225: *steps = -ts->steps;
226: TSMonitor(ts,ts->steps,ts->ptime,sol);
228: for (i=0; i<max_steps; i++) {
229: ts->ptime += ts->time_step;
230: if (ts->ptime > ts->max_time) break;
231: VecCopy(sol,cn->update);
232: SNESSolve(ts->snes,PETSC_NULL,cn->update);
233: SNESGetIterationNumber(ts->snes,&its);
234: SNESGetNumberLinearIterations(ts->snes,&lits);
235: ts->nonlinear_its += its; ts->linear_its += lits;
236: VecCopy(cn->update,sol);
237: ts->steps++;
238: TSMonitor(ts,ts->steps,ts->ptime,sol);
239: }
241: *steps += ts->steps;
242: *ptime = ts->ptime;
243: return(0);
244: }
246: /*------------------------------------------------------------*/
249: static PetscErrorCode TSDestroy_CN(TS ts)
250: {
251: TS_CN *cn = (TS_CN*)ts->data;
255: if (cn->update) {VecDestroy(cn->update);}
256: if (cn->func) {VecDestroy(cn->func);}
257: if (cn->rhsfunc) {VecDestroy(cn->rhsfunc);}
258: if (cn->rhsfunc_old) {VecDestroy(cn->rhsfunc_old);}
259: if (cn->rhs) {VecDestroy(cn->rhs);}
260: PetscFree(cn);
261: return(0);
262: }
264: /*
265: This defines the nonlinear equation that is to be solved with SNES
266: 1/dt*Alhs*(U^{n+1} - U^{n}) - 0.5*(F(U^{n+1}) + F(U^{n}))
267: */
270: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
271: {
272: TS ts = (TS) ctx;
273: PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1,*Fun,*yarray;
275: PetscInt i,n;
276: TS_CN *cn = (TS_CN*)ts->data;
279: /* apply user provided function */
280: if (cn->rhsfunc_time == (ts->ptime - ts->time_step)){
281: /* printf(" copy rhsfunc to rhsfunc_old, then eval rhsfunc\n"); */
282: VecCopy(cn->rhsfunc,cn->rhsfunc_old);
283: cn->rhsfunc_old_time = cn->rhsfunc_time;
284: } else if (cn->rhsfunc_time != ts->ptime && cn->rhsfunc_old_time != ts->ptime-ts->time_step){
285: /* printf(" eval both rhsfunc_old and rhsfunc\n"); */
286: TSComputeRHSFunction(ts,ts->ptime-ts->time_step,ts->vec_sol,cn->rhsfunc_old); /* rhsfunc_old=F(U^{n}) */
287: cn->rhsfunc_old_time = ts->ptime - ts->time_step;
288: }
289:
290: if (ts->Alhs){
291: /* compute y=Alhs*(U^{n+1} - U^{n}) with cn->rhsfunc as workspce */
292: VecWAXPY(cn->rhsfunc,-1.0,ts->vec_sol,x);
293: MatMult(ts->Alhs,cn->rhsfunc,y);
294: }
296: TSComputeRHSFunction(ts,ts->ptime,x,cn->rhsfunc); /* rhsfunc = F(U^{n+1}) */
297: cn->rhsfunc_time = ts->ptime;
298:
299: VecGetArray(ts->vec_sol,&un); /* U^{n} */
300: VecGetArray(x,&unp1); /* U^{n+1} */
301: VecGetArray(cn->rhsfunc,&Funp1);
302: VecGetArray(cn->rhsfunc_old,&Fun);
303: VecGetArray(y,&yarray);
304: VecGetLocalSize(x,&n);
305: if (ts->Alhs){
306: for (i=0; i<n; i++) {
307: yarray[i] = mdt*yarray[i] - 0.5*(Funp1[i] + Fun[i]);
308: }
309: } else {
310: for (i=0; i<n; i++) {
311: yarray[i] = mdt*(unp1[i] - un[i]) - 0.5*(Funp1[i] + Fun[i]);
312: }
313: }
314: VecRestoreArray(ts->vec_sol,&un);
315: VecRestoreArray(x,&unp1);
316: VecRestoreArray(cn->rhsfunc,&Funp1);
317: VecRestoreArray(cn->rhsfunc_old,&Fun);
318: VecRestoreArray(y,&yarray);
319: return(0);
320: }
322: /* Set A = B = 1/dt*A - 0.5*A */
325: PetscErrorCode TSScaleShiftMatrices_CN(TS ts,Mat A,Mat B,MatStructure str)
326: {
327: PetscTruth flg;
329: PetscScalar mdt = 1.0/ts->time_step;
332: /* this function requires additional work! */
333: PetscTypeCompare((PetscObject)A,MATMFFD,&flg);
334: if (!flg) {
335: MatScale(A,-0.5);
336: if (ts->Alhs){
337: MatAXPY(A,mdt,ts->Alhs,DIFFERENT_NONZERO_PATTERN); /* DIFFERENT_NONZERO_PATTERN? */
338: } else {
339: MatShift(A,mdt);
340: }
341: } else {
342: SETERRQ(PETSC_ERR_SUP,"Matrix type MATMFFD is not supported yet"); /* ref TSScaleShiftMatrices() */
343: }
344: if (B != A && str != SAME_PRECONDITIONER) {
345: SETERRQ(PETSC_ERR_SUP,"not supported yet");
346: }
347: return(0);
348: }
350: /*
351: This constructs the Jacobian needed for SNES
352: J = I/dt - 0.5*J_{F} where J_{F} is the given Jacobian of F.
353: x - input vector
354: AA - Jacobian matrix
355: BB - preconditioner matrix, usually the same as AA
356: */
359: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
360: {
361: TS ts = (TS) ctx;
365: /* construct user's Jacobian */
366: TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str); /* AA = J_{F} */
368: /* shift and scale Jacobian */
369: TSScaleShiftMatrices_CN(ts,*AA,*BB,*str); /* Set AA = 1/dt*Alhs - 0.5*AA */
370: return(0);
371: }
373: /* ------------------------------------------------------------*/
376: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
377: {
378: TS_CN *cn = (TS_CN*)ts->data;
380: PetscTruth shelltype;
383: KSPSetFromOptions(ts->ksp);
384: VecDuplicate(ts->vec_sol,&cn->update);
385: VecDuplicate(ts->vec_sol,&cn->rhs);
386:
387: /* build linear system to be solved */
388: /* scale ts->Alhs = 1/dt*Alhs, ts->Arhs = 0.5*Arhs; set ts->A = Alhs - Arhs */
389: PetscTypeCompare((PetscObject)ts->Arhs,MATSHELL,&shelltype);
390: if (shelltype){
391: TSSetKSPOperators_CN_No_Matrix(ts);
392: } else {
393: TSSetKSPOperators_CN_Matrix(ts);
394: }
395: KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
396: return(0);
397: }
401: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
402: {
403: TS_CN *cn = (TS_CN*)ts->data;
407: VecDuplicate(ts->vec_sol,&cn->update);
408: VecDuplicate(ts->vec_sol,&cn->rhs);
409: return(0);
410: }
414: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
415: {
416: TS_CN *cn = (TS_CN*)ts->data;
420: VecDuplicate(ts->vec_sol,&cn->update);
421: VecDuplicate(ts->vec_sol,&cn->func);
422: VecDuplicate(ts->vec_sol,&cn->rhsfunc);
423: VecDuplicate(ts->vec_sol,&cn->rhsfunc_old);
424: SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
425: SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
426: cn->rhsfunc_time = -100.0; /* cn->rhsfunc is not evaluated yet */
427: cn->rhsfunc_old_time = -100.0;
428: return(0);
429: }
430: /*------------------------------------------------------------*/
434: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
435: {
437: return(0);
438: }
442: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
443: {
445: return(0);
446: }
450: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
451: {
453: return(0);
454: }
456: /* ------------------------------------------------------------ */
457: /*MC
458: TS_CN - ODE solver using the implicit Crank-Nicholson method
460: Level: beginner
462: .seealso: TSCreate(), TS, TSSetType()
464: M*/
468: PetscErrorCode TSCreate_CN(TS ts)
469: {
470: TS_CN *cn;
474: ts->ops->destroy = TSDestroy_CN;
475: ts->ops->view = TSView_CN;
477: if (ts->problem_type == TS_LINEAR) {
478: if (!ts->Arhs) {
479: SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
480: }
481: if (!ts->ops->rhsmatrix) {
482: ts->ops->setup = TSSetUp_CN_Linear_Constant_Matrix;
483: ts->ops->step = TSStep_CN_Linear_Constant_Matrix;
484: } else {
485: ts->ops->setup = TSSetUp_CN_Linear_Variable_Matrix;
486: ts->ops->step = TSStep_CN_Linear_Variable_Matrix;
487: }
488: ts->ops->setfromoptions = TSSetFromOptions_CN_Linear;
489: KSPCreate(ts->comm,&ts->ksp);
490: KSPSetInitialGuessNonzero(ts->ksp,PETSC_TRUE);
491: } else if (ts->problem_type == TS_NONLINEAR) {
492: ts->ops->setup = TSSetUp_CN_Nonlinear;
493: ts->ops->step = TSStep_CN_Nonlinear;
494: ts->ops->setfromoptions = TSSetFromOptions_CN_Nonlinear;
495: SNESCreate(ts->comm,&ts->snes);
496: } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");
498: PetscNew(TS_CN,&cn);
499: PetscLogObjectMemory(ts,sizeof(TS_CN));
500: ts->data = (void*)cn;
501: return(0);
502: }