1 /* 2 Code for Timestepping with my makeshift IMEX. 3 */ 4 #include <petsc/private/tsimpl.h> /*I "petscts.h" I*/ 5 #include <petscds.h> 6 #include <petscsection.h> 7 #include <petscdmplex.h> 8 9 typedef struct { 10 Vec Xdot, update; 11 PetscReal stage_time; 12 PetscInt version; 13 } TS_Mimex; 14 15 static PetscErrorCode TSMimexGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot) 16 { 17 TS_Mimex *mimex = (TS_Mimex *) ts->data; 18 19 PetscFunctionBegin; 20 if (X0) { 21 if (dm && dm != ts->dm) PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_X0", X0)); 22 else {*X0 = ts->vec_sol;} 23 } 24 if (Xdot) { 25 if (dm && dm != ts->dm) PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot)); 26 else {*Xdot = mimex->Xdot;} 27 } 28 PetscFunctionReturn(0); 29 } 30 31 static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot) 32 { 33 PetscFunctionBegin; 34 if (X0) if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0)); 35 if (Xdot) if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot)); 36 PetscFunctionReturn(0); 37 } 38 39 static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G) 40 { 41 PetscFunctionBegin; 42 PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar)); 43 PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_G", G)); 44 PetscFunctionReturn(0); 45 } 46 47 static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G) 48 { 49 PetscFunctionBegin; 50 PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar)); 51 PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_G", G)); 52 PetscFunctionReturn(0); 53 } 54 55 /* 56 This defines the nonlinear equation that is to be solved with SNES 57 G(U) = F[t0+dt, U, (U-U0)*shift] = 0 58 */ 59 static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts) 60 { 61 TS_Mimex *mimex = (TS_Mimex *) ts->data; 62 DM dm, dmsave; 63 Vec X0, Xdot; 64 PetscReal shift = 1./ts->time_step; 65 66 PetscFunctionBegin; 67 PetscCall(SNESGetDM(snes, &dm)); 68 PetscCall(TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot)); 69 PetscCall(VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x)); 70 71 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 72 dmsave = ts->dm; 73 ts->dm = dm; 74 PetscCall(TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE)); 75 if (mimex->version == 1) { 76 DM dm; 77 PetscDS prob; 78 PetscSection s; 79 Vec Xstar = NULL, G = NULL; 80 const PetscScalar *ax; 81 PetscScalar *axstar; 82 PetscInt Nf, f, pStart, pEnd, p; 83 84 PetscCall(TSGetDM(ts, &dm)); 85 PetscCall(DMGetDS(dm, &prob)); 86 PetscCall(DMGetLocalSection(dm, &s)); 87 PetscCall(PetscDSGetNumFields(prob, &Nf)); 88 PetscCall(PetscSectionGetChart(s, &pStart, &pEnd)); 89 PetscCall(TSMimexGetXstarAndG(ts, dm, &Xstar, &G)); 90 PetscCall(VecCopy(X0, Xstar)); 91 PetscCall(VecGetArrayRead(x, &ax)); 92 PetscCall(VecGetArray(Xstar, &axstar)); 93 for (f = 0; f < Nf; ++f) { 94 PetscBool implicit; 95 96 PetscCall(PetscDSGetImplicit(prob, f, &implicit)); 97 if (!implicit) continue; 98 for (p = pStart; p < pEnd; ++p) { 99 PetscScalar *a, *axs; 100 PetscInt fdof, fcdof, d; 101 102 PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof)); 103 PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof)); 104 PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, ax, &a)); 105 PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs)); 106 for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d]; 107 } 108 } 109 PetscCall(VecRestoreArrayRead(x, &ax)); 110 PetscCall(VecRestoreArray(Xstar, &axstar)); 111 PetscCall(TSComputeRHSFunction(ts, ts->ptime, Xstar, G)); 112 PetscCall(VecAXPY(y, -1.0, G)); 113 PetscCall(TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G)); 114 } 115 ts->dm = dmsave; 116 PetscCall(TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot)); 117 PetscFunctionReturn(0); 118 } 119 120 static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts) 121 { 122 TS_Mimex *mimex = (TS_Mimex *) ts->data; 123 DM dm, dmsave; 124 Vec Xdot; 125 PetscReal shift = 1./ts->time_step; 126 127 PetscFunctionBegin; 128 /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */ 129 PetscCall(SNESGetDM(snes, &dm)); 130 PetscCall(TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot)); 131 132 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 133 dmsave = ts->dm; 134 ts->dm = dm; 135 PetscCall(TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE)); 136 ts->dm = dmsave; 137 PetscCall(TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot)); 138 PetscFunctionReturn(0); 139 } 140 141 static PetscErrorCode TSStep_Mimex_Split(TS ts) 142 { 143 TS_Mimex *mimex = (TS_Mimex *) ts->data; 144 DM dm; 145 PetscDS prob; 146 PetscSection s; 147 Vec sol = ts->vec_sol, update = mimex->update; 148 const PetscScalar *aupdate; 149 PetscScalar *asol, dt = ts->time_step; 150 PetscInt Nf, f, pStart, pEnd, p; 151 152 PetscFunctionBegin; 153 PetscCall(TSGetDM(ts, &dm)); 154 PetscCall(DMGetDS(dm, &prob)); 155 PetscCall(DMGetLocalSection(dm, &s)); 156 PetscCall(PetscDSGetNumFields(prob, &Nf)); 157 PetscCall(PetscSectionGetChart(s, &pStart, &pEnd)); 158 PetscCall(TSPreStage(ts, ts->ptime)); 159 /* Compute implicit update */ 160 mimex->stage_time = ts->ptime + ts->time_step; 161 PetscCall(VecCopy(sol, update)); 162 PetscCall(SNESSolve(ts->snes, NULL, update)); 163 PetscCall(VecGetArrayRead(update, &aupdate)); 164 PetscCall(VecGetArray(sol, &asol)); 165 for (f = 0; f < Nf; ++f) { 166 PetscBool implicit; 167 168 PetscCall(PetscDSGetImplicit(prob, f, &implicit)); 169 if (!implicit) continue; 170 for (p = pStart; p < pEnd; ++p) { 171 PetscScalar *au, *as; 172 PetscInt fdof, fcdof, d; 173 174 PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof)); 175 PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof)); 176 PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au)); 177 PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, asol, &as)); 178 for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d]; 179 } 180 } 181 PetscCall(VecRestoreArrayRead(update, &aupdate)); 182 PetscCall(VecRestoreArray(sol, &asol)); 183 /* Compute explicit update */ 184 PetscCall(TSComputeRHSFunction(ts, ts->ptime, sol, update)); 185 PetscCall(VecGetArrayRead(update, &aupdate)); 186 PetscCall(VecGetArray(sol, &asol)); 187 for (f = 0; f < Nf; ++f) { 188 PetscBool implicit; 189 190 PetscCall(PetscDSGetImplicit(prob, f, &implicit)); 191 if (implicit) continue; 192 for (p = pStart; p < pEnd; ++p) { 193 PetscScalar *au, *as; 194 PetscInt fdof, fcdof, d; 195 196 PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof)); 197 PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof)); 198 PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au)); 199 PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, asol, &as)); 200 for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d]; 201 } 202 } 203 PetscCall(VecRestoreArrayRead(update, &aupdate)); 204 PetscCall(VecRestoreArray(sol, &asol)); 205 PetscCall(TSPostStage(ts, ts->ptime, 0, &sol)); 206 ts->ptime += ts->time_step; 207 PetscFunctionReturn(0); 208 } 209 210 /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */ 211 static PetscErrorCode TSStep_Mimex_Implicit(TS ts) 212 { 213 TS_Mimex *mimex = (TS_Mimex *) ts->data; 214 Vec sol = ts->vec_sol; 215 Vec update = mimex->update; 216 217 PetscFunctionBegin; 218 PetscCall(TSPreStage(ts, ts->ptime)); 219 /* Compute implicit update */ 220 mimex->stage_time = ts->ptime + ts->time_step; 221 ts->ptime += ts->time_step; 222 PetscCall(VecCopy(sol, update)); 223 PetscCall(SNESSolve(ts->snes, NULL, update)); 224 PetscCall(VecCopy(update, sol)); 225 PetscCall(TSPostStage(ts, ts->ptime, 0, &sol)); 226 PetscFunctionReturn(0); 227 } 228 229 static PetscErrorCode TSStep_Mimex(TS ts) 230 { 231 TS_Mimex *mimex = (TS_Mimex*)ts->data; 232 233 PetscFunctionBegin; 234 switch(mimex->version) { 235 case 0: 236 PetscCall(TSStep_Mimex_Split(ts)); break; 237 case 1: 238 PetscCall(TSStep_Mimex_Implicit(ts)); break; 239 default: 240 SETERRQ(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %" PetscInt_FMT, mimex->version); 241 } 242 PetscFunctionReturn(0); 243 } 244 245 /*------------------------------------------------------------*/ 246 247 static PetscErrorCode TSSetUp_Mimex(TS ts) 248 { 249 TS_Mimex *mimex = (TS_Mimex*)ts->data; 250 251 PetscFunctionBegin; 252 PetscCall(VecDuplicate(ts->vec_sol, &mimex->update)); 253 PetscCall(VecDuplicate(ts->vec_sol, &mimex->Xdot)); 254 PetscFunctionReturn(0); 255 } 256 257 static PetscErrorCode TSReset_Mimex(TS ts) 258 { 259 TS_Mimex *mimex = (TS_Mimex*)ts->data; 260 261 PetscFunctionBegin; 262 PetscCall(VecDestroy(&mimex->update)); 263 PetscCall(VecDestroy(&mimex->Xdot)); 264 PetscFunctionReturn(0); 265 } 266 267 static PetscErrorCode TSDestroy_Mimex(TS ts) 268 { 269 PetscFunctionBegin; 270 PetscCall(TSReset_Mimex(ts)); 271 PetscCall(PetscFree(ts->data)); 272 PetscFunctionReturn(0); 273 } 274 /*------------------------------------------------------------*/ 275 276 static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts) 277 { 278 TS_Mimex *mimex = (TS_Mimex *) ts->data; 279 280 PetscFunctionBegin; 281 PetscOptionsHeadBegin(PetscOptionsObject, "MIMEX ODE solver options"); 282 { 283 PetscCall(PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL)); 284 } 285 PetscOptionsHeadEnd(); 286 PetscFunctionReturn(0); 287 } 288 289 static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer) 290 { 291 TS_Mimex *mimex = (TS_Mimex *) ts->data; 292 PetscBool iascii; 293 294 PetscFunctionBegin; 295 PetscCall(PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii)); 296 if (iascii) { 297 PetscCall(PetscViewerASCIIPrintf(viewer, " Version = %" PetscInt_FMT "\n", mimex->version)); 298 } 299 PetscFunctionReturn(0); 300 } 301 302 static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X) 303 { 304 PetscReal alpha = (ts->ptime - t)/ts->time_step; 305 306 PetscFunctionBegin; 307 PetscCall(VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X)); 308 PetscFunctionReturn(0); 309 } 310 311 static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi) 312 { 313 PetscFunctionBegin; 314 *yr = 1.0 + xr; 315 *yi = xi; 316 PetscFunctionReturn(0); 317 } 318 /* ------------------------------------------------------------ */ 319 320 /*MC 321 TSMIMEX - ODE solver using the explicit forward Mimex method 322 323 Level: beginner 324 325 .seealso: TSCreate(), TS, TSSetType(), TSBEULER 326 327 M*/ 328 PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts) 329 { 330 TS_Mimex *mimex; 331 332 PetscFunctionBegin; 333 ts->ops->setup = TSSetUp_Mimex; 334 ts->ops->step = TSStep_Mimex; 335 ts->ops->reset = TSReset_Mimex; 336 ts->ops->destroy = TSDestroy_Mimex; 337 ts->ops->setfromoptions = TSSetFromOptions_Mimex; 338 ts->ops->view = TSView_Mimex; 339 ts->ops->interpolate = TSInterpolate_Mimex; 340 ts->ops->linearstability = TSComputeLinearStability_Mimex; 341 ts->ops->snesfunction = SNESTSFormFunction_Mimex; 342 ts->ops->snesjacobian = SNESTSFormJacobian_Mimex; 343 ts->default_adapt_type = TSADAPTNONE; 344 345 PetscCall(PetscNewLog(ts,&mimex)); 346 ts->data = (void*)mimex; 347 348 mimex->version = 1; 349 PetscFunctionReturn(0); 350 } 351