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 static PetscErrorCode TSSetUp_Mimex(TS ts) 250 { 251 TS_Mimex *mimex = (TS_Mimex *)ts->data; 252 253 PetscFunctionBegin; 254 PetscCall(VecDuplicate(ts->vec_sol, &mimex->update)); 255 PetscCall(VecDuplicate(ts->vec_sol, &mimex->Xdot)); 256 PetscFunctionReturn(PETSC_SUCCESS); 257 } 258 259 static PetscErrorCode TSReset_Mimex(TS ts) 260 { 261 TS_Mimex *mimex = (TS_Mimex *)ts->data; 262 263 PetscFunctionBegin; 264 PetscCall(VecDestroy(&mimex->update)); 265 PetscCall(VecDestroy(&mimex->Xdot)); 266 PetscFunctionReturn(PETSC_SUCCESS); 267 } 268 269 static PetscErrorCode TSDestroy_Mimex(TS ts) 270 { 271 PetscFunctionBegin; 272 PetscCall(TSReset_Mimex(ts)); 273 PetscCall(PetscFree(ts->data)); 274 PetscFunctionReturn(PETSC_SUCCESS); 275 } 276 277 static PetscErrorCode TSSetFromOptions_Mimex(TS ts, PetscOptionItems PetscOptionsObject) 278 { 279 TS_Mimex *mimex = (TS_Mimex *)ts->data; 280 281 PetscFunctionBegin; 282 PetscOptionsHeadBegin(PetscOptionsObject, "MIMEX ODE solver options"); 283 { 284 PetscCall(PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL)); 285 } 286 PetscOptionsHeadEnd(); 287 PetscFunctionReturn(PETSC_SUCCESS); 288 } 289 290 static PetscErrorCode TSView_Mimex(TS ts, PetscViewer viewer) 291 { 292 TS_Mimex *mimex = (TS_Mimex *)ts->data; 293 PetscBool isascii; 294 295 PetscFunctionBegin; 296 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii)); 297 if (isascii) PetscCall(PetscViewerASCIIPrintf(viewer, " Version = %" PetscInt_FMT "\n", mimex->version)); 298 PetscFunctionReturn(PETSC_SUCCESS); 299 } 300 301 static PetscErrorCode TSInterpolate_Mimex(TS ts, PetscReal t, Vec X) 302 { 303 PetscReal alpha = (ts->ptime - t) / ts->time_step; 304 305 PetscFunctionBegin; 306 PetscCall(VecAXPBY(ts->vec_sol, 1.0 - alpha, alpha, X)); 307 PetscFunctionReturn(PETSC_SUCCESS); 308 } 309 310 static PetscErrorCode TSComputeLinearStability_Mimex(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi) 311 { 312 PetscFunctionBegin; 313 *yr = 1.0 + xr; 314 *yi = xi; 315 PetscFunctionReturn(PETSC_SUCCESS); 316 } 317 318 /*MC 319 TSMIMEX - ODE solver using the explicit forward Mimex method 320 321 Level: beginner 322 323 .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER` 324 M*/ 325 PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts) 326 { 327 TS_Mimex *mimex; 328 329 PetscFunctionBegin; 330 ts->ops->setup = TSSetUp_Mimex; 331 ts->ops->step = TSStep_Mimex; 332 ts->ops->reset = TSReset_Mimex; 333 ts->ops->destroy = TSDestroy_Mimex; 334 ts->ops->setfromoptions = TSSetFromOptions_Mimex; 335 ts->ops->view = TSView_Mimex; 336 ts->ops->interpolate = TSInterpolate_Mimex; 337 ts->ops->linearstability = TSComputeLinearStability_Mimex; 338 ts->ops->snesfunction = SNESTSFormFunction_Mimex; 339 ts->ops->snesjacobian = SNESTSFormJacobian_Mimex; 340 ts->default_adapt_type = TSADAPTNONE; 341 342 PetscCall(PetscNew(&mimex)); 343 ts->data = (void *)mimex; 344 345 mimex->version = 1; 346 PetscFunctionReturn(PETSC_SUCCESS); 347 } 348