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