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