1 static char help[] = "Simple wrapper object to solve DAE of the form:\n\ 2 \\dot{U} = f(U,V)\n\ 3 F(U,V) = 0\n\n"; 4 5 #include <petscts.h> 6 7 /* ----------------------------------------------------------------------------*/ 8 9 typedef struct _p_TSDAESimple *TSDAESimple; 10 struct _p_TSDAESimple { 11 MPI_Comm comm; 12 PetscErrorCode (*setfromoptions)(TSDAESimple, PetscOptionItems *); 13 PetscErrorCode (*solve)(TSDAESimple, Vec); 14 PetscErrorCode (*destroy)(TSDAESimple); 15 Vec U, V; 16 PetscErrorCode (*f)(PetscReal, Vec, Vec, Vec, void *); 17 PetscErrorCode (*F)(PetscReal, Vec, Vec, Vec, void *); 18 void *fctx, *Fctx; 19 void *data; 20 }; 21 22 PetscErrorCode TSDAESimpleCreate(MPI_Comm comm, TSDAESimple *tsdae) { 23 PetscFunctionBeginUser; 24 PetscCall(PetscNew(tsdae)); 25 (*tsdae)->comm = comm; 26 PetscFunctionReturn(0); 27 } 28 29 PetscErrorCode TSDAESimpleSetRHSFunction(TSDAESimple tsdae, Vec U, PetscErrorCode (*f)(PetscReal, Vec, Vec, Vec, void *), void *ctx) { 30 PetscFunctionBeginUser; 31 tsdae->f = f; 32 tsdae->U = U; 33 PetscCall(PetscObjectReference((PetscObject)U)); 34 tsdae->fctx = ctx; 35 PetscFunctionReturn(0); 36 } 37 38 PetscErrorCode TSDAESimpleSetIFunction(TSDAESimple tsdae, Vec V, PetscErrorCode (*F)(PetscReal, Vec, Vec, Vec, void *), void *ctx) { 39 PetscFunctionBeginUser; 40 tsdae->F = F; 41 tsdae->V = V; 42 PetscCall(PetscObjectReference((PetscObject)V)); 43 tsdae->Fctx = ctx; 44 PetscFunctionReturn(0); 45 } 46 47 PetscErrorCode TSDAESimpleDestroy(TSDAESimple *tsdae) { 48 PetscFunctionBeginUser; 49 PetscCall((*(*tsdae)->destroy)(*tsdae)); 50 PetscCall(VecDestroy(&(*tsdae)->U)); 51 PetscCall(VecDestroy(&(*tsdae)->V)); 52 PetscCall(PetscFree(*tsdae)); 53 PetscFunctionReturn(0); 54 } 55 56 PetscErrorCode TSDAESimpleSolve(TSDAESimple tsdae, Vec Usolution) { 57 PetscFunctionBeginUser; 58 PetscCall((*tsdae->solve)(tsdae, Usolution)); 59 PetscFunctionReturn(0); 60 } 61 62 PetscErrorCode TSDAESimpleSetFromOptions(TSDAESimple tsdae) { 63 PetscFunctionBeginUser; 64 PetscCall((*tsdae->setfromoptions)(PetscOptionsObject, tsdae)); 65 PetscFunctionReturn(0); 66 } 67 68 /* ----------------------------------------------------------------------------*/ 69 /* 70 Integrates the system by integrating the reduced ODE system and solving the 71 algebraic constraints at each stage with a separate SNES solve. 72 */ 73 74 typedef struct { 75 PetscReal t; 76 TS ts; 77 SNES snes; 78 Vec U; 79 } TSDAESimple_Reduced; 80 81 /* 82 Defines the RHS function that is passed to the time-integrator. 83 84 Solves F(U,V) for V and then computes f(U,V) 85 86 */ 87 PetscErrorCode TSDAESimple_Reduced_TSFunction(TS ts, PetscReal t, Vec U, Vec F, void *actx) { 88 TSDAESimple tsdae = (TSDAESimple)actx; 89 TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data; 90 91 PetscFunctionBeginUser; 92 red->t = t; 93 red->U = U; 94 PetscCall(SNESSolve(red->snes, NULL, tsdae->V)); 95 PetscCall((*tsdae->f)(t, U, tsdae->V, F, tsdae->fctx)); 96 PetscFunctionReturn(0); 97 } 98 99 /* 100 Defines the nonlinear function that is passed to the nonlinear solver 101 102 */ 103 PetscErrorCode TSDAESimple_Reduced_SNESFunction(SNES snes, Vec V, Vec F, void *actx) { 104 TSDAESimple tsdae = (TSDAESimple)actx; 105 TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data; 106 107 PetscFunctionBeginUser; 108 PetscCall((*tsdae->F)(red->t, red->U, V, F, tsdae->Fctx)); 109 PetscFunctionReturn(0); 110 } 111 112 PetscErrorCode TSDAESimpleSolve_Reduced(TSDAESimple tsdae, Vec U) { 113 TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data; 114 115 PetscFunctionBeginUser; 116 PetscCall(TSSolve(red->ts, U)); 117 PetscFunctionReturn(0); 118 } 119 120 PetscErrorCode TSDAESimpleSetFromOptions_Reduced(TSDAESimple tsdae, PetscOptionItems *PetscOptionsObject) { 121 TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data; 122 123 PetscFunctionBeginUser; 124 PetscCall(TSSetFromOptions(red->ts)); 125 PetscCall(SNESSetFromOptions(red->snes)); 126 PetscFunctionReturn(0); 127 } 128 129 PetscErrorCode TSDAESimpleDestroy_Reduced(TSDAESimple tsdae) { 130 TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data; 131 132 PetscFunctionBeginUser; 133 PetscCall(TSDestroy(&red->ts)); 134 PetscCall(SNESDestroy(&red->snes)); 135 PetscCall(PetscFree(red)); 136 PetscFunctionReturn(0); 137 } 138 139 PetscErrorCode TSDAESimpleSetUp_Reduced(TSDAESimple tsdae) { 140 TSDAESimple_Reduced *red; 141 Vec tsrhs; 142 143 PetscFunctionBeginUser; 144 PetscCall(PetscNew(&red)); 145 tsdae->data = red; 146 147 tsdae->setfromoptions = TSDAESimpleSetFromOptions_Reduced; 148 tsdae->solve = TSDAESimpleSolve_Reduced; 149 tsdae->destroy = TSDAESimpleDestroy_Reduced; 150 151 PetscCall(TSCreate(tsdae->comm, &red->ts)); 152 PetscCall(TSSetProblemType(red->ts, TS_NONLINEAR)); 153 PetscCall(TSSetType(red->ts, TSEULER)); 154 PetscCall(TSSetExactFinalTime(red->ts, TS_EXACTFINALTIME_STEPOVER)); 155 PetscCall(VecDuplicate(tsdae->U, &tsrhs)); 156 PetscCall(TSSetRHSFunction(red->ts, tsrhs, TSDAESimple_Reduced_TSFunction, tsdae)); 157 PetscCall(VecDestroy(&tsrhs)); 158 159 PetscCall(SNESCreate(tsdae->comm, &red->snes)); 160 PetscCall(SNESSetOptionsPrefix(red->snes, "tsdaesimple_")); 161 PetscCall(SNESSetFunction(red->snes, NULL, TSDAESimple_Reduced_SNESFunction, tsdae)); 162 PetscCall(SNESSetJacobian(red->snes, NULL, NULL, SNESComputeJacobianDefault, tsdae)); 163 PetscFunctionReturn(0); 164 } 165 166 /* ----------------------------------------------------------------------------*/ 167 168 /* 169 Integrates the system by integrating directly the entire DAE system 170 */ 171 172 typedef struct { 173 TS ts; 174 Vec UV, UF, VF; 175 VecScatter scatterU, scatterV; 176 } TSDAESimple_Full; 177 178 /* 179 Defines the RHS function that is passed to the time-integrator. 180 181 f(U,V) 182 0 183 184 */ 185 PetscErrorCode TSDAESimple_Full_TSRHSFunction(TS ts, PetscReal t, Vec UV, Vec F, void *actx) { 186 TSDAESimple tsdae = (TSDAESimple)actx; 187 TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data; 188 189 PetscFunctionBeginUser; 190 PetscCall(VecSet(F, 0.0)); 191 PetscCall(VecScatterBegin(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE)); 192 PetscCall(VecScatterEnd(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE)); 193 PetscCall(VecScatterBegin(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE)); 194 PetscCall(VecScatterEnd(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE)); 195 PetscCall((*tsdae->f)(t, tsdae->U, tsdae->V, full->UF, tsdae->fctx)); 196 PetscCall(VecScatterBegin(full->scatterU, full->UF, F, INSERT_VALUES, SCATTER_FORWARD)); 197 PetscCall(VecScatterEnd(full->scatterU, full->UF, F, INSERT_VALUES, SCATTER_FORWARD)); 198 PetscFunctionReturn(0); 199 } 200 201 /* 202 Defines the nonlinear function that is passed to the nonlinear solver 203 204 \dot{U} 205 F(U,V) 206 207 */ 208 PetscErrorCode TSDAESimple_Full_TSIFunction(TS ts, PetscReal t, Vec UV, Vec UVdot, Vec F, void *actx) { 209 TSDAESimple tsdae = (TSDAESimple)actx; 210 TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data; 211 212 PetscFunctionBeginUser; 213 PetscCall(VecCopy(UVdot, F)); 214 PetscCall(VecScatterBegin(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE)); 215 PetscCall(VecScatterEnd(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE)); 216 PetscCall(VecScatterBegin(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE)); 217 PetscCall(VecScatterEnd(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE)); 218 PetscCall((*tsdae->F)(t, tsdae->U, tsdae->V, full->VF, tsdae->Fctx)); 219 PetscCall(VecScatterBegin(full->scatterV, full->VF, F, INSERT_VALUES, SCATTER_FORWARD)); 220 PetscCall(VecScatterEnd(full->scatterV, full->VF, F, INSERT_VALUES, SCATTER_FORWARD)); 221 PetscFunctionReturn(0); 222 } 223 224 PetscErrorCode TSDAESimpleSolve_Full(TSDAESimple tsdae, Vec U) { 225 TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data; 226 227 PetscFunctionBeginUser; 228 PetscCall(VecSet(full->UV, 1.0)); 229 PetscCall(VecScatterBegin(full->scatterU, U, full->UV, INSERT_VALUES, SCATTER_FORWARD)); 230 PetscCall(VecScatterEnd(full->scatterU, U, full->UV, INSERT_VALUES, SCATTER_FORWARD)); 231 PetscCall(TSSolve(full->ts, full->UV)); 232 PetscCall(VecScatterBegin(full->scatterU, full->UV, U, INSERT_VALUES, SCATTER_REVERSE)); 233 PetscCall(VecScatterEnd(full->scatterU, full->UV, U, INSERT_VALUES, SCATTER_REVERSE)); 234 PetscFunctionReturn(0); 235 } 236 237 PetscErrorCode TSDAESimpleSetFromOptions_Full(TSDAESimple tsdae, PetscOptionItems *PetscOptionsObject) { 238 TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data; 239 240 PetscFunctionBeginUser; 241 PetscCall(TSSetFromOptions(full->ts)); 242 PetscFunctionReturn(0); 243 } 244 245 PetscErrorCode TSDAESimpleDestroy_Full(TSDAESimple tsdae) { 246 TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data; 247 248 PetscFunctionBeginUser; 249 PetscCall(TSDestroy(&full->ts)); 250 PetscCall(VecDestroy(&full->UV)); 251 PetscCall(VecDestroy(&full->UF)); 252 PetscCall(VecDestroy(&full->VF)); 253 PetscCall(VecScatterDestroy(&full->scatterU)); 254 PetscCall(VecScatterDestroy(&full->scatterV)); 255 PetscCall(PetscFree(full)); 256 PetscFunctionReturn(0); 257 } 258 259 PetscErrorCode TSDAESimpleSetUp_Full(TSDAESimple tsdae) { 260 TSDAESimple_Full *full; 261 Vec tsrhs; 262 PetscInt nU, nV, UVstart; 263 IS is; 264 265 PetscFunctionBeginUser; 266 PetscCall(PetscNew(&full)); 267 tsdae->data = full; 268 269 tsdae->setfromoptions = TSDAESimpleSetFromOptions_Full; 270 tsdae->solve = TSDAESimpleSolve_Full; 271 tsdae->destroy = TSDAESimpleDestroy_Full; 272 273 PetscCall(TSCreate(tsdae->comm, &full->ts)); 274 PetscCall(TSSetProblemType(full->ts, TS_NONLINEAR)); 275 PetscCall(TSSetType(full->ts, TSROSW)); 276 PetscCall(TSSetExactFinalTime(full->ts, TS_EXACTFINALTIME_STEPOVER)); 277 PetscCall(VecDuplicate(tsdae->U, &full->UF)); 278 PetscCall(VecDuplicate(tsdae->V, &full->VF)); 279 280 PetscCall(VecGetLocalSize(tsdae->U, &nU)); 281 PetscCall(VecGetLocalSize(tsdae->V, &nV)); 282 PetscCall(VecCreateMPI(tsdae->comm, nU + nV, PETSC_DETERMINE, &tsrhs)); 283 PetscCall(VecDuplicate(tsrhs, &full->UV)); 284 285 PetscCall(VecGetOwnershipRange(tsrhs, &UVstart, NULL)); 286 PetscCall(ISCreateStride(tsdae->comm, nU, UVstart, 1, &is)); 287 PetscCall(VecScatterCreate(tsdae->U, NULL, tsrhs, is, &full->scatterU)); 288 PetscCall(ISDestroy(&is)); 289 PetscCall(ISCreateStride(tsdae->comm, nV, UVstart + nU, 1, &is)); 290 PetscCall(VecScatterCreate(tsdae->V, NULL, tsrhs, is, &full->scatterV)); 291 PetscCall(ISDestroy(&is)); 292 293 PetscCall(TSSetRHSFunction(full->ts, tsrhs, TSDAESimple_Full_TSRHSFunction, tsdae)); 294 PetscCall(TSSetIFunction(full->ts, NULL, TSDAESimple_Full_TSIFunction, tsdae)); 295 PetscCall(VecDestroy(&tsrhs)); 296 PetscFunctionReturn(0); 297 } 298 299 /* ----------------------------------------------------------------------------*/ 300 301 /* 302 Simple example: f(U,V) = U + V 303 304 */ 305 PetscErrorCode f(PetscReal t, Vec U, Vec V, Vec F, void *ctx) { 306 PetscFunctionBeginUser; 307 PetscCall(VecWAXPY(F, 1.0, U, V)); 308 PetscFunctionReturn(0); 309 } 310 311 /* 312 Simple example: F(U,V) = U - V 313 314 */ 315 PetscErrorCode F(PetscReal t, Vec U, Vec V, Vec F, void *ctx) { 316 PetscFunctionBeginUser; 317 PetscCall(VecWAXPY(F, -1.0, V, U)); 318 PetscFunctionReturn(0); 319 } 320 321 int main(int argc, char **argv) { 322 TSDAESimple tsdae; 323 Vec U, V, Usolution; 324 325 PetscFunctionBeginUser; 326 PetscCall(PetscInitialize(&argc, &argv, (char *)0, help)); 327 PetscCall(TSDAESimpleCreate(PETSC_COMM_WORLD, &tsdae)); 328 329 PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 1, PETSC_DETERMINE, &U)); 330 PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 1, PETSC_DETERMINE, &V)); 331 PetscCall(TSDAESimpleSetRHSFunction(tsdae, U, f, NULL)); 332 PetscCall(TSDAESimpleSetIFunction(tsdae, V, F, NULL)); 333 334 PetscCall(VecDuplicate(U, &Usolution)); 335 PetscCall(VecSet(Usolution, 1.0)); 336 337 /* PetscCall(TSDAESimpleSetUp_Full(tsdae)); */ 338 PetscCall(TSDAESimpleSetUp_Reduced(tsdae)); 339 340 PetscCall(TSDAESimpleSetFromOptions(tsdae)); 341 PetscCall(TSDAESimpleSolve(tsdae, Usolution)); 342 PetscCall(TSDAESimpleDestroy(&tsdae)); 343 344 PetscCall(VecDestroy(&U)); 345 PetscCall(VecDestroy(&Usolution)); 346 PetscCall(VecDestroy(&V)); 347 PetscCall(PetscFinalize()); 348 return 0; 349 } 350