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