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, Xstar, G; 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 TS_Mimex *mimex = (TS_Mimex *) ts->data; 50 PetscErrorCode ierr; 51 52 PetscFunctionBegin; 53 if (!mimex->Xstar) {ierr = VecDuplicate(ts->vec_sol, &mimex->Xstar);CHKERRQ(ierr);} 54 if (!mimex->G) {ierr = VecDuplicate(ts->vec_sol, &mimex->G);CHKERRQ(ierr);} 55 *Xstar = mimex->Xstar; 56 *G = mimex->G; 57 PetscFunctionReturn(0); 58 } 59 60 #undef __FUNCT__ 61 #define __FUNCT__ "TSMimexRestoreXstarAndG" 62 static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G) 63 { 64 PetscFunctionBegin; 65 PetscFunctionReturn(0); 66 } 67 68 /* 69 This defines the nonlinear equation that is to be solved with SNES 70 G(U) = F[t0+dt, U, (U-U0)*shift] = 0 71 */ 72 #undef __FUNCT__ 73 #define __FUNCT__ "SNESTSFormFunction_Mimex" 74 static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts) 75 { 76 TS_Mimex *mimex = (TS_Mimex *) ts->data; 77 DM dm, dmsave; 78 Vec X0, Xdot; 79 PetscReal shift = 1./ts->time_step; 80 PetscErrorCode ierr; 81 82 PetscFunctionBegin; 83 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 84 ierr = TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr); 85 ierr = VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);CHKERRQ(ierr); 86 87 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 88 dmsave = ts->dm; 89 ts->dm = dm; 90 ierr = TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);CHKERRQ(ierr); 91 if (mimex->version == 1) { 92 DM dm; 93 PetscDS prob; 94 PetscSection s; 95 Vec Xstar, G; 96 const PetscScalar *ax; 97 PetscScalar *axstar; 98 PetscInt Nf, f, pStart, pEnd, p; 99 100 ierr = TSGetDM(ts, &dm);CHKERRQ(ierr); 101 ierr = DMGetDS(dm, &prob);CHKERRQ(ierr); 102 ierr = DMGetDefaultSection(dm, &s);CHKERRQ(ierr); 103 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 104 ierr = PetscSectionGetChart(s, &pStart, &pEnd);CHKERRQ(ierr); 105 ierr = TSMimexGetXstarAndG(ts, dm, &Xstar, &G);CHKERRQ(ierr); 106 ierr = VecCopy(X0, Xstar);CHKERRQ(ierr); 107 ierr = VecGetArrayRead(x, &ax);CHKERRQ(ierr); 108 ierr = VecGetArray(Xstar, &axstar);CHKERRQ(ierr); 109 for (f = 0; f < Nf; ++f) { 110 PetscBool implicit; 111 112 ierr = PetscDSGetImplicit(prob, f, &implicit);CHKERRQ(ierr); 113 if (!implicit) continue; 114 for (p = pStart; p < pEnd; ++p) { 115 PetscScalar *a, *axs; 116 PetscInt fdof, fcdof, d; 117 118 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 119 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 120 ierr = DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);CHKERRQ(ierr); 121 ierr = DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);CHKERRQ(ierr); 122 for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d]; 123 } 124 } 125 ierr = VecRestoreArrayRead(x, &ax);CHKERRQ(ierr); 126 ierr = VecRestoreArray(Xstar, &axstar);CHKERRQ(ierr); 127 ierr = TSComputeRHSFunction(ts, ts->ptime, Xstar, G);CHKERRQ(ierr); 128 ierr = VecAXPY(y, -1.0, G);CHKERRQ(ierr); 129 ierr = TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);CHKERRQ(ierr); 130 } 131 ts->dm = dmsave; 132 ierr = TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr); 133 PetscFunctionReturn(0); 134 } 135 136 #undef __FUNCT__ 137 #define __FUNCT__ "SNESTSFormJacobian_Mimex" 138 static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts) 139 { 140 TS_Mimex *mimex = (TS_Mimex *) ts->data; 141 DM dm, dmsave; 142 Vec Xdot; 143 PetscReal shift = 1./ts->time_step; 144 PetscErrorCode ierr; 145 146 PetscFunctionBegin; 147 /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */ 148 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 149 ierr = TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr); 150 151 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 152 dmsave = ts->dm; 153 ts->dm = dm; 154 ierr = TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);CHKERRQ(ierr); 155 ts->dm = dmsave; 156 ierr = TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr); 157 PetscFunctionReturn(0); 158 } 159 160 #undef __FUNCT__ 161 #define __FUNCT__ "TSStep_Mimex_Split" 162 static PetscErrorCode TSStep_Mimex_Split(TS ts) 163 { 164 TS_Mimex *mimex = (TS_Mimex *) ts->data; 165 DM dm; 166 PetscDS prob; 167 PetscSection s; 168 Vec sol = ts->vec_sol, update = mimex->update; 169 const PetscScalar *aupdate; 170 PetscScalar *asol, dt = ts->time_step; 171 PetscInt Nf, f, pStart, pEnd, p; 172 PetscErrorCode ierr; 173 174 PetscFunctionBegin; 175 ierr = TSGetDM(ts, &dm);CHKERRQ(ierr); 176 ierr = DMGetDS(dm, &prob);CHKERRQ(ierr); 177 ierr = DMGetDefaultSection(dm, &s);CHKERRQ(ierr); 178 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 179 ierr = PetscSectionGetChart(s, &pStart, &pEnd);CHKERRQ(ierr); 180 ierr = TSPreStep(ts);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 = TSPreStep(ts);CHKERRQ(ierr); 247 ierr = TSPreStage(ts, ts->ptime);CHKERRQ(ierr); 248 /* Compute implicit update */ 249 mimex->stage_time = ts->ptime + ts->time_step; 250 ts->ptime += ts->time_step; 251 ierr = VecCopy(sol, update);CHKERRQ(ierr); 252 ierr = SNESSolve(ts->snes, NULL, update);CHKERRQ(ierr); 253 ierr = VecCopy(update, sol);CHKERRQ(ierr); 254 ierr = TSPostStage(ts, ts->ptime, 0, &sol);CHKERRQ(ierr); 255 ts->steps++; 256 PetscFunctionReturn(0); 257 } 258 259 #undef __FUNCT__ 260 #define __FUNCT__ "TSStep_Mimex" 261 static PetscErrorCode TSStep_Mimex(TS ts) 262 { 263 TS_Mimex *mimex = (TS_Mimex*)ts->data; 264 PetscErrorCode ierr; 265 266 PetscFunctionBegin; 267 switch(mimex->version) { 268 case 0: 269 ierr = TSStep_Mimex_Split(ts);CHKERRQ(ierr); break; 270 case 1: 271 ierr = TSStep_Mimex_Implicit(ts);CHKERRQ(ierr); break; 272 default: 273 SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version); 274 } 275 PetscFunctionReturn(0); 276 } 277 278 /*------------------------------------------------------------*/ 279 280 #undef __FUNCT__ 281 #define __FUNCT__ "TSSetUp_Mimex" 282 static PetscErrorCode TSSetUp_Mimex(TS ts) 283 { 284 TS_Mimex *mimex = (TS_Mimex*)ts->data; 285 PetscErrorCode ierr; 286 287 PetscFunctionBegin; 288 ierr = VecDuplicate(ts->vec_sol, &mimex->update);CHKERRQ(ierr); 289 ierr = VecDuplicate(ts->vec_sol, &mimex->Xdot);CHKERRQ(ierr); 290 PetscFunctionReturn(0); 291 } 292 293 #undef __FUNCT__ 294 #define __FUNCT__ "TSReset_Mimex" 295 static PetscErrorCode TSReset_Mimex(TS ts) 296 { 297 TS_Mimex *mimex = (TS_Mimex*)ts->data; 298 PetscErrorCode ierr; 299 300 PetscFunctionBegin; 301 ierr = VecDestroy(&mimex->Xstar);CHKERRQ(ierr); 302 ierr = VecDestroy(&mimex->G);CHKERRQ(ierr); 303 ierr = VecDestroy(&mimex->update);CHKERRQ(ierr); 304 ierr = VecDestroy(&mimex->Xdot);CHKERRQ(ierr); 305 PetscFunctionReturn(0); 306 } 307 308 #undef __FUNCT__ 309 #define __FUNCT__ "TSDestroy_Mimex" 310 static PetscErrorCode TSDestroy_Mimex(TS ts) 311 { 312 PetscErrorCode ierr; 313 314 PetscFunctionBegin; 315 ierr = TSReset_Mimex(ts);CHKERRQ(ierr); 316 ierr = PetscFree(ts->data);CHKERRQ(ierr); 317 PetscFunctionReturn(0); 318 } 319 /*------------------------------------------------------------*/ 320 321 #undef __FUNCT__ 322 #define __FUNCT__ "TSSetFromOptions_Mimex" 323 static PetscErrorCode TSSetFromOptions_Mimex(PetscOptions *PetscOptionsObject, TS ts) 324 { 325 TS_Mimex *mimex = (TS_Mimex *) ts->data; 326 PetscErrorCode ierr; 327 328 PetscFunctionBegin; 329 ierr = PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");CHKERRQ(ierr); 330 { 331 ierr = PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);CHKERRQ(ierr); 332 } 333 ierr = PetscOptionsTail();CHKERRQ(ierr); 334 PetscFunctionReturn(0); 335 } 336 337 #undef __FUNCT__ 338 #define __FUNCT__ "TSView_Mimex" 339 static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer) 340 { 341 PetscFunctionBegin; 342 PetscFunctionReturn(0); 343 } 344 345 #undef __FUNCT__ 346 #define __FUNCT__ "TSInterpolate_Mimex" 347 static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X) 348 { 349 PetscReal alpha = (ts->ptime - t)/ts->time_step; 350 PetscErrorCode ierr; 351 352 PetscFunctionBegin; 353 ierr = VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);CHKERRQ(ierr); 354 PetscFunctionReturn(0); 355 } 356 357 #undef __FUNCT__ 358 #define __FUNCT__ "TSComputeLinearStability_Mimex" 359 PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi) 360 { 361 PetscFunctionBegin; 362 *yr = 1.0 + xr; 363 *yi = xi; 364 PetscFunctionReturn(0); 365 } 366 /* ------------------------------------------------------------ */ 367 368 /*MC 369 TSMIMEX - ODE solver using the explicit forward Mimex method 370 371 Level: beginner 372 373 .seealso: TSCreate(), TS, TSSetType(), TSBEULER 374 375 M*/ 376 #undef __FUNCT__ 377 #define __FUNCT__ "TSCreate_Mimex" 378 PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts) 379 { 380 TS_Mimex *mimex; 381 PetscErrorCode ierr; 382 383 PetscFunctionBegin; 384 ts->ops->setup = TSSetUp_Mimex; 385 ts->ops->step = TSStep_Mimex; 386 ts->ops->reset = TSReset_Mimex; 387 ts->ops->destroy = TSDestroy_Mimex; 388 ts->ops->setfromoptions = TSSetFromOptions_Mimex; 389 ts->ops->view = TSView_Mimex; 390 ts->ops->interpolate = TSInterpolate_Mimex; 391 ts->ops->linearstability = TSComputeLinearStability_Mimex; 392 ts->ops->snesfunction = SNESTSFormFunction_Mimex; 393 ts->ops->snesjacobian = SNESTSFormJacobian_Mimex; 394 395 ierr = PetscNewLog(ts,&mimex);CHKERRQ(ierr); 396 ts->data = (void*)mimex; 397 398 mimex->version = 1; 399 PetscFunctionReturn(0); 400 } 401