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 = TSPreStep(ts);CHKERRQ(ierr); 182 ierr = TSPreStage(ts, ts->ptime);CHKERRQ(ierr); 183 /* Compute implicit update */ 184 mimex->stage_time = ts->ptime + ts->time_step; 185 ierr = VecCopy(sol, update);CHKERRQ(ierr); 186 ierr = SNESSolve(ts->snes, NULL, update);CHKERRQ(ierr); 187 ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr); 188 ierr = VecGetArray(sol, &asol);CHKERRQ(ierr); 189 for (f = 0; f < Nf; ++f) { 190 PetscBool implicit; 191 192 ierr = PetscDSGetImplicit(prob, f, &implicit);CHKERRQ(ierr); 193 if (!implicit) continue; 194 for (p = pStart; p < pEnd; ++p) { 195 PetscScalar *au, *as; 196 PetscInt fdof, fcdof, d; 197 198 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 199 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 200 ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr); 201 ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr); 202 for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d]; 203 } 204 } 205 ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr); 206 ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr); 207 /* Compute explicit update */ 208 ierr = TSComputeRHSFunction(ts, ts->ptime, sol, update);CHKERRQ(ierr); 209 ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr); 210 ierr = VecGetArray(sol, &asol);CHKERRQ(ierr); 211 for (f = 0; f < Nf; ++f) { 212 PetscBool implicit; 213 214 ierr = PetscDSGetImplicit(prob, f, &implicit);CHKERRQ(ierr); 215 if (implicit) continue; 216 for (p = pStart; p < pEnd; ++p) { 217 PetscScalar *au, *as; 218 PetscInt fdof, fcdof, d; 219 220 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 221 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 222 ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr); 223 ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr); 224 for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d]; 225 } 226 } 227 ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr); 228 ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr); 229 ierr = TSPostStage(ts, ts->ptime, 0, &sol);CHKERRQ(ierr); 230 ts->ptime += ts->time_step; 231 ts->steps++; 232 PetscFunctionReturn(0); 233 } 234 235 236 #undef __FUNCT__ 237 #define __FUNCT__ "TSStep_Mimex_Implicit" 238 /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */ 239 static PetscErrorCode TSStep_Mimex_Implicit(TS ts) 240 { 241 TS_Mimex *mimex = (TS_Mimex *) ts->data; 242 Vec sol = ts->vec_sol; 243 Vec update = mimex->update; 244 PetscErrorCode ierr; 245 246 PetscFunctionBegin; 247 ierr = TSPreStep(ts);CHKERRQ(ierr); 248 ierr = TSPreStage(ts, ts->ptime);CHKERRQ(ierr); 249 /* Compute implicit update */ 250 mimex->stage_time = ts->ptime + ts->time_step; 251 ts->ptime += ts->time_step; 252 ierr = VecCopy(sol, update);CHKERRQ(ierr); 253 ierr = SNESSolve(ts->snes, NULL, update);CHKERRQ(ierr); 254 ierr = VecCopy(update, sol);CHKERRQ(ierr); 255 ierr = TSPostStage(ts, ts->ptime, 0, &sol);CHKERRQ(ierr); 256 ts->steps++; 257 PetscFunctionReturn(0); 258 } 259 260 #undef __FUNCT__ 261 #define __FUNCT__ "TSStep_Mimex" 262 static PetscErrorCode TSStep_Mimex(TS ts) 263 { 264 TS_Mimex *mimex = (TS_Mimex*)ts->data; 265 PetscErrorCode ierr; 266 267 PetscFunctionBegin; 268 switch(mimex->version) { 269 case 0: 270 ierr = TSStep_Mimex_Split(ts);CHKERRQ(ierr); break; 271 case 1: 272 ierr = TSStep_Mimex_Implicit(ts);CHKERRQ(ierr); break; 273 default: 274 SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version); 275 } 276 PetscFunctionReturn(0); 277 } 278 279 /*------------------------------------------------------------*/ 280 281 #undef __FUNCT__ 282 #define __FUNCT__ "TSSetUp_Mimex" 283 static PetscErrorCode TSSetUp_Mimex(TS ts) 284 { 285 TS_Mimex *mimex = (TS_Mimex*)ts->data; 286 PetscErrorCode ierr; 287 288 PetscFunctionBegin; 289 ierr = VecDuplicate(ts->vec_sol, &mimex->update);CHKERRQ(ierr); 290 ierr = VecDuplicate(ts->vec_sol, &mimex->Xdot);CHKERRQ(ierr); 291 PetscFunctionReturn(0); 292 } 293 294 #undef __FUNCT__ 295 #define __FUNCT__ "TSReset_Mimex" 296 static PetscErrorCode TSReset_Mimex(TS ts) 297 { 298 TS_Mimex *mimex = (TS_Mimex*)ts->data; 299 PetscErrorCode ierr; 300 301 PetscFunctionBegin; 302 ierr = VecDestroy(&mimex->update);CHKERRQ(ierr); 303 ierr = VecDestroy(&mimex->Xdot);CHKERRQ(ierr); 304 PetscFunctionReturn(0); 305 } 306 307 #undef __FUNCT__ 308 #define __FUNCT__ "TSDestroy_Mimex" 309 static PetscErrorCode TSDestroy_Mimex(TS ts) 310 { 311 PetscErrorCode ierr; 312 313 PetscFunctionBegin; 314 ierr = TSReset_Mimex(ts);CHKERRQ(ierr); 315 ierr = PetscFree(ts->data);CHKERRQ(ierr); 316 PetscFunctionReturn(0); 317 } 318 /*------------------------------------------------------------*/ 319 320 #undef __FUNCT__ 321 #define __FUNCT__ "TSSetFromOptions_Mimex" 322 static PetscErrorCode TSSetFromOptions_Mimex(PetscOptions *PetscOptionsObject, TS ts) 323 { 324 TS_Mimex *mimex = (TS_Mimex *) ts->data; 325 PetscErrorCode ierr; 326 327 PetscFunctionBegin; 328 ierr = PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");CHKERRQ(ierr); 329 { 330 ierr = PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);CHKERRQ(ierr); 331 } 332 ierr = PetscOptionsTail();CHKERRQ(ierr); 333 PetscFunctionReturn(0); 334 } 335 336 #undef __FUNCT__ 337 #define __FUNCT__ "TSView_Mimex" 338 static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer) 339 { 340 TS_Mimex *mimex = (TS_Mimex *) ts->data; 341 PetscBool iascii; 342 PetscErrorCode ierr; 343 344 PetscFunctionBegin; 345 ierr = PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);CHKERRQ(ierr); 346 if (iascii) { 347 ierr = PetscViewerASCIIPrintf(viewer, " Version = %D\n", mimex->version);CHKERRQ(ierr); 348 } 349 if (ts->snes) {ierr = SNESView(ts->snes, viewer);CHKERRQ(ierr);} 350 PetscFunctionReturn(0); 351 } 352 353 #undef __FUNCT__ 354 #define __FUNCT__ "TSInterpolate_Mimex" 355 static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X) 356 { 357 PetscReal alpha = (ts->ptime - t)/ts->time_step; 358 PetscErrorCode ierr; 359 360 PetscFunctionBegin; 361 ierr = VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);CHKERRQ(ierr); 362 PetscFunctionReturn(0); 363 } 364 365 #undef __FUNCT__ 366 #define __FUNCT__ "TSComputeLinearStability_Mimex" 367 PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi) 368 { 369 PetscFunctionBegin; 370 *yr = 1.0 + xr; 371 *yi = xi; 372 PetscFunctionReturn(0); 373 } 374 /* ------------------------------------------------------------ */ 375 376 /*MC 377 TSMIMEX - ODE solver using the explicit forward Mimex method 378 379 Level: beginner 380 381 .seealso: TSCreate(), TS, TSSetType(), TSBEULER 382 383 M*/ 384 #undef __FUNCT__ 385 #define __FUNCT__ "TSCreate_Mimex" 386 PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts) 387 { 388 TS_Mimex *mimex; 389 PetscErrorCode ierr; 390 391 PetscFunctionBegin; 392 ts->ops->setup = TSSetUp_Mimex; 393 ts->ops->step = TSStep_Mimex; 394 ts->ops->reset = TSReset_Mimex; 395 ts->ops->destroy = TSDestroy_Mimex; 396 ts->ops->setfromoptions = TSSetFromOptions_Mimex; 397 ts->ops->view = TSView_Mimex; 398 ts->ops->interpolate = TSInterpolate_Mimex; 399 ts->ops->linearstability = TSComputeLinearStability_Mimex; 400 ts->ops->snesfunction = SNESTSFormFunction_Mimex; 401 ts->ops->snesjacobian = SNESTSFormJacobian_Mimex; 402 403 ierr = PetscNewLog(ts,&mimex);CHKERRQ(ierr); 404 ts->data = (void*)mimex; 405 406 mimex->version = 1; 407 PetscFunctionReturn(0); 408 } 409