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