xref: /petsc/src/ts/impls/implicit/theta/theta.c (revision 0f5c6efe1f3f3c860ea310cfc5c66eb104cb8809)
1 #define PETSCTS_DLL
2 
3 /*
4   Code for timestepping with implicit Theta method
5 
6   Notes:
7   This method can be applied to DAE.
8 
9   This method is cast as a 1-stage implicit Runge-Kutta method.
10 
11   Theta | Theta
12   -------------
13         |  1
14 
15   To apply a diagonally implicit RK method to DAE, the stage formula
16 
17   X_i = x + h sum_j a_ij X'_j
18 
19   is interpreted as a formula for X'_i in terms of X_i and known stuff (X'_j, j<i)
20 */
21 #include "private/tsimpl.h"                /*I   "petscts.h"   I*/
22 
23 typedef struct {
24   Vec X,Xdot;                   /* Storage for one stage */
25   Vec res;                      /* DAE residuals */
26   PetscTruth extrapolate;
27   PetscReal Theta;
28   PetscReal shift;
29   PetscReal stage_time;
30 } TS_Theta;
31 
32 #undef __FUNCT__
33 #define __FUNCT__ "TSStep_Theta"
34 static PetscErrorCode TSStep_Theta(TS ts,PetscInt *steps,PetscReal *ptime)
35 {
36   TS_Theta       *th = (TS_Theta*)ts->data;
37   PetscInt       i,max_steps = ts->max_steps,its,lits;
38   PetscErrorCode ierr;
39 
40   PetscFunctionBegin;
41   *steps = -ts->steps;
42   *ptime = ts->ptime;
43 
44   ierr = TSMonitor(ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
45 
46   for (i=0; i<max_steps; i++) {
47     if (ts->ptime + ts->time_step > ts->max_time) break;
48     ierr = TSPreStep(ts);CHKERRQ(ierr);
49 
50     th->stage_time = ts->ptime + th->Theta*ts->time_step;
51     th->shift = 1./(th->Theta*ts->time_step);
52 
53     if (th->extrapolate) {
54       ierr = VecWAXPY(th->X,1./th->shift,th->Xdot,ts->vec_sol);CHKERRQ(ierr);
55     } else {
56       ierr = VecCopy(ts->vec_sol,th->X);CHKERRQ(ierr);
57     }
58     ierr = SNESSolve(ts->snes,PETSC_NULL,th->X);CHKERRQ(ierr);
59     ierr = SNESGetIterationNumber(ts->snes,&its);CHKERRQ(ierr);
60     ierr = SNESGetLinearSolveIterations(ts->snes,&lits);CHKERRQ(ierr);
61     ts->nonlinear_its += its; ts->linear_its += lits;
62 
63     ierr = VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,ts->vec_sol,th->X);CHKERRQ(ierr);
64     ierr = VecAXPY(ts->vec_sol,ts->time_step,th->Xdot);CHKERRQ(ierr);
65     ts->ptime += ts->time_step;
66     ts->steps++;
67 
68     ierr = TSPostStep(ts);CHKERRQ(ierr);
69     ierr = TSMonitor(ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
70   }
71 
72   *steps += ts->steps;
73   *ptime  = ts->ptime;
74   PetscFunctionReturn(0);
75 }
76 
77 /*------------------------------------------------------------*/
78 #undef __FUNCT__
79 #define __FUNCT__ "TSDestroy_Theta"
80 static PetscErrorCode TSDestroy_Theta(TS ts)
81 {
82   TS_Theta       *th = (TS_Theta*)ts->data;
83   PetscErrorCode  ierr;
84 
85   PetscFunctionBegin;
86   if (th->X)    {ierr = VecDestroy(th->X);CHKERRQ(ierr);}
87   if (th->Xdot) {ierr = VecDestroy(th->Xdot);CHKERRQ(ierr);}
88   if (th->res)  {ierr = VecDestroy(th->res);CHKERRQ(ierr);}
89   ierr = PetscFree(th);CHKERRQ(ierr);
90   PetscFunctionReturn(0);
91 }
92 
93 /*
94   This defines the nonlinear equation that is to be solved with SNES
95   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
96 */
97 #undef __FUNCT__
98 #define __FUNCT__ "SNESTSFormFunction_Theta"
99 static PetscErrorCode SNESTSFormFunction_Theta(SNES snes,Vec x,Vec y,TS ts)
100 {
101   TS_Theta *th = (TS_Theta*)ts->data;
102   PetscErrorCode ierr;
103 
104   PetscFunctionBegin;
105   ierr = VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,ts->vec_sol,x);CHKERRQ(ierr);
106   ierr = TSComputeIFunction(ts,th->stage_time,x,th->Xdot,y);CHKERRQ(ierr);
107   PetscFunctionReturn(0);
108 }
109 
110 #undef __FUNCT__
111 #define __FUNCT__ "SNESTSFormJacobian_Theta"
112 static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat *A,Mat *B,MatStructure *str,TS ts)
113 {
114   TS_Theta *th = (TS_Theta*)ts->data;
115   PetscErrorCode ierr;
116 
117   PetscFunctionBegin;
118   /* th->Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
119   ierr = TSComputeIJacobian(ts,th->stage_time,x,th->Xdot,th->shift,A,B,str);CHKERRQ(ierr);
120   PetscFunctionReturn(0);
121 }
122 
123 
124 #undef __FUNCT__
125 #define __FUNCT__ "TSSetUp_Theta"
126 static PetscErrorCode TSSetUp_Theta(TS ts)
127 {
128   TS_Theta *th = (TS_Theta*)ts->data;
129   PetscErrorCode ierr;
130 
131   PetscFunctionBegin;
132   if (ts->problem_type == TS_LINEAR) {
133     SETERRQ(PETSC_ERR_ARG_WRONG,"Only for nonlinear problems");
134   }
135   ierr = VecDuplicate(ts->vec_sol,&th->X);CHKERRQ(ierr);
136   ierr = VecDuplicate(ts->vec_sol,&th->Xdot);CHKERRQ(ierr);
137   ierr = VecDuplicate(ts->vec_sol,&th->res);CHKERRQ(ierr);
138   ierr = SNESSetFunction(ts->snes,th->res,SNESTSFormFunction,ts);CHKERRQ(ierr);
139   /* This is nasty.  SNESSetFromOptions() is usually called in TSSetFromOptions().  With -snes_mf_operator, it will
140   replace A and we don't want to mess with that.  With -snes_mf, A and B will be replaced as well as the function and
141   context.  Note that SNESSetFunction() normally has not been called before SNESSetFromOptions(), so when -snes_mf sets
142   the Jacobian user context to snes->funP, it will actually be NULL.  This is not a problem because both snes->funP and
143   snes->jacP should be the TS. */
144   {
145     Mat A,B;
146     PetscErrorCode (*func)(SNES,Vec,Mat*,Mat*,MatStructure*,void*);
147     void *ctx;
148     ierr = SNESGetJacobian(ts->snes,&A,&B,&func,&ctx);CHKERRQ(ierr);
149     ierr = SNESSetJacobian(ts->snes,A?A:ts->A,B?B:ts->B,func?func:&SNESTSFormJacobian,ctx?ctx:ts);CHKERRQ(ierr);
150   }
151   PetscFunctionReturn(0);
152 }
153 /*------------------------------------------------------------*/
154 
155 #undef __FUNCT__
156 #define __FUNCT__ "TSSetFromOptions_Theta"
157 static PetscErrorCode TSSetFromOptions_Theta(TS ts)
158 {
159   TS_Theta *th = (TS_Theta*)ts->data;
160   PetscErrorCode ierr;
161 
162   PetscFunctionBegin;
163   ierr = PetscOptionsHead("Theta ODE solver options");CHKERRQ(ierr);
164   {
165     ierr = PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,PETSC_NULL);CHKERRQ(ierr);
166     ierr = PetscOptionsTruth("-ts_theta_extrapolate","Extrapolate stage solution from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,PETSC_NULL);CHKERRQ(ierr);
167   }
168   ierr = PetscOptionsTail();CHKERRQ(ierr);
169   PetscFunctionReturn(0);
170 }
171 
172 #undef __FUNCT__
173 #define __FUNCT__ "TSView_Theta"
174 static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
175 {
176   TS_Theta       *th = (TS_Theta*)ts->data;
177   PetscTruth      iascii;
178   PetscErrorCode  ierr;
179 
180   PetscFunctionBegin;
181   ierr = PetscTypeCompare((PetscObject)viewer,PETSC_VIEWER_ASCII,&iascii);CHKERRQ(ierr);
182   if (iascii) {
183     ierr = PetscViewerASCIIPrintf(viewer,"  Theta=%G\n",th->Theta);CHKERRQ(ierr);
184     ierr = PetscViewerASCIIPrintf(viewer,"  Extrapolation=%s\n",th->extrapolate?"yes":"no");CHKERRQ(ierr);
185   } else {
186     SETERRQ1(PETSC_ERR_SUP,"Viewer type %s not supported for TS_Theta",((PetscObject)viewer)->type_name);
187   }
188   PetscFunctionReturn(0);
189 }
190 
191 /* ------------------------------------------------------------ */
192 /*MC
193       TSTHETA - DAE solver using the implicit Theta method
194 
195   Level: beginner
196 
197 .seealso:  TSCreate(), TS, TSSetType()
198 
199 M*/
200 EXTERN_C_BEGIN
201 #undef __FUNCT__
202 #define __FUNCT__ "TSCreate_Theta"
203 PetscErrorCode PETSCTS_DLLEXPORT TSCreate_Theta(TS ts)
204 {
205   TS_Theta       *th;
206   PetscErrorCode ierr;
207 
208   PetscFunctionBegin;
209   ts->ops->destroy        = TSDestroy_Theta;
210   ts->ops->view           = TSView_Theta;
211   ts->ops->setup          = TSSetUp_Theta;
212   ts->ops->step           = TSStep_Theta;
213   ts->ops->setfromoptions = TSSetFromOptions_Theta;
214   ts->ops->snesfunction   = SNESTSFormFunction_Theta;
215   ts->ops->snesjacobian   = SNESTSFormJacobian_Theta;
216 
217   ts->problem_type = TS_NONLINEAR;
218   ierr = SNESCreate(((PetscObject)ts)->comm,&ts->snes);CHKERRQ(ierr);
219   ierr = PetscObjectIncrementTabLevel((PetscObject)ts->snes,(PetscObject)ts,1);CHKERRQ(ierr);
220 
221   ierr = PetscNewLog(ts,TS_Theta,&th);CHKERRQ(ierr);
222   ts->data = (void*)th;
223 
224   th->extrapolate = PETSC_TRUE;
225   th->Theta       = 0.5;
226 
227   PetscFunctionReturn(0);
228 }
229 EXTERN_C_END
230