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