Actual source code: ex4.h
1: #pragma once
3: #include <petscdmda.h>
4: #include <petscts.h>
6: /* Helpers are `static inline` for small math; bare `static` for address-taken callbacks
7: (e.g. ShallowWaterRHS2D registered via TSSetRHSFunction) and for one-shot setup helpers
8: that are intentionally duplicated per translation unit (tutorial isolation - ex4.c and
9: ex4fwd.c stay independent). */
11: /* Wet/dry threshold: cells with h below this fall back to zero flux to avoid division by ~0. */
12: #define EX4_DRY_TOL 1e-10
14: /* User-supplied problem parameters; passed by const-pointer to the setup helpers below to keep
15: call sites readable. Internal/derived state (DM, TS handle, dx, dy) lives in ShallowWater2DCtx. */
16: typedef struct {
17: PetscInt nx, ny;
18: PetscReal Lx, Ly;
19: PetscReal g;
20: PetscReal dt;
21: PetscReal h0;
22: PetscReal Ax, Ay; /* wave amplitudes in x and y */
23: PetscBool verify_mms;
24: } ShallowWater2DConfig;
26: typedef struct {
27: DM da;
28: PetscInt nx, ny;
29: PetscReal Lx, Ly;
30: PetscReal dx, dy;
31: PetscReal g;
32: PetscReal dt;
33: TS ts;
34: PetscReal h0;
35: PetscReal Ax, Ay; /* wave amplitudes in x and y */
36: PetscBool verify_mms;
37: } ShallowWater2DCtx;
39: static inline void ManufacturedSolution2D(PetscReal Lx, PetscReal Ly, PetscReal x, PetscReal y, PetscReal t, PetscReal h0, PetscReal A, PetscReal *h, PetscReal *hu, PetscReal *hv)
40: {
41: PetscReal sx = PetscSinReal(2.0 * PETSC_PI * x / Lx);
42: PetscReal sy = PetscSinReal(2.0 * PETSC_PI * y / Ly);
43: PetscReal cx = PetscCosReal(2.0 * PETSC_PI * x / Lx);
44: PetscReal cy = PetscCosReal(2.0 * PETSC_PI * y / Ly);
46: *h = h0 + A * PetscSinReal(t) * sx * sy;
47: *hu = A * PetscCosReal(t) * cx * sy;
48: *hv = A * PetscSinReal(t) * sx * cy;
49: }
51: /* MMS uses a single scalar amplitude A: callers pass sw->Ax with the invariant Ax == Ay enforced
52: by SetInitialCondition() before the integration begins. ShallowWater2DContextCreate()
53: collectively enforces h0 - A > EX4_DRY_TOL when verify_mms is enabled (h = h0 + A*sin(t)*sx*sy
54: ranges in [h0-A, h0+A] and the source divides by h*h), so this hot path skips the per-cell guard. */
55: static inline void ManufacturedSource2D(PetscReal Lx, PetscReal Ly, PetscReal x, PetscReal y, PetscReal t, PetscReal g, PetscReal h0, PetscReal A, PetscReal *S_h, PetscReal *S_hu, PetscReal *S_hv)
56: {
57: PetscReal h, hu, hv;
58: PetscReal kx = 2.0 * PETSC_PI / Lx;
59: PetscReal ky = 2.0 * PETSC_PI / Ly;
60: PetscReal sx = PetscSinReal(kx * x);
61: PetscReal sy = PetscSinReal(ky * y);
62: PetscReal cx = PetscCosReal(kx * x);
63: PetscReal cy = PetscCosReal(ky * y);
64: PetscReal ht = A * PetscCosReal(t) * sx * sy;
65: PetscReal hut = -A * PetscSinReal(t) * cx * sy;
66: PetscReal hvt = A * PetscCosReal(t) * sx * cy;
67: PetscReal hx = A * PetscSinReal(t) * kx * cx * sy;
68: PetscReal hy = A * PetscSinReal(t) * ky * sx * cy;
69: PetscReal hux = -A * PetscCosReal(t) * kx * sx * sy;
70: PetscReal huy = A * PetscCosReal(t) * ky * cx * cy;
71: PetscReal hvx = A * PetscSinReal(t) * kx * cx * cy;
72: PetscReal hvy = -A * PetscSinReal(t) * ky * sx * sy;
73: PetscReal hu2_over_h_x, huhv_over_h_y, huhv_over_h_x, hv2_over_h_y;
75: ManufacturedSolution2D(Lx, Ly, x, y, t, h0, A, &h, &hu, &hv);
76: hu2_over_h_x = (2.0 * hu * hux * h - hu * hu * hx) / (h * h);
77: huhv_over_h_y = ((huy * hv + hu * hvy) * h - hu * hv * hy) / (h * h);
78: huhv_over_h_x = ((hux * hv + hu * hvx) * h - hu * hv * hx) / (h * h);
79: hv2_over_h_y = (2.0 * hv * hvy * h - hv * hv * hy) / (h * h);
81: *S_h = ht + hux + hvy;
82: *S_hu = hut + hu2_over_h_x + g * h * hx + huhv_over_h_y;
83: *S_hv = hvt + huhv_over_h_x + hv2_over_h_y + g * h * hy;
84: }
86: static inline void ComputeFluxX(PetscReal g, PetscReal h, PetscReal hu, PetscReal hv, PetscReal *F_h, PetscReal *F_hu, PetscReal *F_hv, PetscReal *u, PetscReal *c)
87: {
88: if (h > EX4_DRY_TOL) {
89: *u = hu / h;
90: *c = PetscSqrtReal(g * h);
91: *F_h = hu;
92: *F_hu = hu * *u + 0.5 * g * h * h;
93: *F_hv = hu * (hv / h);
94: } else {
95: *u = 0.0;
96: *c = 0.0;
97: *F_h = 0.0;
98: *F_hu = 0.0;
99: *F_hv = 0.0;
100: }
101: }
103: static inline void ComputeFluxY(PetscReal g, PetscReal h, PetscReal hu, PetscReal hv, PetscReal *G_h, PetscReal *G_hu, PetscReal *G_hv, PetscReal *v, PetscReal *c)
104: {
105: if (h > EX4_DRY_TOL) {
106: *v = hv / h;
107: *c = PetscSqrtReal(g * h);
108: *G_h = hv;
109: *G_hu = hv * (hu / h);
110: *G_hv = hv * *v + 0.5 * g * h * h;
111: } else {
112: *v = 0.0;
113: *c = 0.0;
114: *G_h = 0.0;
115: *G_hu = 0.0;
116: *G_hv = 0.0;
117: }
118: }
120: static PetscErrorCode ShallowWaterRHS2D(TS ts, PetscReal t, Vec X, Vec F_vec, PetscCtx ctx)
121: {
122: ShallowWater2DCtx *sw = (ShallowWater2DCtx *)ctx;
123: Vec X_local;
124: const PetscScalar ***x;
125: PetscScalar ***f;
126: PetscInt xs, ys, xm, ym;
128: PetscFunctionBeginUser;
129: PetscCall(DMDAGetCorners(sw->da, &xs, &ys, NULL, &xm, &ym, NULL));
130: PetscCall(DMGetLocalVector(sw->da, &X_local));
131: PetscCall(DMGlobalToLocalBegin(sw->da, X, INSERT_VALUES, X_local));
132: PetscCall(DMGlobalToLocalEnd(sw->da, X, INSERT_VALUES, X_local));
133: PetscCall(DMDAVecGetArrayDOFRead(sw->da, X_local, (void *)&x));
134: PetscCall(DMDAVecGetArrayDOFWrite(sw->da, F_vec, &f));
136: /* Only the Rusanov flux is implemented. */
137: for (PetscInt j = ys; j < ys + ym; j++) {
138: for (PetscInt i = xs; i < xs + xm; i++) {
139: PetscReal h = PetscRealPart(x[j][i][0]);
140: PetscReal hu = PetscRealPart(x[j][i][1]);
141: PetscReal hv = PetscRealPart(x[j][i][2]);
142: PetscReal h_im1 = PetscRealPart(x[j][i - 1][0]);
143: PetscReal hu_im1 = PetscRealPart(x[j][i - 1][1]);
144: PetscReal hv_im1 = PetscRealPart(x[j][i - 1][2]);
145: PetscReal h_ip1 = PetscRealPart(x[j][i + 1][0]);
146: PetscReal hu_ip1 = PetscRealPart(x[j][i + 1][1]);
147: PetscReal hv_ip1 = PetscRealPart(x[j][i + 1][2]);
148: PetscReal h_jm1 = PetscRealPart(x[j - 1][i][0]);
149: PetscReal hu_jm1 = PetscRealPart(x[j - 1][i][1]);
150: PetscReal hv_jm1 = PetscRealPart(x[j - 1][i][2]);
151: PetscReal h_jp1 = PetscRealPart(x[j + 1][i][0]);
152: PetscReal hu_jp1 = PetscRealPart(x[j + 1][i][1]);
153: PetscReal hv_jp1 = PetscRealPart(x[j + 1][i][2]);
154: PetscReal F_h_i, F_hu_i, F_hv_i, u, c;
155: PetscReal F_h_im1, F_hu_im1, F_hv_im1, u_im1, c_im1;
156: PetscReal F_h_ip1, F_hu_ip1, F_hv_ip1, u_ip1, c_ip1;
157: PetscReal G_h_j, G_hu_j, G_hv_j, v, c_y;
158: PetscReal G_h_jm1, G_hu_jm1, G_hv_jm1, v_jm1, c_jm1;
159: PetscReal G_h_jp1, G_hu_jp1, G_hv_jp1, v_jp1, c_jp1;
160: PetscReal alpha_left, alpha_right, beta_bottom, beta_top;
161: PetscReal flux_h_left, flux_hu_left, flux_hv_left;
162: PetscReal flux_h_right, flux_hu_right, flux_hv_right;
163: PetscReal flux_h_bottom, flux_hu_bottom, flux_hv_bottom;
164: PetscReal flux_h_top, flux_hu_top, flux_hv_top;
166: ComputeFluxX(sw->g, h, hu, hv, &F_h_i, &F_hu_i, &F_hv_i, &u, &c);
167: ComputeFluxX(sw->g, h_im1, hu_im1, hv_im1, &F_h_im1, &F_hu_im1, &F_hv_im1, &u_im1, &c_im1);
168: ComputeFluxX(sw->g, h_ip1, hu_ip1, hv_ip1, &F_h_ip1, &F_hu_ip1, &F_hv_ip1, &u_ip1, &c_ip1);
169: ComputeFluxY(sw->g, h, hu, hv, &G_h_j, &G_hu_j, &G_hv_j, &v, &c_y);
170: ComputeFluxY(sw->g, h_jm1, hu_jm1, hv_jm1, &G_h_jm1, &G_hu_jm1, &G_hv_jm1, &v_jm1, &c_jm1);
171: ComputeFluxY(sw->g, h_jp1, hu_jp1, hv_jp1, &G_h_jp1, &G_hu_jp1, &G_hv_jp1, &v_jp1, &c_jp1);
173: alpha_left = PetscMax(PetscAbsReal(u_im1) + c_im1, PetscAbsReal(u) + c);
174: alpha_right = PetscMax(PetscAbsReal(u) + c, PetscAbsReal(u_ip1) + c_ip1);
175: beta_bottom = PetscMax(PetscAbsReal(v_jm1) + c_jm1, PetscAbsReal(v) + c_y);
176: beta_top = PetscMax(PetscAbsReal(v) + c_y, PetscAbsReal(v_jp1) + c_jp1);
178: flux_h_left = 0.5 * (F_h_im1 + F_h_i - alpha_left * (h - h_im1));
179: flux_hu_left = 0.5 * (F_hu_im1 + F_hu_i - alpha_left * (hu - hu_im1));
180: flux_hv_left = 0.5 * (F_hv_im1 + F_hv_i - alpha_left * (hv - hv_im1));
181: flux_h_right = 0.5 * (F_h_i + F_h_ip1 - alpha_right * (h_ip1 - h));
182: flux_hu_right = 0.5 * (F_hu_i + F_hu_ip1 - alpha_right * (hu_ip1 - hu));
183: flux_hv_right = 0.5 * (F_hv_i + F_hv_ip1 - alpha_right * (hv_ip1 - hv));
184: flux_h_bottom = 0.5 * (G_h_jm1 + G_h_j - beta_bottom * (h - h_jm1));
185: flux_hu_bottom = 0.5 * (G_hu_jm1 + G_hu_j - beta_bottom * (hu - hu_jm1));
186: flux_hv_bottom = 0.5 * (G_hv_jm1 + G_hv_j - beta_bottom * (hv - hv_jm1));
187: flux_h_top = 0.5 * (G_h_j + G_h_jp1 - beta_top * (h_jp1 - h));
188: flux_hu_top = 0.5 * (G_hu_j + G_hu_jp1 - beta_top * (hu_jp1 - hu));
189: flux_hv_top = 0.5 * (G_hv_j + G_hv_jp1 - beta_top * (hv_jp1 - hv));
191: f[j][i][0] = -(flux_h_right - flux_h_left) / sw->dx - (flux_h_top - flux_h_bottom) / sw->dy;
192: f[j][i][1] = -(flux_hu_right - flux_hu_left) / sw->dx - (flux_hu_top - flux_hu_bottom) / sw->dy;
193: f[j][i][2] = -(flux_hv_right - flux_hv_left) / sw->dx - (flux_hv_top - flux_hv_bottom) / sw->dy;
195: if (sw->verify_mms) {
196: PetscReal x_coord = ((PetscReal)i + 0.5) * sw->dx;
197: PetscReal y_coord = ((PetscReal)j + 0.5) * sw->dy;
198: PetscReal S_h = 0.0, S_hu = 0.0, S_hv = 0.0;
200: ManufacturedSource2D(sw->Lx, sw->Ly, x_coord, y_coord, t, sw->g, sw->h0, sw->Ax, &S_h, &S_hu, &S_hv);
201: f[j][i][0] += S_h;
202: f[j][i][1] += S_hu;
203: f[j][i][2] += S_hv;
204: }
205: }
206: }
208: PetscCall(DMDAVecRestoreArrayDOFRead(sw->da, X_local, (void *)&x));
209: PetscCall(DMDAVecRestoreArrayDOFWrite(sw->da, F_vec, &f));
210: PetscCall(DMRestoreLocalVector(sw->da, &X_local));
211: PetscFunctionReturn(PETSC_SUCCESS);
212: }
214: static PetscErrorCode ShallowWater2DContextCreate(DM da, const ShallowWater2DConfig *cfg, ShallowWater2DCtx **ctx)
215: {
216: ShallowWater2DCtx *sw;
218: PetscFunctionBeginUser;
219: /* ManufacturedSource2D divides by h*h, where h = h0 + A*sin(t)*sx*sy ranges in [h0-A, h0+A].
220: Enforce h0 - Ax > EX4_DRY_TOL once on the DA's comm so misuse fails collectively at setup
221: instead of deadlocking on a per-cell SETERRQ from one rank during the RHS evaluation. */
222: PetscCheck(!cfg->verify_mms || cfg->h0 - cfg->Ax > EX4_DRY_TOL, PetscObjectComm((PetscObject)da), PETSC_ERR_ARG_OUTOFRANGE, "MMS amplitude Ax (%g) must leave h0 (%g) > EX4_DRY_TOL (%g)", (double)cfg->Ax, (double)cfg->h0, (double)EX4_DRY_TOL);
223: PetscCall(PetscNew(&sw));
224: /* Borrowed reference; caller owns the DM and ShallowWater2DContextDestroy() does not free it. */
225: sw->da = da;
226: sw->nx = cfg->nx;
227: sw->ny = cfg->ny;
228: sw->Lx = cfg->Lx;
229: sw->Ly = cfg->Ly;
230: sw->g = cfg->g;
231: sw->dx = cfg->Lx / cfg->nx;
232: sw->dy = cfg->Ly / cfg->ny;
233: sw->dt = cfg->dt;
234: sw->h0 = cfg->h0;
235: sw->Ax = cfg->Ax;
236: sw->Ay = cfg->Ay;
237: sw->verify_mms = cfg->verify_mms;
239: PetscCall(TSCreate(PetscObjectComm((PetscObject)da), &sw->ts));
240: PetscCall(TSSetProblemType(sw->ts, TS_NONLINEAR));
241: PetscCall(TSSetRHSFunction(sw->ts, NULL, ShallowWaterRHS2D, sw));
242: PetscCall(TSSetType(sw->ts, TSRK));
243: PetscCall(TSRKSetType(sw->ts, TSRK4));
244: PetscCall(TSSetTimeStep(sw->ts, cfg->dt));
245: PetscCall(TSSetMaxSteps(sw->ts, 1));
246: PetscCall(TSSetMaxTime(sw->ts, cfg->dt));
247: PetscCall(TSSetExactFinalTime(sw->ts, TS_EXACTFINALTIME_MATCHSTEP));
248: PetscCall(TSSetFromOptions(sw->ts));
250: *ctx = sw;
251: PetscFunctionReturn(PETSC_SUCCESS);
252: }
254: static inline PetscErrorCode ShallowWater2DContextDestroy(ShallowWater2DCtx **ctx)
255: {
256: PetscFunctionBeginUser;
257: if (!*ctx) PetscFunctionReturn(PETSC_SUCCESS);
258: PetscCall(TSDestroy(&(*ctx)->ts));
259: PetscCall(PetscFree(*ctx));
260: PetscFunctionReturn(PETSC_SUCCESS);
261: }
263: /* Advance a single state vector one TS step starting from physical time t_start. The MMS source
264: (when verify_mms is enabled in the context) is evaluated against the TS time, so callers driving
265: multi-step verification runs must pass the cumulative start time of each step rather than 0. */
266: static PetscErrorCode ShallowWaterStep2DVec(ShallowWater2DCtx *sw, PetscReal t_start, Vec x)
267: {
268: PetscFunctionBeginUser;
269: /* The TSSetTimeStep and TSSetMaxSteps calls below look redundant with ShallowWater2DContextCreate()
270: but are load-bearing: TSSetMaxTime forces the last sub-step to shrink to land exactly on
271: t_start+dt (and TSAdapt may have shrunk dt during the step), so the trailing dt persists into the
272: next call unless reset; and -ts_max_steps from the command line overrides the create-time 1, so
273: re-asserting it per call enforces this wrapper's "exactly one step per call" contract. */
274: PetscCall(TSSetTime(sw->ts, t_start));
275: PetscCall(TSSetStepNumber(sw->ts, 0));
276: PetscCall(TSSetTimeStep(sw->ts, sw->dt));
277: PetscCall(TSSetMaxSteps(sw->ts, 1));
278: PetscCall(TSSetMaxTime(sw->ts, t_start + sw->dt));
279: PetscCall(TSSolve(sw->ts, x));
280: PetscFunctionReturn(PETSC_SUCCESS);
281: }
283: static inline void ShallowWaterSolution_Wave2D(PetscReal Lx, PetscReal Ly, PetscReal x, PetscReal y, PetscReal t, PetscReal g, PetscReal h0, PetscReal Ax, PetscReal Ay, PetscReal *h, PetscReal *hu, PetscReal *hv)
284: {
285: PetscReal c = PetscSqrtReal(g * h0);
286: PetscReal kx = 2.0 * PETSC_PI / Lx;
287: PetscReal ky = 2.0 * PETSC_PI / Ly;
288: PetscReal omega_x = c * kx;
289: PetscReal omega_y = c * ky;
291: *h = h0 + Ax * PetscSinReal(kx * x - omega_x * t) + Ay * PetscSinReal(ky * y - omega_y * t);
292: *hu = (*h) * (c / h0) * Ax * PetscCosReal(kx * x - omega_x * t);
293: *hv = (*h) * (c / h0) * Ay * PetscCosReal(ky * y - omega_y * t);
294: }
296: /*
297: SetupForwardProblem - Create DM, shallow water context, and solution vector
298: */
299: static PetscErrorCode SetupForwardProblem(const ShallowWater2DConfig *cfg, DM *da_state, ShallowWater2DCtx **sw_ctx, Vec *x)
300: {
301: PetscFunctionBeginUser;
302: PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, DMDA_STENCIL_STAR, cfg->nx, cfg->ny, PETSC_DECIDE, PETSC_DECIDE, 3, 1, NULL, NULL, da_state));
303: PetscCall(DMSetFromOptions(*da_state));
304: PetscCall(DMSetUp(*da_state));
305: PetscCall(ShallowWater2DContextCreate(*da_state, cfg, sw_ctx));
306: PetscCall(DMCreateGlobalVector(*da_state, x));
307: PetscFunctionReturn(PETSC_SUCCESS);
308: }
310: /*
311: SetInitialCondition - Set initial condition on solution vector from analytic solution
312: */
313: static PetscErrorCode SetInitialCondition(DM da_state, Vec x, ShallowWater2DCtx *sw, PetscBool use_mms)
314: {
315: PetscScalar ***x_array;
316: PetscInt xs, ys, xm, ym;
318: PetscFunctionBeginUser;
319: PetscCheck(!use_mms || PetscAbsReal(sw->Ax - sw->Ay) <= 100 * PETSC_MACHINE_EPSILON * PetscMax(PetscAbsReal(sw->Ax), PetscAbsReal(sw->Ay)), PetscObjectComm((PetscObject)da_state), PETSC_ERR_ARG_INCOMP, "MMS requires isotropic amplitude (Ax == Ay); got Ax=%g, Ay=%g",
320: (double)sw->Ax, (double)sw->Ay);
321: PetscCall(DMDAGetCorners(da_state, &xs, &ys, NULL, &xm, &ym, NULL));
322: PetscCall(DMDAVecGetArrayDOFWrite(da_state, x, &x_array));
323: for (PetscInt j = ys; j < ys + ym; j++) {
324: for (PetscInt i = xs; i < xs + xm; i++) {
325: PetscReal xc = ((PetscReal)i + 0.5) * sw->dx;
326: PetscReal yc = ((PetscReal)j + 0.5) * sw->dy;
327: PetscReal h, hu, hv;
329: if (use_mms) ManufacturedSolution2D(sw->Lx, sw->Ly, xc, yc, 0.0, sw->h0, sw->Ax, &h, &hu, &hv);
330: else ShallowWaterSolution_Wave2D(sw->Lx, sw->Ly, xc, yc, 0.0, sw->g, sw->h0, sw->Ax, sw->Ay, &h, &hu, &hv);
331: x_array[j][i][0] = h;
332: x_array[j][i][1] = hu;
333: x_array[j][i][2] = hv;
334: }
335: }
336: PetscCall(DMDAVecRestoreArrayDOFWrite(da_state, x, &x_array));
337: PetscFunctionReturn(PETSC_SUCCESS);
338: }