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 ts->steps++; 231 PetscFunctionReturn(0); 232 } 233 234 235 #undef __FUNCT__ 236 #define __FUNCT__ "TSStep_Mimex_Implicit" 237 /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */ 238 static PetscErrorCode TSStep_Mimex_Implicit(TS ts) 239 { 240 TS_Mimex *mimex = (TS_Mimex *) ts->data; 241 Vec sol = ts->vec_sol; 242 Vec update = mimex->update; 243 PetscErrorCode ierr; 244 245 PetscFunctionBegin; 246 ierr = TSPreStage(ts, ts->ptime);CHKERRQ(ierr); 247 /* Compute implicit update */ 248 mimex->stage_time = ts->ptime + ts->time_step; 249 ts->ptime += ts->time_step; 250 ierr = VecCopy(sol, update);CHKERRQ(ierr); 251 ierr = SNESSolve(ts->snes, NULL, update);CHKERRQ(ierr); 252 ierr = VecCopy(update, sol);CHKERRQ(ierr); 253 ierr = TSPostStage(ts, ts->ptime, 0, &sol);CHKERRQ(ierr); 254 ts->steps++; 255 PetscFunctionReturn(0); 256 } 257 258 #undef __FUNCT__ 259 #define __FUNCT__ "TSStep_Mimex" 260 static PetscErrorCode TSStep_Mimex(TS ts) 261 { 262 TS_Mimex *mimex = (TS_Mimex*)ts->data; 263 PetscErrorCode ierr; 264 265 PetscFunctionBegin; 266 switch(mimex->version) { 267 case 0: 268 ierr = TSStep_Mimex_Split(ts);CHKERRQ(ierr); break; 269 case 1: 270 ierr = TSStep_Mimex_Implicit(ts);CHKERRQ(ierr); break; 271 default: 272 SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version); 273 } 274 PetscFunctionReturn(0); 275 } 276 277 /*------------------------------------------------------------*/ 278 279 #undef __FUNCT__ 280 #define __FUNCT__ "TSSetUp_Mimex" 281 static PetscErrorCode TSSetUp_Mimex(TS ts) 282 { 283 TS_Mimex *mimex = (TS_Mimex*)ts->data; 284 PetscErrorCode ierr; 285 286 PetscFunctionBegin; 287 ierr = VecDuplicate(ts->vec_sol, &mimex->update);CHKERRQ(ierr); 288 ierr = VecDuplicate(ts->vec_sol, &mimex->Xdot);CHKERRQ(ierr); 289 PetscFunctionReturn(0); 290 } 291 292 #undef __FUNCT__ 293 #define __FUNCT__ "TSReset_Mimex" 294 static PetscErrorCode TSReset_Mimex(TS ts) 295 { 296 TS_Mimex *mimex = (TS_Mimex*)ts->data; 297 PetscErrorCode ierr; 298 299 PetscFunctionBegin; 300 ierr = VecDestroy(&mimex->update);CHKERRQ(ierr); 301 ierr = VecDestroy(&mimex->Xdot);CHKERRQ(ierr); 302 PetscFunctionReturn(0); 303 } 304 305 #undef __FUNCT__ 306 #define __FUNCT__ "TSDestroy_Mimex" 307 static PetscErrorCode TSDestroy_Mimex(TS ts) 308 { 309 PetscErrorCode ierr; 310 311 PetscFunctionBegin; 312 ierr = TSReset_Mimex(ts);CHKERRQ(ierr); 313 ierr = PetscFree(ts->data);CHKERRQ(ierr); 314 PetscFunctionReturn(0); 315 } 316 /*------------------------------------------------------------*/ 317 318 #undef __FUNCT__ 319 #define __FUNCT__ "TSSetFromOptions_Mimex" 320 static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts) 321 { 322 TS_Mimex *mimex = (TS_Mimex *) ts->data; 323 PetscErrorCode ierr; 324 325 PetscFunctionBegin; 326 ierr = PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");CHKERRQ(ierr); 327 { 328 ierr = PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);CHKERRQ(ierr); 329 } 330 ierr = PetscOptionsTail();CHKERRQ(ierr); 331 PetscFunctionReturn(0); 332 } 333 334 #undef __FUNCT__ 335 #define __FUNCT__ "TSView_Mimex" 336 static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer) 337 { 338 TS_Mimex *mimex = (TS_Mimex *) ts->data; 339 PetscBool iascii; 340 PetscErrorCode ierr; 341 342 PetscFunctionBegin; 343 ierr = PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);CHKERRQ(ierr); 344 if (iascii) { 345 ierr = PetscViewerASCIIPrintf(viewer, " Version = %D\n", mimex->version);CHKERRQ(ierr); 346 } 347 if (ts->snes) {ierr = SNESView(ts->snes, viewer);CHKERRQ(ierr);} 348 PetscFunctionReturn(0); 349 } 350 351 #undef __FUNCT__ 352 #define __FUNCT__ "TSInterpolate_Mimex" 353 static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X) 354 { 355 PetscReal alpha = (ts->ptime - t)/ts->time_step; 356 PetscErrorCode ierr; 357 358 PetscFunctionBegin; 359 ierr = VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);CHKERRQ(ierr); 360 PetscFunctionReturn(0); 361 } 362 363 #undef __FUNCT__ 364 #define __FUNCT__ "TSComputeLinearStability_Mimex" 365 static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi) 366 { 367 PetscFunctionBegin; 368 *yr = 1.0 + xr; 369 *yi = xi; 370 PetscFunctionReturn(0); 371 } 372 /* ------------------------------------------------------------ */ 373 374 /*MC 375 TSMIMEX - ODE solver using the explicit forward Mimex method 376 377 Level: beginner 378 379 .seealso: TSCreate(), TS, TSSetType(), TSBEULER 380 381 M*/ 382 #undef __FUNCT__ 383 #define __FUNCT__ "TSCreate_Mimex" 384 PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts) 385 { 386 TS_Mimex *mimex; 387 PetscErrorCode ierr; 388 389 PetscFunctionBegin; 390 ts->ops->setup = TSSetUp_Mimex; 391 ts->ops->step = TSStep_Mimex; 392 ts->ops->reset = TSReset_Mimex; 393 ts->ops->destroy = TSDestroy_Mimex; 394 ts->ops->setfromoptions = TSSetFromOptions_Mimex; 395 ts->ops->view = TSView_Mimex; 396 ts->ops->interpolate = TSInterpolate_Mimex; 397 ts->ops->linearstability = TSComputeLinearStability_Mimex; 398 ts->ops->snesfunction = SNESTSFormFunction_Mimex; 399 ts->ops->snesjacobian = SNESTSFormJacobian_Mimex; 400 401 ierr = PetscNewLog(ts,&mimex);CHKERRQ(ierr); 402 ts->data = (void*)mimex; 403 404 mimex->version = 1; 405 PetscFunctionReturn(0); 406 } 407