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 <petscdmplex.h> 7 8 typedef struct { 9 Vec Xdot, update; 10 PetscReal stage_time; 11 PetscInt version; 12 } TS_Mimex; 13 14 static PetscErrorCode TSMimexGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot) 15 { 16 TS_Mimex *mimex = (TS_Mimex *) ts->data; 17 PetscErrorCode ierr; 18 19 PetscFunctionBegin; 20 if (X0) { 21 if (dm && dm != ts->dm) {ierr = DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);CHKERRQ(ierr);} 22 else {*X0 = ts->vec_sol;} 23 } 24 if (Xdot) { 25 if (dm && dm != ts->dm) {ierr = DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);CHKERRQ(ierr);} 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 PetscErrorCode ierr; 34 35 PetscFunctionBegin; 36 if (X0) if (dm && dm != ts->dm) {ierr = DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);CHKERRQ(ierr);} 37 if (Xdot) if (dm && dm != ts->dm) {ierr = DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);CHKERRQ(ierr);} 38 PetscFunctionReturn(0); 39 } 40 41 static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G) 42 { 43 PetscErrorCode ierr; 44 45 PetscFunctionBegin; 46 ierr = DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);CHKERRQ(ierr); 47 ierr = DMGetNamedGlobalVector(dm, "TSMimex_G", G);CHKERRQ(ierr); 48 PetscFunctionReturn(0); 49 } 50 51 static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G) 52 { 53 PetscErrorCode ierr; 54 55 PetscFunctionBegin; 56 ierr = DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);CHKERRQ(ierr); 57 ierr = DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);CHKERRQ(ierr); 58 PetscFunctionReturn(0); 59 } 60 61 /* 62 This defines the nonlinear equation that is to be solved with SNES 63 G(U) = F[t0+dt, U, (U-U0)*shift] = 0 64 */ 65 static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts) 66 { 67 TS_Mimex *mimex = (TS_Mimex *) ts->data; 68 DM dm, dmsave; 69 Vec X0, Xdot; 70 PetscReal shift = 1./ts->time_step; 71 PetscErrorCode ierr; 72 73 PetscFunctionBegin; 74 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 75 ierr = TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr); 76 ierr = VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);CHKERRQ(ierr); 77 78 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 79 dmsave = ts->dm; 80 ts->dm = dm; 81 ierr = TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);CHKERRQ(ierr); 82 if (mimex->version == 1) { 83 DM dm; 84 PetscDS prob; 85 PetscSection s; 86 Vec Xstar = NULL, G = NULL; 87 const PetscScalar *ax; 88 PetscScalar *axstar; 89 PetscInt Nf, f, pStart, pEnd, p; 90 91 ierr = TSGetDM(ts, &dm);CHKERRQ(ierr); 92 ierr = DMGetDS(dm, &prob);CHKERRQ(ierr); 93 ierr = DMGetSection(dm, &s);CHKERRQ(ierr); 94 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 95 ierr = PetscSectionGetChart(s, &pStart, &pEnd);CHKERRQ(ierr); 96 ierr = TSMimexGetXstarAndG(ts, dm, &Xstar, &G);CHKERRQ(ierr); 97 ierr = VecCopy(X0, Xstar);CHKERRQ(ierr); 98 ierr = VecGetArrayRead(x, &ax);CHKERRQ(ierr); 99 ierr = VecGetArray(Xstar, &axstar);CHKERRQ(ierr); 100 for (f = 0; f < Nf; ++f) { 101 PetscBool implicit; 102 103 ierr = PetscDSGetImplicit(prob, f, &implicit);CHKERRQ(ierr); 104 if (!implicit) continue; 105 for (p = pStart; p < pEnd; ++p) { 106 PetscScalar *a, *axs; 107 PetscInt fdof, fcdof, d; 108 109 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 110 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 111 ierr = DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);CHKERRQ(ierr); 112 ierr = DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);CHKERRQ(ierr); 113 for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d]; 114 } 115 } 116 ierr = VecRestoreArrayRead(x, &ax);CHKERRQ(ierr); 117 ierr = VecRestoreArray(Xstar, &axstar);CHKERRQ(ierr); 118 ierr = TSComputeRHSFunction(ts, ts->ptime, Xstar, G);CHKERRQ(ierr); 119 ierr = VecAXPY(y, -1.0, G);CHKERRQ(ierr); 120 ierr = TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);CHKERRQ(ierr); 121 } 122 ts->dm = dmsave; 123 ierr = TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr); 124 PetscFunctionReturn(0); 125 } 126 127 static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts) 128 { 129 TS_Mimex *mimex = (TS_Mimex *) ts->data; 130 DM dm, dmsave; 131 Vec Xdot; 132 PetscReal shift = 1./ts->time_step; 133 PetscErrorCode ierr; 134 135 PetscFunctionBegin; 136 /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */ 137 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 138 ierr = TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr); 139 140 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 141 dmsave = ts->dm; 142 ts->dm = dm; 143 ierr = TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);CHKERRQ(ierr); 144 ts->dm = dmsave; 145 ierr = TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr); 146 PetscFunctionReturn(0); 147 } 148 149 static PetscErrorCode TSStep_Mimex_Split(TS ts) 150 { 151 TS_Mimex *mimex = (TS_Mimex *) ts->data; 152 DM dm; 153 PetscDS prob; 154 PetscSection s; 155 Vec sol = ts->vec_sol, update = mimex->update; 156 const PetscScalar *aupdate; 157 PetscScalar *asol, dt = ts->time_step; 158 PetscInt Nf, f, pStart, pEnd, p; 159 PetscErrorCode ierr; 160 161 PetscFunctionBegin; 162 ierr = TSGetDM(ts, &dm);CHKERRQ(ierr); 163 ierr = DMGetDS(dm, &prob);CHKERRQ(ierr); 164 ierr = DMGetSection(dm, &s);CHKERRQ(ierr); 165 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 166 ierr = PetscSectionGetChart(s, &pStart, &pEnd);CHKERRQ(ierr); 167 ierr = TSPreStage(ts, ts->ptime);CHKERRQ(ierr); 168 /* Compute implicit update */ 169 mimex->stage_time = ts->ptime + ts->time_step; 170 ierr = VecCopy(sol, update);CHKERRQ(ierr); 171 ierr = SNESSolve(ts->snes, NULL, update);CHKERRQ(ierr); 172 ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr); 173 ierr = VecGetArray(sol, &asol);CHKERRQ(ierr); 174 for (f = 0; f < Nf; ++f) { 175 PetscBool implicit; 176 177 ierr = PetscDSGetImplicit(prob, f, &implicit);CHKERRQ(ierr); 178 if (!implicit) continue; 179 for (p = pStart; p < pEnd; ++p) { 180 PetscScalar *au, *as; 181 PetscInt fdof, fcdof, d; 182 183 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 184 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 185 ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr); 186 ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr); 187 for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d]; 188 } 189 } 190 ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr); 191 ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr); 192 /* Compute explicit update */ 193 ierr = TSComputeRHSFunction(ts, ts->ptime, sol, update);CHKERRQ(ierr); 194 ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr); 195 ierr = VecGetArray(sol, &asol);CHKERRQ(ierr); 196 for (f = 0; f < Nf; ++f) { 197 PetscBool implicit; 198 199 ierr = PetscDSGetImplicit(prob, f, &implicit);CHKERRQ(ierr); 200 if (implicit) continue; 201 for (p = pStart; p < pEnd; ++p) { 202 PetscScalar *au, *as; 203 PetscInt fdof, fcdof, d; 204 205 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 206 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 207 ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr); 208 ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr); 209 for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d]; 210 } 211 } 212 ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr); 213 ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr); 214 ierr = TSPostStage(ts, ts->ptime, 0, &sol);CHKERRQ(ierr); 215 ts->ptime += ts->time_step; 216 PetscFunctionReturn(0); 217 } 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