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; 10 PetscReal stage_time; 11 } TS_Mimex; 12 13 #undef __FUNCT__ 14 #define __FUNCT__ "TSMimexGetX0AndXdot" 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 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 /* 46 This defines the nonlinear equation that is to be solved with SNES 47 G(U) = F[t0+dt, U, (U-U0)*shift] = 0 48 */ 49 #undef __FUNCT__ 50 #define __FUNCT__ "SNESTSFormFunction_Mimex" 51 static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts) 52 { 53 TS_Mimex *mimex = (TS_Mimex *) ts->data; 54 DM dm, dmsave; 55 Vec X0, Xdot; 56 PetscReal shift = 1./ts->time_step; 57 PetscErrorCode ierr; 58 59 PetscFunctionBegin; 60 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 61 ierr = TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr); 62 ierr = VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);CHKERRQ(ierr); 63 64 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 65 dmsave = ts->dm; 66 ts->dm = dm; 67 ierr = TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);CHKERRQ(ierr); 68 ts->dm = dmsave; 69 ierr = TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr); 70 PetscFunctionReturn(0); 71 } 72 73 #undef __FUNCT__ 74 #define __FUNCT__ "SNESTSFormJacobian_Mimex" 75 static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts) 76 { 77 TS_Mimex *mimex = (TS_Mimex *) ts->data; 78 DM dm, dmsave; 79 Vec Xdot; 80 PetscReal shift = 1./ts->time_step; 81 PetscErrorCode ierr; 82 83 PetscFunctionBegin; 84 /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */ 85 ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr); 86 ierr = TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr); 87 88 /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */ 89 dmsave = ts->dm; 90 ts->dm = dm; 91 ierr = TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);CHKERRQ(ierr); 92 ts->dm = dmsave; 93 ierr = TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr); 94 PetscFunctionReturn(0); 95 } 96 97 #undef __FUNCT__ 98 #define __FUNCT__ "TSStep_Mimex" 99 static PetscErrorCode TSStep_Mimex(TS ts) 100 { 101 TS_Mimex *mimex = (TS_Mimex *) ts->data; 102 DM dm; 103 PetscDS prob; 104 PetscSection s; 105 Vec sol = ts->vec_sol, update = mimex->update; 106 const PetscScalar *aupdate; 107 PetscScalar *asol, dt = ts->time_step; 108 PetscInt Nf, f, pStart, pEnd, p; 109 PetscErrorCode ierr; 110 111 PetscFunctionBegin; 112 ierr = TSGetDM(ts, &dm);CHKERRQ(ierr); 113 ierr = DMGetDS(dm, &prob);CHKERRQ(ierr); 114 ierr = DMGetDefaultSection(dm, &s);CHKERRQ(ierr); 115 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 116 ierr = PetscSectionGetChart(s, &pStart, &pEnd);CHKERRQ(ierr); 117 ierr = TSPreStep(ts);CHKERRQ(ierr); 118 ierr = TSPreStage(ts, ts->ptime);CHKERRQ(ierr); 119 /* Compute implicit update */ 120 mimex->stage_time = ts->ptime + ts->time_step; 121 ierr = VecCopy(sol, update);CHKERRQ(ierr); 122 ierr = SNESSolve(ts->snes, NULL, update);CHKERRQ(ierr); 123 ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr); 124 ierr = VecGetArray(sol, &asol);CHKERRQ(ierr); 125 for (f = 0; f < Nf; ++f) { 126 PetscObject obj; 127 PetscClassId id; 128 129 ierr = PetscDSGetDiscretization(prob, f, &obj);CHKERRQ(ierr); 130 ierr = PetscObjectGetClassId(obj, &id);CHKERRQ(ierr); 131 if (id == PETSCFV_CLASSID) continue; 132 for (p = pStart; p < pEnd; ++p) { 133 PetscScalar *au, *as; 134 PetscInt fdof, fcdof, d; 135 136 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 137 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 138 ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr); 139 ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr); 140 for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d]; 141 } 142 } 143 ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr); 144 ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr); 145 /* Compute explicit update */ 146 ierr = TSComputeRHSFunction(ts, ts->ptime, sol, update);CHKERRQ(ierr); 147 ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr); 148 ierr = VecGetArray(sol, &asol);CHKERRQ(ierr); 149 for (f = 0; f < Nf; ++f) { 150 PetscObject obj; 151 PetscClassId id; 152 153 ierr = PetscDSGetDiscretization(prob, f, &obj);CHKERRQ(ierr); 154 ierr = PetscObjectGetClassId(obj, &id);CHKERRQ(ierr); 155 if (id != PETSCFV_CLASSID) continue; 156 for (p = pStart; p < pEnd; ++p) { 157 PetscScalar *au, *as; 158 PetscInt fdof, fcdof, d; 159 160 ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr); 161 ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr); 162 ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr); 163 ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr); 164 for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d]; 165 } 166 } 167 ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr); 168 ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr); 169 ierr = TSPostStage(ts, ts->ptime, 0, &sol);CHKERRQ(ierr); 170 ts->ptime += ts->time_step; 171 ts->steps++; 172 PetscFunctionReturn(0); 173 } 174 175 /*------------------------------------------------------------*/ 176 177 #undef __FUNCT__ 178 #define __FUNCT__ "TSSetUp_Mimex" 179 static PetscErrorCode TSSetUp_Mimex(TS ts) 180 { 181 TS_Mimex *mimex = (TS_Mimex*)ts->data; 182 PetscErrorCode ierr; 183 184 PetscFunctionBegin; 185 ierr = VecDuplicate(ts->vec_sol, &mimex->update);CHKERRQ(ierr); 186 ierr = VecDuplicate(ts->vec_sol, &mimex->Xdot);CHKERRQ(ierr); 187 PetscFunctionReturn(0); 188 } 189 190 #undef __FUNCT__ 191 #define __FUNCT__ "TSReset_Mimex" 192 static PetscErrorCode TSReset_Mimex(TS ts) 193 { 194 TS_Mimex *mimex = (TS_Mimex*)ts->data; 195 PetscErrorCode ierr; 196 197 PetscFunctionBegin; 198 ierr = VecDestroy(&mimex->update);CHKERRQ(ierr); 199 ierr = VecDestroy(&mimex->Xdot);CHKERRQ(ierr); 200 PetscFunctionReturn(0); 201 } 202 203 #undef __FUNCT__ 204 #define __FUNCT__ "TSDestroy_Mimex" 205 static PetscErrorCode TSDestroy_Mimex(TS ts) 206 { 207 PetscErrorCode ierr; 208 209 PetscFunctionBegin; 210 ierr = TSReset_Mimex(ts);CHKERRQ(ierr); 211 ierr = PetscFree(ts->data);CHKERRQ(ierr); 212 PetscFunctionReturn(0); 213 } 214 /*------------------------------------------------------------*/ 215 216 #undef __FUNCT__ 217 #define __FUNCT__ "TSSetFromOptions_Mimex" 218 static PetscErrorCode TSSetFromOptions_Mimex(PetscOptions *PetscOptionsObject,TS ts) 219 { 220 PetscFunctionBegin; 221 PetscFunctionReturn(0); 222 } 223 224 #undef __FUNCT__ 225 #define __FUNCT__ "TSView_Mimex" 226 static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer) 227 { 228 PetscFunctionBegin; 229 PetscFunctionReturn(0); 230 } 231 232 #undef __FUNCT__ 233 #define __FUNCT__ "TSInterpolate_Mimex" 234 static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X) 235 { 236 PetscReal alpha = (ts->ptime - t)/ts->time_step; 237 PetscErrorCode ierr; 238 239 PetscFunctionBegin; 240 ierr = VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);CHKERRQ(ierr); 241 PetscFunctionReturn(0); 242 } 243 244 #undef __FUNCT__ 245 #define __FUNCT__ "TSComputeLinearStability_Mimex" 246 PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi) 247 { 248 PetscFunctionBegin; 249 *yr = 1.0 + xr; 250 *yi = xi; 251 PetscFunctionReturn(0); 252 } 253 /* ------------------------------------------------------------ */ 254 255 /*MC 256 TSMIMEX - ODE solver using the explicit forward Mimex method 257 258 Level: beginner 259 260 .seealso: TSCreate(), TS, TSSetType(), TSBEULER 261 262 M*/ 263 #undef __FUNCT__ 264 #define __FUNCT__ "TSCreate_Mimex" 265 PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts) 266 { 267 TS_Mimex *mimex; 268 PetscErrorCode ierr; 269 270 PetscFunctionBegin; 271 ts->ops->setup = TSSetUp_Mimex; 272 ts->ops->step = TSStep_Mimex; 273 ts->ops->reset = TSReset_Mimex; 274 ts->ops->destroy = TSDestroy_Mimex; 275 ts->ops->setfromoptions = TSSetFromOptions_Mimex; 276 ts->ops->view = TSView_Mimex; 277 ts->ops->interpolate = TSInterpolate_Mimex; 278 ts->ops->linearstability = TSComputeLinearStability_Mimex; 279 ts->ops->snesfunction = SNESTSFormFunction_Mimex; 280 ts->ops->snesjacobian = SNESTSFormJacobian_Mimex; 281 282 ierr = PetscNewLog(ts,&mimex);CHKERRQ(ierr); 283 ts->data = (void*)mimex; 284 PetscFunctionReturn(0); 285 } 286