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: if (ts->ptime + ts->time_step > ts->max_time) break;
130: /* set rhs = (1/dt*Alhs + 0.5*Arhs)*sol */
131: MatMult(ts->Arhs,sol,rhs); /* rhs = 0.5*Arhs*sol */
132: if (ts->Alhs){
133: MatMultAdd(ts->Alhs,sol,rhs,rhs); /* rhs = rhs + 1/dt*Alhs*sol */
134: } else {
135: VecAXPY(rhs,mdt,sol); /* rhs = rhs + 1/dt*sol */
136: }
138: ts->ptime += ts->time_step;
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: if (ts->ptime + ts->time_step > ts->max_time) break;
176: /* set rhs = (1/dt*Alhs(t_mid) + 0.5*Arhs(t_n)) * sol */
177: if (i==0){
178: /* evaluate 0.5*Arhs(t_0) */
179: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->Arhs,PETSC_NULL,&str,ts->jacP);
180: MatScale(ts->Arhs,0.5);
181: }
182: if (ts->Alhs){
183: /* evaluate Alhs(t_mid) */
184: t_mid = ts->ptime+ts->time_step/2.0;
185: (*ts->ops->lhsmatrix)(ts,t_mid,&ts->Alhs,PETSC_NULL,&str,ts->jacPlhs);
186: MatMult(ts->Alhs,sol,rhs); /* rhs = Alhs_mid*sol */
187: VecScale(rhs,mdt); /* rhs = 1/dt*Alhs_mid*sol */
188: MatMultAdd(ts->Arhs,sol,rhs,rhs); /* rhs = rhs + 0.5*Arhs_mid*sol */
189: } else {
190: MatMult(ts->Arhs,sol,rhs); /* rhs = 0.5*Arhs_n*sol */
191: VecAXPY(rhs,mdt,sol); /* rhs = rhs + 1/dt*sol */
192: }
194: ts->ptime += ts->time_step;
196: /* evaluate Arhs at current ptime t_{n+1} */
197: (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->Arhs,PETSC_NULL,&str,ts->jacP);
198: TSSetKSPOperators_CN_Matrix(ts);
200: KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
201: KSPSolve(ts->ksp,rhs,update);
202: KSPGetIterationNumber(ts->ksp,&its);
203: ts->linear_its += PetscAbsInt(its);
204: VecCopy(update,sol);
205: ts->steps++;
206: TSMonitor(ts,ts->steps,ts->ptime,sol);
207: }
209: *steps += ts->steps;
210: *ptime = ts->ptime;
211: return(0);
212: }
213: /*
214: Version for nonlinear PDE.
215: */
218: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
219: {
220: Vec sol = ts->vec_sol;
222: PetscInt i,max_steps = ts->max_steps,its,lits;
223: TS_CN *cn = (TS_CN*)ts->data;
224:
226: *steps = -ts->steps;
227: TSMonitor(ts,ts->steps,ts->ptime,sol);
229: for (i=0; i<max_steps; i++) {
230: if (ts->ptime + ts->time_step > ts->max_time) break;
231: ts->ptime += ts->time_step;
232:
233: VecCopy(sol,cn->update);
234: SNESSolve(ts->snes,PETSC_NULL,cn->update);
235: SNESGetIterationNumber(ts->snes,&its);
236: SNESGetLinearSolveIterations(ts->snes,&lits);
237: ts->nonlinear_its += its; ts->linear_its += lits;
238: VecCopy(cn->update,sol);
239: ts->steps++;
240: TSMonitor(ts,ts->steps,ts->ptime,sol);
241: }
243: *steps += ts->steps;
244: *ptime = ts->ptime;
245: return(0);
246: }
248: /*------------------------------------------------------------*/
251: static PetscErrorCode TSDestroy_CN(TS ts)
252: {
253: TS_CN *cn = (TS_CN*)ts->data;
257: if (cn->update) {VecDestroy(cn->update);}
258: if (cn->func) {VecDestroy(cn->func);}
259: if (cn->rhsfunc) {VecDestroy(cn->rhsfunc);}
260: if (cn->rhsfunc_old) {VecDestroy(cn->rhsfunc_old);}
261: if (cn->rhs) {VecDestroy(cn->rhs);}
262: PetscFree(cn);
263: return(0);
264: }
266: /*
267: This defines the nonlinear equation that is to be solved with SNES
268: 1/dt*Alhs*(U^{n+1} - U^{n}) - 0.5*(F(U^{n+1}) + F(U^{n}))
269: */
272: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
273: {
274: TS ts = (TS) ctx;
275: PetscScalar mdt = 1.0/ts->time_step,*unp1,*un,*Funp1,*Fun,*yarray;
277: PetscInt i,n;
278: TS_CN *cn = (TS_CN*)ts->data;
281: /* apply user provided function */
282: if (cn->rhsfunc_time == (ts->ptime - ts->time_step)){
283: /* printf(" copy rhsfunc to rhsfunc_old, then eval rhsfunc\n"); */
284: VecCopy(cn->rhsfunc,cn->rhsfunc_old);
285: cn->rhsfunc_old_time = cn->rhsfunc_time;
286: } else if (cn->rhsfunc_time != ts->ptime && cn->rhsfunc_old_time != ts->ptime-ts->time_step){
287: /* printf(" eval both rhsfunc_old and rhsfunc\n"); */
288: TSComputeRHSFunction(ts,ts->ptime-ts->time_step,ts->vec_sol,cn->rhsfunc_old); /* rhsfunc_old=F(U^{n}) */
289: cn->rhsfunc_old_time = ts->ptime - ts->time_step;
290: }
291:
292: if (ts->Alhs){
293: /* compute y=Alhs*(U^{n+1} - U^{n}) with cn->rhsfunc as workspce */
294: VecWAXPY(cn->rhsfunc,-1.0,ts->vec_sol,x);
295: MatMult(ts->Alhs,cn->rhsfunc,y);
296: }
298: TSComputeRHSFunction(ts,ts->ptime,x,cn->rhsfunc); /* rhsfunc = F(U^{n+1}) */
299: cn->rhsfunc_time = ts->ptime;
300:
301: VecGetArray(ts->vec_sol,&un); /* U^{n} */
302: VecGetArray(x,&unp1); /* U^{n+1} */
303: VecGetArray(cn->rhsfunc,&Funp1);
304: VecGetArray(cn->rhsfunc_old,&Fun);
305: VecGetArray(y,&yarray);
306: VecGetLocalSize(x,&n);
307: if (ts->Alhs){
308: for (i=0; i<n; i++) {
309: yarray[i] = mdt*yarray[i] - 0.5*(Funp1[i] + Fun[i]);
310: }
311: } else {
312: for (i=0; i<n; i++) {
313: yarray[i] = mdt*(unp1[i] - un[i]) - 0.5*(Funp1[i] + Fun[i]);
314: }
315: }
316: VecRestoreArray(ts->vec_sol,&un);
317: VecRestoreArray(x,&unp1);
318: VecRestoreArray(cn->rhsfunc,&Funp1);
319: VecRestoreArray(cn->rhsfunc_old,&Fun);
320: VecRestoreArray(y,&yarray);
321: return(0);
322: }
324: /* Set A = B = 1/dt*A - 0.5*A */
327: PetscErrorCode TSScaleShiftMatrices_CN(TS ts,Mat A,Mat B,MatStructure str)
328: {
329: PetscTruth flg;
331: PetscScalar mdt = 1.0/ts->time_step;
334: /* this function requires additional work! */
335: PetscTypeCompare((PetscObject)A,MATMFFD,&flg);
336: if (!flg) {
337: MatScale(A,-0.5);
338: if (ts->Alhs){
339: MatAXPY(A,mdt,ts->Alhs,DIFFERENT_NONZERO_PATTERN); /* DIFFERENT_NONZERO_PATTERN? */
340: } else {
341: MatShift(A,mdt);
342: }
343: } else {
344: SETERRQ(PETSC_ERR_SUP,"Matrix type MATMFFD is not supported yet"); /* ref TSScaleShiftMatrices() */
345: }
346: if (B != A && str != SAME_PRECONDITIONER) {
347: SETERRQ(PETSC_ERR_SUP,"not supported yet");
348: }
349: return(0);
350: }
352: /*
353: This constructs the Jacobian needed for SNES
354: J = I/dt - 0.5*J_{F} where J_{F} is the given Jacobian of F.
355: x - input vector
356: AA - Jacobian matrix
357: BB - preconditioner matrix, usually the same as AA
358: */
361: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
362: {
363: TS ts = (TS) ctx;
367: /* construct user's Jacobian */
368: TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str); /* AA = J_{F} */
370: /* shift and scale Jacobian */
371: TSScaleShiftMatrices_CN(ts,*AA,*BB,*str); /* Set AA = 1/dt*Alhs - 0.5*AA */
372: return(0);
373: }
375: /* ------------------------------------------------------------*/
378: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
379: {
380: TS_CN *cn = (TS_CN*)ts->data;
382: PetscTruth shelltype;
385: VecDuplicate(ts->vec_sol,&cn->update);
386: VecDuplicate(ts->vec_sol,&cn->rhs);
387:
388: /* build linear system to be solved */
389: /* scale ts->Alhs = 1/dt*Alhs, ts->Arhs = 0.5*Arhs; set ts->A = Alhs - Arhs */
390: PetscTypeCompare((PetscObject)ts->Arhs,MATSHELL,&shelltype);
391: if (shelltype){
392: TSSetKSPOperators_CN_No_Matrix(ts);
393: } else {
394: TSSetKSPOperators_CN_Matrix(ts);
395: }
396: KSPSetOperators(ts->ksp,ts->A,ts->A,SAME_NONZERO_PATTERN);
397: return(0);
398: }
402: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
403: {
404: TS_CN *cn = (TS_CN*)ts->data;
408: VecDuplicate(ts->vec_sol,&cn->update);
409: VecDuplicate(ts->vec_sol,&cn->rhs);
410: return(0);
411: }
415: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
416: {
417: TS_CN *cn = (TS_CN*)ts->data;
421: VecDuplicate(ts->vec_sol,&cn->update);
422: VecDuplicate(ts->vec_sol,&cn->func);
423: VecDuplicate(ts->vec_sol,&cn->rhsfunc);
424: VecDuplicate(ts->vec_sol,&cn->rhsfunc_old);
425: SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
426: SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
427: cn->rhsfunc_time = -100.0; /* cn->rhsfunc is not evaluated yet */
428: cn->rhsfunc_old_time = -100.0;
429: return(0);
430: }
431: /*------------------------------------------------------------*/
435: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
436: {
438: return(0);
439: }
443: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
444: {
446: return(0);
447: }
451: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
452: {
454: return(0);
455: }
457: /* ------------------------------------------------------------ */
458: /*MC
459: TS_CN - ODE solver using the implicit Crank-Nicholson method
461: Level: beginner
463: .seealso: TSCreate(), TS, TSSetType()
465: M*/
469: PetscErrorCode TSCreate_CN(TS ts)
470: {
471: TS_CN *cn;
475: ts->ops->destroy = TSDestroy_CN;
476: ts->ops->view = TSView_CN;
478: if (ts->problem_type == TS_LINEAR) {
479: if (!ts->Arhs) {
480: SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
481: }
482: if (!ts->ops->rhsmatrix) {
483: ts->ops->setup = TSSetUp_CN_Linear_Constant_Matrix;
484: ts->ops->step = TSStep_CN_Linear_Constant_Matrix;
485: } else {
486: ts->ops->setup = TSSetUp_CN_Linear_Variable_Matrix;
487: ts->ops->step = TSStep_CN_Linear_Variable_Matrix;
488: }
489: ts->ops->setfromoptions = TSSetFromOptions_CN_Linear;
490: KSPCreate(ts->comm,&ts->ksp);
491: KSPSetInitialGuessNonzero(ts->ksp,PETSC_TRUE);
492: } else if (ts->problem_type == TS_NONLINEAR) {
493: ts->ops->setup = TSSetUp_CN_Nonlinear;
494: ts->ops->step = TSStep_CN_Nonlinear;
495: ts->ops->setfromoptions = TSSetFromOptions_CN_Nonlinear;
496: SNESCreate(ts->comm,&ts->snes);
497: } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");
499: PetscNew(TS_CN,&cn);
500: PetscLogObjectMemory(ts,sizeof(TS_CN));
501: ts->data = (void*)cn;
502: return(0);
503: }