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